From ced3b778e282f2c64fcf8af9aa4aad3f19f717e8 Mon Sep 17 00:00:00 2001
From: Guy Korland <gkorland@gmail.com>
Date: Fri, 17 Oct 2025 13:34:47 +0300
Subject: [PATCH] fix falkordb add triplet issue

* Fix FalkorDB edge creation: Add missing UUID fields
---
 graphiti_core/edges.py                        |  2 +
 graphiti_core/models/edges/edge_db_queries.py |  2 +-
 tests/test_edge_int.py                        | 67 +++++++++++++++++++
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/graphiti_core/edges.py b/graphiti_core/edges.py
index 82066e73..a3822594 100644
--- a/graphiti_core/edges.py
+++ b/graphiti_core/edges.py
@@ -286,6 +286,8 @@ class EntityEdge(Edge):
         edge_data: dict[str, Any] = {
             'source_uuid': self.source_node_uuid,
             'target_uuid': self.target_node_uuid,
+            'source_node_uuid': self.source_node_uuid,
+            'target_node_uuid': self.target_node_uuid,
             'uuid': self.uuid,
             'name': self.name,
             'group_id': self.group_id,
diff --git a/graphiti_core/models/edges/edge_db_queries.py b/graphiti_core/models/edges/edge_db_queries.py
index ef4578e2..ce8aa915 100644
--- a/graphiti_core/models/edges/edge_db_queries.py
+++ b/graphiti_core/models/edges/edge_db_queries.py
@@ -129,7 +129,7 @@ def get_entity_edge_save_bulk_query(provider: GraphProvider, has_aoss: bool = Fa
                 MATCH (source:Entity {uuid: edge.source_node_uuid})
                 MATCH (target:Entity {uuid: edge.target_node_uuid})
                 MERGE (source)-[r:RELATES_TO {uuid: edge.uuid}]->(target)
-                SET r = {uuid: edge.uuid, name: edge.name, group_id: edge.group_id, fact: edge.fact, episodes: edge.episodes,
+                SET r = {uuid: edge.uuid, source_node_uuid: edge.source_node_uuid, target_node_uuid: edge.target_node_uuid, name: edge.name, group_id: edge.group_id, fact: edge.fact, episodes: edge.episodes,
                 created_at: edge.created_at, expired_at: edge.expired_at, valid_at: edge.valid_at, invalid_at: edge.invalid_at, fact_embedding: vecf32(edge.fact_embedding)}
                 WITH r, edge
                 RETURN edge.uuid AS uuid
diff --git a/tests/test_edge_int.py b/tests/test_edge_int.py
index 15555d72..ec51e684 100644
--- a/tests/test_edge_int.py
+++ b/tests/test_edge_int.py
@@ -395,3 +395,70 @@ async def test_community_edge(graph_driver, mock_embedder):
     assert node_count == 0
 
     await graph_driver.close()
+
+
+@pytest.mark.asyncio
+async def test_entity_edge_stores_node_uuids(graph_driver, mock_embedder):
+    """Test that entity edges store source_node_uuid and target_node_uuid properties."""
+    now = datetime.now()
+
+    # Create entity nodes
+    alice_node = EntityNode(
+        name='Alice',
+        labels=[],
+        created_at=now,
+        summary='Alice summary',
+        group_id=group_id,
+    )
+    await alice_node.generate_name_embedding(mock_embedder)
+    await alice_node.save(graph_driver)
+
+    bob_node = EntityNode(
+        name='Bob', labels=[], created_at=now, summary='Bob summary', group_id=group_id
+    )
+    await bob_node.generate_name_embedding(mock_embedder)
+    await bob_node.save(graph_driver)
+
+    # Create entity edge
+    entity_edge = EntityEdge(
+        source_node_uuid=alice_node.uuid,
+        target_node_uuid=bob_node.uuid,
+        created_at=now,
+        name='knows',
+        fact='Alice knows Bob',
+        episodes=[],
+        expired_at=now,
+        valid_at=now,
+        invalid_at=now,
+        group_id=group_id,
+    )
+    await entity_edge.generate_embedding(mock_embedder)
+    await entity_edge.save(graph_driver)
+
+    # Verify edge properties directly from the database
+    # This query checks if the edge relationship itself has the UUID properties
+    query = """
+        MATCH (source:Entity)-[r:RELATES_TO]->(target:Entity)
+        WHERE r.uuid = $edge_uuid
+        RETURN
+            source.uuid AS source_uuid,
+            target.uuid AS target_uuid,
+            r.source_node_uuid AS stored_source_uuid,
+            r.target_node_uuid AS stored_target_uuid
+    """
+
+    results, _, _ = await graph_driver.execute_query(query, edge_uuid=entity_edge.uuid)
+
+    # Verify results
+    assert len(results) == 1
+    result = results[0]
+
+    # The edge properties should contain source_node_uuid and target_node_uuid
+    assert result['stored_source_uuid'] == alice_node.uuid
+    assert result['stored_target_uuid'] == bob_node.uuid
+
+    # Cleanup
+    await entity_edge.delete(graph_driver)
+    await alice_node.delete(graph_driver)
+    await bob_node.delete(graph_driver)
+    await graph_driver.close()