fix: preserve relationship direction in get_by_node_uuid and BFS edge search

graphiti_core/models/edges/edge_db_queries.py

@@ -188,13 +188,36 @@ def get_entity_edge_save_bulk_query(provider: GraphProvider, has_aoss: bool = Fa
 
 def get_entity_edge_return_query(provider: GraphProvider) -> str:
     # `fact_embedding` is not returned by default and must be manually loaded using `load_fact_embedding()`.
+    # Direction semantics: NEO4J/FALKORDB/NEPTUNE read source/target from the relationship itself
+    # via Cypher's built-in `startNode(e)`/`endNode(e)`, so the returned direction matches the
+    # stored relationship regardless of whether the surrounding MATCH is directed or undirected.
+    # KUZU models relationships as a `RelatesToNode_` intermediate node plus two `:RELATES_TO`
+    # edges, so `e` is a node, not a relationship, and `startNode`/`endNode` are not applicable.
+    # KUZU's MATCH is always directed `(n)-[:RELATES_TO]->(e:RelatesToNode_)-[:RELATES_TO]->(m)`,
+    # so `n`/`m` already reflect the true source/target.
 
-    if provider == GraphProvider.NEPTUNE:
+    if provider == GraphProvider.KUZU:
         return """
         e.uuid AS uuid,
         n.uuid AS source_node_uuid,
         m.uuid AS target_node_uuid,
         e.group_id AS group_id,
+        e.created_at AS created_at,
+        e.name AS name,
+        e.fact AS fact,
+        e.episodes AS episodes,
+        e.expired_at AS expired_at,
+        e.valid_at AS valid_at,
+        e.invalid_at AS invalid_at,
+        e.attributes AS attributes
+    """
+
+    if provider == GraphProvider.NEPTUNE:
+        return """
+        e.uuid AS uuid,
+        startNode(e).uuid AS source_node_uuid,
+        endNode(e).uuid AS target_node_uuid,
+        e.group_id AS group_id,
         e.name AS name,
         e.fact AS fact,
         split(e.episodes, ',') AS episodes,
@@ -207,8 +230,8 @@ def get_entity_edge_return_query(provider: GraphProvider) -> str:
 
     return """
         e.uuid AS uuid,
-        n.uuid AS source_node_uuid,
-        m.uuid AS target_node_uuid,
+        startNode(e).uuid AS source_node_uuid,
+        endNode(e).uuid AS target_node_uuid,
         e.group_id AS group_id,
         e.created_at AS created_at,
         e.name AS name,
@@ -217,11 +240,8 @@ def get_entity_edge_return_query(provider: GraphProvider) -> str:
         e.expired_at AS expired_at,
         e.valid_at AS valid_at,
         e.invalid_at AS invalid_at,
-    """ + (
-        'e.attributes AS attributes'
-        if provider == GraphProvider.KUZU
-        else 'properties(e) AS attributes'
-    )
+        properties(e) AS attributes
+    """
 
 
 def get_community_edge_save_query(provider: GraphProvider) -> str:

tests/test_edge_db_queries.py

@@ -0,0 +1,52 @@
+"""Unit tests for the entity edge return-query builder.
+
+Locks in the direction-correctness fix in `get_entity_edge_return_query`:
+NEO4J/FALKORDB/NEPTUNE must read source/target from the relationship itself
+(``startNode(e)``/``endNode(e)``) so the returned direction reflects the stored
+relationship regardless of whether the surrounding MATCH is directed or undirected.
+KUZU keeps ``n.uuid``/``m.uuid`` because its `e` is an intermediate node, not a
+relationship.
+"""
+
+from graphiti_core.driver.driver import GraphProvider
+from graphiti_core.models.edges.edge_db_queries import get_entity_edge_return_query
+
+
+def test_neo4j_uses_start_end_node():
+    query = get_entity_edge_return_query(GraphProvider.NEO4J)
+    assert 'startNode(e).uuid AS source_node_uuid' in query
+    assert 'endNode(e).uuid AS target_node_uuid' in query
+    assert 'n.uuid AS source_node_uuid' not in query
+    assert 'm.uuid AS target_node_uuid' not in query
+
+
+def test_falkordb_uses_start_end_node():
+    query = get_entity_edge_return_query(GraphProvider.FALKORDB)
+    assert 'startNode(e).uuid AS source_node_uuid' in query
+    assert 'endNode(e).uuid AS target_node_uuid' in query
+
+
+def test_neptune_uses_start_end_node_with_split_episodes():
+    query = get_entity_edge_return_query(GraphProvider.NEPTUNE)
+    assert 'startNode(e).uuid AS source_node_uuid' in query
+    assert 'endNode(e).uuid AS target_node_uuid' in query
+    assert "split(e.episodes, ',') AS episodes" in query
+
+
+def test_kuzu_keeps_match_variables():
+    # KUZU models relationships as an intermediate RelatesToNode_ node, so `e` is a node
+    # and startNode/endNode are not applicable. Its MATCH is always directed, so n/m
+    # already reflect the true source/target.
+    query = get_entity_edge_return_query(GraphProvider.KUZU)
+    assert 'n.uuid AS source_node_uuid' in query
+    assert 'm.uuid AS target_node_uuid' in query
+    assert 'startNode' not in query
+    assert 'endNode' not in query
+    assert 'e.attributes AS attributes' in query
+
+
+def test_default_branch_returns_neo4j_or_falkordb_attributes():
+    # Both NEO4J and FALKORDB use `properties(e)`, unlike KUZU's `e.attributes`.
+    for provider in (GraphProvider.NEO4J, GraphProvider.FALKORDB):
+        query = get_entity_edge_return_query(provider)
+        assert 'properties(e) AS attributes' in query

tests/test_edge_int.py

@@ -240,6 +240,15 @@ async def test_entity_edge(graph_driver, mock_embedder):
     assert retrieved[0].created_at == now
     assert retrieved[0].group_id == group_id
 
+    # Get edge by node uuid from the target side — direction must be preserved
+    # (regression test: previously the queried node was always bound to `n`,
+    # flipping source/target when the queried node was the relationship target).
+    retrieved = await EntityEdge.get_by_node_uuid(graph_driver, bob_node.uuid)
+    assert len(retrieved) == 1
+    assert retrieved[0].uuid == entity_edge.uuid
+    assert retrieved[0].source_node_uuid == alice_node.uuid
+    assert retrieved[0].target_node_uuid == bob_node.uuid
+
     # Get fact embedding
     await entity_edge.load_fact_embedding(graph_driver)
     assert np.allclose(entity_edge.fact_embedding, edge_embedding)