Do not pollute the graph with internal/temp properties

Author: stellasia

CHANGELOG.md

@@ -28,8 +28,8 @@
 
 #### Other
 
-- The node internal `id` property that's used to create relationships between nodes has been renamed to `__kg_builder_id`. This releases the `id` name for domain-specific meaningful `id` property.
-
+- The `id` property on `__KG_Builder__` nodes is removed.
+- The `chunk_index` property on `__Entity__` nodes is removed. Use the `FROM_CHUNK` relationship instead.
 
 ## 1.7.0
 

src/neo4j_graphrag/experimental/components/entity_relation_extractor.py

@@ -146,12 +146,11 @@ def update_ids(
         """Make node IDs unique across chunks, document and pipeline runs
         by prefixing them with a unique prefix.
         """
-        prefix = f"{chunk.chunk_id}"
+        prefix = chunk.chunk_id
         for node in graph.nodes:
             node.id = f"{prefix}:{node.id}"
             if node.properties is None:
                 node.properties = {}
-            node.properties.update({"chunk_index": chunk.index})
         for rel in graph.relationships:
             rel.start_node_id = f"{prefix}:{rel.start_node_id}"
             rel.end_node_id = f"{prefix}:{rel.end_node_id}"

src/neo4j_graphrag/experimental/components/kg_writer.py

@@ -125,12 +125,8 @@ def __init__(
         self.is_version_5_23_or_above = is_version_5_23_or_above(version_tuple)
 
     def _db_setup(self) -> None:
-        # create index on __KGBuilder__.id
-        # used when creating the relationships
-        self.driver.execute_query(
-            "CREATE INDEX __entity__id IF NOT EXISTS  FOR (n:__KGBuilder__) ON (n.id)",
-            database_=self.neo4j_database,
-        )
+        # not used for now
+        pass
 
     @staticmethod
     def _nodes_to_rows(
@@ -148,45 +144,62 @@ def _nodes_to_rows(
 
     def _upsert_nodes(
         self, nodes: list[Neo4jNode], lexical_graph_config: LexicalGraphConfig
-    ) -> None:
+    ) -> dict[str, str]:
         """Upserts a single node into the Neo4j database."
 
         Args:
             nodes (list[Neo4jNode]): The nodes batch to upsert into the database.
         """
         parameters = {"rows": self._nodes_to_rows(nodes, lexical_graph_config)}
-        if self.is_version_5_23_or_above:
-            self.driver.execute_query(
-                UPSERT_NODE_QUERY_VARIABLE_SCOPE_CLAUSE,
-                parameters_=parameters,
-                database_=self.neo4j_database,
-            )
-        else:
-            self.driver.execute_query(
-                UPSERT_NODE_QUERY,
-                parameters_=parameters,
-                database_=self.neo4j_database,
-            )
+        query = (
+            UPSERT_NODE_QUERY_VARIABLE_SCOPE_CLAUSE
+            if self.is_version_5_23_or_above
+            else UPSERT_NODE_QUERY
+        )
+        records, _, _ = self.driver.execute_query(
+            query,
+            parameters_=parameters,
+            database_=self.neo4j_database,
+        )
+        print("RECORDS", records)
+        return {r["_internal_id"]: r["element_id"] for r in records}
 
-    def _upsert_relationships(self, rels: list[Neo4jRelationship]) -> None:
+    @staticmethod
+    def _relationships_to_rows(
+        relationships: list[Neo4jRelationship], node_id_mapping: dict[str, str]
+    ) -> list[dict[str, Any]]:
+        return [
+            {
+                **relationship.model_dump(),
+                "start_node_element_id": node_id_mapping.get(
+                    relationship.start_node_id, ""
+                ),
+                "end_node_element_id": node_id_mapping.get(
+                    relationship.end_node_id, ""
+                ),
+            }
+            for relationship in relationships
+        ]
+
+    def _upsert_relationships(
+        self, rels: list[Neo4jRelationship], node_id_mapping: dict[str, str]
+    ) -> None:
         """Upserts a single relationship into the Neo4j database.
 
         Args:
             rels (list[Neo4jRelationship]): The relationships batch to upsert into the database.
         """
-        parameters = {"rows": [rel.model_dump() for rel in rels]}
-        if self.is_version_5_23_or_above:
-            self.driver.execute_query(
-                UPSERT_RELATIONSHIP_QUERY_VARIABLE_SCOPE_CLAUSE,
-                parameters_=parameters,
-                database_=self.neo4j_database,
-            )
-        else:
-            self.driver.execute_query(
-                UPSERT_RELATIONSHIP_QUERY,
-                parameters_=parameters,
-                database_=self.neo4j_database,
-            )
+        parameters = {"rows": self._relationships_to_rows(rels, node_id_mapping)}
+        query = (
+            UPSERT_RELATIONSHIP_QUERY_VARIABLE_SCOPE_CLAUSE
+            if self.is_version_5_23_or_above
+            else UPSERT_RELATIONSHIP_QUERY
+        )
+        self.driver.execute_query(
+            query,
+            parameters_=parameters,
+            database_=self.neo4j_database,
+        )
 
     @validate_call
     async def run(
@@ -203,11 +216,14 @@ async def run(
         try:
             self._db_setup()
 
+            node_id_mapping = {}
+
             for batch in batched(graph.nodes, self.batch_size):
-                self._upsert_nodes(batch, lexical_graph_config)
+                batch_mapping = self._upsert_nodes(batch, lexical_graph_config)
+                node_id_mapping.update(batch_mapping)
 
             for batch in batched(graph.relationships, self.batch_size):
-                self._upsert_relationships(batch)
+                self._upsert_relationships(batch, node_id_mapping)
 
             return KGWriterModel(
                 status="SUCCESS",

src/neo4j_graphrag/neo4j_queries.py

@@ -54,7 +54,7 @@
 
 UPSERT_NODE_QUERY = (
     "UNWIND $rows AS row "
-    "CREATE (n:__KGBuilder__ {__kg_builder_id: row.id}) "
+    "CREATE (n:__KGBuilder__) "
     "SET n += row.properties "
     "WITH n, row CALL apoc.create.addLabels(n, row.labels) YIELD node "
     "WITH node as n, row CALL { "
@@ -63,12 +63,12 @@
     "CALL db.create.setNodeVectorProperty(n, emb, row.embedding_properties[emb]) "
     "RETURN count(*) as nbEmb "
     "} "
-    "RETURN elementId(n)"
+    "RETURN row.id as _internal_id, elementId(n) as element_id"
 )
 
 UPSERT_NODE_QUERY_VARIABLE_SCOPE_CLAUSE = (
     "UNWIND $rows AS row "
-    "CREATE (n:__KGBuilder__ {__kg_builder_id: row.id}) "
+    "CREATE (n:__KGBuilder__) "
     "SET n += row.properties "
     "WITH n, row CALL apoc.create.addLabels(n, row.labels) YIELD node "
     "WITH node as n, row CALL (n, row) { "
@@ -77,13 +77,13 @@
     "CALL db.create.setNodeVectorProperty(n, emb, row.embedding_properties[emb]) "
     "RETURN count(*) as nbEmb "
     "} "
-    "RETURN elementId(n)"
+    "RETURN row.id as _internal_id, elementId(n) as element_id"
 )
 
 UPSERT_RELATIONSHIP_QUERY = (
     "UNWIND $rows as row "
-    "MATCH (start:__KGBuilder__ {__kg_builder_id: row.start_node_id}) "
-    "MATCH (end:__KGBuilder__ {__kg_builder_id: row.end_node_id}) "
+    "MATCH (start:__KGBuilder__), (end:__KGBuilder__) "
+    "WHERE elementId(start) = row.start_node_element_id AND elementId(end) = row.end_node_element_id "
     "WITH start, end, row "
     "CALL apoc.merge.relationship(start, row.type, {}, row.properties, end, row.properties) YIELD rel  "
     "WITH rel, row CALL { "
@@ -96,8 +96,8 @@
 
 UPSERT_RELATIONSHIP_QUERY_VARIABLE_SCOPE_CLAUSE = (
     "UNWIND $rows as row "
-    "MATCH (start:__KGBuilder__ {__kg_builder_id: row.start_node_id}) "
-    "MATCH (end:__KGBuilder__ {__kg_builder_id: row.end_node_id}) "
+    "MATCH (start:__KGBuilder__), (end:__KGBuilder__) "
+    "WHERE elementId(start) = row.start_node_element_id AND elementId(end) = row.end_node_element_id "
     "WITH start, end, row "
     "CALL apoc.merge.relationship(start, row.type, {}, row.properties, end, row.properties) YIELD rel  "
     "WITH rel, row CALL (rel, row) { "

tests/unit/experimental/components/test_entity_relation_extractor.py

@@ -109,7 +109,6 @@ async def test_extractor_happy_path_non_empty_result() -> None:
     entity = result.nodes[2]
     assert entity.id == f"{chunk_entity.id}:0"
     assert entity.label == "Person"
-    assert entity.properties == {"chunk_index": 0}
     assert len(result.relationships) == 2
     assert result.relationships[0].type == "FROM_DOCUMENT"
     assert result.relationships[0].start_node_id == f"{chunk_entity.id}"
@@ -213,7 +212,6 @@ async def test_extractor_llm_badly_formatted_json_gets_fixed() -> None:
 
     assert len(res.nodes) == 1
     assert res.nodes[0].label == "Person"
-    assert res.nodes[0].properties == {"chunk_index": 0}
     assert res.nodes[0].embedding_properties is None
     assert res.relationships == []
 

tests/unit/experimental/components/test_kg_writer.py

@@ -56,9 +56,17 @@ def test_batched() -> None:
     return_value=None,
 )
 def test_upsert_nodes(_: Mock, driver: MagicMock) -> None:
+    driver.execute_query.return_value = (
+        [{"_internal_id": "1", "element_id": "#1"}],
+        None,
+        None,
+    )
     neo4j_writer = Neo4jWriter(driver=driver)
     node = Neo4jNode(id="1", label="Label", properties={"key": "value"})
-    neo4j_writer._upsert_nodes(nodes=[node], lexical_graph_config=LexicalGraphConfig())
+    result = neo4j_writer._upsert_nodes(
+        nodes=[node], lexical_graph_config=LexicalGraphConfig()
+    )
+    assert result == {"1": "#1"}
     driver.execute_query.assert_called_once_with(
         UPSERT_NODE_QUERY,
         parameters_={
@@ -88,14 +96,18 @@ def test_upsert_nodes_with_embedding(
     _: Mock,
     driver: MagicMock,
 ) -> None:
+    driver.execute_query.return_value = (
+        [{"_internal_id": "1", "element_id": "#1"}],
+        None,
+        None,
+    )
     neo4j_writer = Neo4jWriter(driver=driver)
     node = Neo4jNode(
         id="1",
         label="Label",
         properties={"key": "value"},
         embedding_properties={"embeddingProp": [1.0, 2.0, 3.0]},
     )
-    driver.execute_query.return_value.records = [{"elementId(n)": 1}]
     neo4j_writer._upsert_nodes(nodes=[node], lexical_graph_config=LexicalGraphConfig())
     driver.execute_query.assert_any_call(
         UPSERT_NODE_QUERY,
@@ -130,7 +142,9 @@ def test_upsert_relationship(_: Mock, driver: MagicMock) -> None:
         type="RELATIONSHIP",
         properties={"key": "value"},
     )
-    neo4j_writer._upsert_relationships(rels=[rel])
+    neo4j_writer._upsert_relationships(
+        rels=[rel], node_id_mapping={"1": "#1", "2": "#2"}
+    )
     parameters = {
         "rows": [
             {
@@ -139,6 +153,8 @@ def test_upsert_relationship(_: Mock, driver: MagicMock) -> None:
                 "end_node_id": "2",
                 "properties": {"key": "value"},
                 "embedding_properties": None,
+                "start_node_element_id": "#1",
+                "end_node_element_id": "#2",
             }
         ]
     }
@@ -167,7 +183,9 @@ def test_upsert_relationship_with_embedding(_: Mock, driver: MagicMock) -> None:
         embedding_properties={"embeddingProp": [1.0, 2.0, 3.0]},
     )
     driver.execute_query.return_value.records = [{"elementId(r)": "rel_elem_id"}]
-    neo4j_writer._upsert_relationships(rels=[rel])
+    neo4j_writer._upsert_relationships(
+        rels=[rel], node_id_mapping={"1": "#1", "2": "#2"}
+    )
     parameters = {
         "rows": [
             {
@@ -176,6 +194,8 @@ def test_upsert_relationship_with_embedding(_: Mock, driver: MagicMock) -> None:
                 "end_node_id": "2",
                 "properties": {"key": "value"},
                 "embedding_properties": {"embeddingProp": [1.0, 2.0, 3.0]},
+                "start_node_element_id": "#1",
+                "end_node_element_id": "#2",
             }
         ]
     }
@@ -196,6 +216,14 @@ def test_upsert_relationship_with_embedding(_: Mock, driver: MagicMock) -> None:
     return_value=None,
 )
 async def test_run(_: Mock, driver: MagicMock) -> None:
+    driver.execute_query.return_value = (
+        [
+            {"_internal_id": "1", "element_id": "#1"},
+            {"_internal_id": "2", "element_id": "#2"},
+        ],
+        None,
+        None,
+    )
     neo4j_writer = Neo4jWriter(driver=driver)
     node = Neo4jNode(id="1", label="Label")
     rel = Neo4jRelationship(start_node_id="1", end_node_id="2", type="RELATIONSHIP")
@@ -224,6 +252,8 @@ async def test_run(_: Mock, driver: MagicMock) -> None:
                 "end_node_id": "2",
                 "properties": {},
                 "embedding_properties": None,
+                "start_node_element_id": "#1",
+                "end_node_element_id": "#2",
             }
         ]
     }
@@ -242,7 +272,14 @@ async def test_run(_: Mock, driver: MagicMock) -> None:
 async def test_run_is_version_below_5_23(_: Mock) -> None:
     driver = MagicMock()
     driver.execute_query = Mock(
-        return_value=([{"versions": ["5.22.0"], "edition": "enterprise"}], None, None)
+        side_effect=(
+            # get_version
+            ([{"versions": ["5.22.0"], "edition": "enterpise"}], None, None),
+            # upsert nodes
+            ([{"_internal_id": "1", "element_id": "#1"}], None, None),
+            # upsert relationships
+            (None, None, None),
+        )
     )
 
     neo4j_writer = Neo4jWriter(driver=driver)
@@ -252,6 +289,8 @@ async def test_run_is_version_below_5_23(_: Mock) -> None:
     graph = Neo4jGraph(nodes=[node], relationships=[rel])
     await neo4j_writer.run(graph=graph)
 
+    print(driver.execute_query.call_args_list)
+
     driver.execute_query.assert_any_call(
         UPSERT_NODE_QUERY,
         parameters_={
@@ -275,6 +314,8 @@ async def test_run_is_version_below_5_23(_: Mock) -> None:
                 "end_node_id": "2",
                 "properties": {},
                 "embedding_properties": None,
+                "start_node_element_id": "#1",
+                "end_node_element_id": "",
             }
         ]
     }
@@ -293,7 +334,14 @@ async def test_run_is_version_below_5_23(_: Mock) -> None:
 async def test_run_is_version_5_23_or_above(_: Mock) -> None:
     driver = MagicMock()
     driver.execute_query = Mock(
-        return_value=([{"versions": ["5.23.0"], "edition": "enterpise"}], None, None)
+        side_effect=(
+            # get_version
+            ([{"versions": ["5.23.0"], "edition": "enterpise"}], None, None),
+            # upsert nodes
+            ([{"_internal_id": "1", "element_id": "#1"}], None, None),
+            # upsert relationships
+            (None, None, None),
+        )
     )
 
     neo4j_writer = Neo4jWriter(driver=driver)
@@ -327,6 +375,8 @@ async def test_run_is_version_5_23_or_above(_: Mock) -> None:
                 "end_node_id": "2",
                 "properties": {},
                 "embedding_properties": None,
+                "start_node_element_id": "#1",
+                "end_node_element_id": "",
             }
         ]
     }