diff --git a/graphiti_core/edges.py b/graphiti_core/edges.py index 72c16d7a..491afa2b 100644 --- a/graphiti_core/edges.py +++ b/graphiti_core/edges.py @@ -223,7 +223,7 @@ class EntityEdge(Edge): async def load_fact_embedding(self, driver: AsyncDriver): query: LiteralString = """ - MATCH (n:Entity)-[e:MENTIONS {uuid: $uuid}]->(m:Entity) + MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity) RETURN e.fact_embedding AS fact_embedding """ records, _, _ = await driver.execute_query( @@ -334,8 +334,8 @@ class EntityEdge(Edge): async def get_by_node_uuid(cls, driver: AsyncDriver, node_uuid: str): query: LiteralString = ( """ - MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity) - """ + MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity) + """ + ENTITY_EDGE_RETURN ) records, _, _ = await driver.execute_query( diff --git a/tests/test_graphiti_int.py b/tests/test_graphiti_int.py index 3efd2697..838c6f70 100644 --- a/tests/test_graphiti_int.py +++ b/tests/test_graphiti_int.py @@ -26,6 +26,7 @@ from graphiti_core.edges import EntityEdge, EpisodicEdge from graphiti_core.graphiti import Graphiti from graphiti_core.helpers import semaphore_gather from graphiti_core.nodes import EntityNode, EpisodicNode +from graphiti_core.search.search_config_recipes import COMBINED_HYBRID_SEARCH_MMR from graphiti_core.search.search_helpers import search_results_to_context_string pytestmark = pytest.mark.integration @@ -65,7 +66,11 @@ async def test_graphiti_init(): logger = setup_logging() graphiti = Graphiti(NEO4J_URI, NEO4j_USER, NEO4j_PASSWORD) - results = await graphiti.search_(query='Who is the User?') + results = await graphiti.search_( + query='Who is the Tania?', + group_ids=['98602bec-2c05-4f34-876e-741b39fe6ae8'], + config=COMBINED_HYBRID_SEARCH_MMR, + ) pretty_results = search_results_to_context_string(results)