diff --git a/graphiti_core/driver/falkordb/operations/search_ops.py b/graphiti_core/driver/falkordb/operations/search_ops.py index c6be500c5a..de6ddfda89 100644 --- a/graphiti_core/driver/falkordb/operations/search_ops.py +++ b/graphiti_core/driver/falkordb/operations/search_ops.py @@ -302,8 +302,8 @@ async def edge_fulltext_search( 'edge_name_and_fact', limit=limit, provider=GraphProvider.FALKORDB ) + """ - YIELD relationship AS rel, score - MATCH (n:Entity)-[e:RELATES_TO {uuid: rel.uuid}]->(m:Entity) + YIELD relationship AS e, score + WITH e, score, startNode(e) AS n, endNode(e) AS m """ + filter_query + """ @@ -415,7 +415,9 @@ async def edge_bfs_search( UNWIND $bfs_origin_node_uuids AS origin_uuid MATCH path = (origin {{uuid: origin_uuid}})-[:RELATES_TO|MENTIONS*1..{max_depth}]->(:Entity) UNWIND relationships(path) AS rel - MATCH (n:Entity)-[e:RELATES_TO {{uuid: rel.uuid}}]-(m:Entity) + WITH rel + WHERE type(rel) = 'RELATES_TO' + WITH rel AS e, startNode(rel) AS n, endNode(rel) AS m """ + filter_query + """ diff --git a/tests/driver/test_falkordb_search_ops.py b/tests/driver/test_falkordb_search_ops.py new file mode 100644 index 0000000000..0a030fdfa4 --- /dev/null +++ b/tests/driver/test_falkordb_search_ops.py @@ -0,0 +1,52 @@ +from typing import Any + +import pytest + +from graphiti_core.driver.falkordb.operations.search_ops import FalkorSearchOperations +from graphiti_core.search.search_filters import SearchFilters + + +class RecordingExecutor: + def __init__(self): + self.cypher_query = '' + self.params: dict[str, Any] = {} + + async def execute_query(self, cypher_query_: str, **kwargs: Any): + self.cypher_query = cypher_query_ + self.params = kwargs + return [], None, None + + +@pytest.mark.asyncio +async def test_edge_fulltext_search_uses_returned_relationship_endpoints(): + executor = RecordingExecutor() + operations = FalkorSearchOperations() + + await operations.edge_fulltext_search( + executor, + 'api test system', + SearchFilters(), + group_ids=['group-a'], + ) + + assert 'YIELD relationship AS e, score' in executor.cypher_query + assert 'WITH e, score, startNode(e) AS n, endNode(e) AS m' in executor.cypher_query + assert 'uuid: rel.uuid' not in executor.cypher_query + + +@pytest.mark.asyncio +async def test_edge_bfs_search_uses_path_relationship_endpoints(): + executor = RecordingExecutor() + operations = FalkorSearchOperations() + + await operations.edge_bfs_search( + executor, + ['origin-uuid'], + 2, + SearchFilters(), + group_ids=['group-a'], + ) + + assert "WHERE type(rel) = 'RELATES_TO'" in executor.cypher_query + assert 'WITH rel AS e, startNode(rel) AS n, endNode(rel) AS m' in executor.cypher_query + assert 'uuid: rel.uuid' not in executor.cypher_query