Rename ID property and more cleaning (#366)

* Allow an ID property by renaming our internal ID to _kg_builder_id

* Prune nodes with invalid label/id and relationships with invalid type before inserting (APOC raises errors on them)

* Do not pollute the graph with internal/temp properties

* Fix e2e

* Improve e2e test

* Improve e2e test

* Rename property, clean db

* Doc update

* Review comments + e2e test

Author: stellasia

CHANGELOG.md

@@ -26,6 +26,12 @@
 - The `SchemaProperty` model has been renamed `PropertyType`.
 - `SchemaConfig` has been removed in favor of `GraphSchema` (used in the `SchemaBuilder` and `EntityRelationExtractor` classes). `entities`, `relations` and `potential_schema` fields have also been renamed `node_types`, `relationship_types` and `patterns` respectively.
 
+#### Other
+
+- The reserved `id` property on `__KGBuilder__` nodes is removed.
+- The `chunk_index` property on `__Entity__` nodes is removed. Use the `FROM_CHUNK` relationship instead.
+- The `__entity__id` index is not used anymore and can be dropped from the database (it has been replaced by `__entity__tmp_internal_id`).
+
 
 ## 1.7.0
 

docs/source/user_guide_kg_builder.rst

@@ -1048,8 +1048,10 @@ _________________
 In addition to the user-defined configuration options described above,
 the `GraphPruning` component performs the following cleanup operations:
 
+- Nodes with empty label or ID are pruned.
 - Nodes with missing required properties are pruned.
 - Nodes with no remaining properties are pruned.
+- Relationships with empty type are pruned.
 - Relationships with invalid source or target nodes (i.e., nodes no longer present in the graph) are pruned.
 - Relationships with incorrect direction have their direction corrected.
 
@@ -1078,6 +1080,11 @@ to a Neo4j database:
 Adjust the batch_size parameter of `Neo4jWriter` to optimize insert performance.
 This parameter controls the number of nodes or relationships inserted per batch, with a default value of 1000.
 
+.. note:: Index
+
+    In order to improve the ingestion performances, an index called `__entity__tmp_internal_id` is automatically added to the database.
+
+
 See :ref:`neo4jgraph`.
 
 

examples/customize/build_graph/components/pruners/graph_pruner.py

@@ -1,4 +1,4 @@
-"""This example demonstrates how to use the GraphPruner component."""
+"""This example demonstrates how to use the GraphPruning component."""
 
 import asyncio
 

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/graph_pruning.py

@@ -41,6 +41,7 @@ class PruningReason(str, enum.Enum):
     NO_PROPERTY_LEFT = "NO_PROPERTY_LEFT"
     INVALID_START_OR_END_NODE = "INVALID_START_OR_END_NODE"
     INVALID_PATTERN = "INVALID_PATTERN"
+    MISSING_LABEL = "MISSING_LABEL"
 
 
 ItemType = TypeVar("ItemType")
@@ -198,6 +199,17 @@ def _validate_node(
         schema_entity: Optional[NodeType],
         additional_node_types: bool,
     ) -> Optional[Neo4jNode]:
+        if not node.label:
+            pruning_stats.add_pruned_node(node, reason=PruningReason.MISSING_LABEL)
+            return None
+        if not node.id:
+            pruning_stats.add_pruned_node(
+                node,
+                reason=PruningReason.MISSING_REQUIRED_PROPERTY,
+                missing_required_properties=["id"],
+                details="The node was extracted without a valid ID.",
+            )
+            return None
         if not schema_entity:
             # node type not declared in the schema
             if additional_node_types:
@@ -262,6 +274,11 @@ def _validate_relationship(
         patterns: tuple[tuple[str, str, str], ...],
         additional_patterns: bool,
     ) -> Optional[Neo4jRelationship]:
+        if not rel.type:
+            pruning_stats.add_pruned_relationship(
+                rel, reason=PruningReason.MISSING_LABEL
+            )
+            return None
         # validate start/end node IDs are valid nodes
         if rel.start_node_id not in valid_nodes or rel.end_node_id not in valid_nodes:
             logger.debug(

src/neo4j_graphrag/experimental/components/kg_writer.py

@@ -29,10 +29,9 @@
 )
 from neo4j_graphrag.experimental.pipeline.component import Component, DataModel
 from neo4j_graphrag.neo4j_queries import (
-    UPSERT_NODE_QUERY,
-    UPSERT_NODE_QUERY_VARIABLE_SCOPE_CLAUSE,
-    UPSERT_RELATIONSHIP_QUERY,
-    UPSERT_RELATIONSHIP_QUERY_VARIABLE_SCOPE_CLAUSE,
+    upsert_node_query,
+    upsert_relationship_query,
+    db_cleaning_query,
 )
 from neo4j_graphrag.utils.version_utils import (
     get_version,
@@ -117,20 +116,19 @@ def __init__(
         driver: neo4j.Driver,
         neo4j_database: Optional[str] = None,
         batch_size: int = 1000,
+        clean_db: bool = True,
     ):
         self.driver = driver_config.override_user_agent(driver)
         self.neo4j_database = neo4j_database
         self.batch_size = batch_size
+        self._clean_db = clean_db
         version_tuple, _, _ = get_version(self.driver, self.neo4j_database)
         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,
-        )
+        self.driver.execute_query("""
+        CREATE INDEX __entity__tmp_internal_id IF NOT EXISTS FOR (n:__KGBuilder__) ON (n.__tmp_internal_id)
+        """)
 
     @staticmethod
     def _nodes_to_rows(
@@ -149,44 +147,51 @@ def _nodes_to_rows(
     def _upsert_nodes(
         self, nodes: list[Neo4jNode], lexical_graph_config: LexicalGraphConfig
     ) -> None:
-        """Upserts a single node into the Neo4j database."
+        """Upserts a batch of nodes 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(
+            support_variable_scope_clause=self.is_version_5_23_or_above
+        )
+        self.driver.execute_query(
+            query,
+            parameters_=parameters,
+            database_=self.neo4j_database,
+        )
+        return None
+
+    @staticmethod
+    def _relationships_to_rows(
+        relationships: list[Neo4jRelationship],
+    ) -> list[dict[str, Any]]:
+        return [relationship.model_dump() for relationship in relationships]
 
     def _upsert_relationships(self, rels: list[Neo4jRelationship]) -> None:
-        """Upserts a single relationship into the Neo4j database.
+        """Upserts a batch of relationships 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)}
+        query = upsert_relationship_query(
+            support_variable_scope_clause=self.is_version_5_23_or_above
+        )
+        self.driver.execute_query(
+            query,
+            parameters_=parameters,
+            database_=self.neo4j_database,
+        )
+
+    def _db_cleaning(self) -> None:
+        query = db_cleaning_query(
+            support_variable_scope_clause=self.is_version_5_23_or_above,
+            batch_size=self.batch_size,
+        )
+        with self.driver.session() as session:
+            session.run(query)
 
     @validate_call
     async def run(
@@ -209,6 +214,9 @@ async def run(
             for batch in batched(graph.relationships, self.batch_size):
                 self._upsert_relationships(batch)
 
+            if self._clean_db:
+                self._db_cleaning()
+
             return KGWriterModel(
                 status="SUCCESS",
                 metadata={

src/neo4j_graphrag/experimental/components/types.py

@@ -14,14 +14,18 @@
 #  limitations under the License.
 from __future__ import annotations
 
+import logging
 import uuid
 from typing import Any, Dict, Optional
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, Field
 
 from neo4j_graphrag.experimental.pipeline.component import DataModel
 
 
+logger = logging.getLogger(__name__)
+
+
 class DocumentInfo(DataModel):
     """A document loaded by a DataLoader.
 
@@ -79,7 +83,7 @@ class Neo4jNode(BaseModel):
     """Represents a Neo4j node.
 
     Attributes:
-        id (str): The element ID of the node.
+        id (str): The ID of the node. This ID is used to refer to the node for relationship creation.
         label (str): The label of the node.
         properties (dict[str, Any]): A dictionary of properties attached to the node.
         embedding_properties (Optional[dict[str, list[float]]]): A list of embedding properties attached to the node.
@@ -90,15 +94,6 @@ class Neo4jNode(BaseModel):
     properties: dict[str, Any] = {}
     embedding_properties: Optional[dict[str, list[float]]] = None
 
-    @field_validator("properties", "embedding_properties")
-    @classmethod
-    def check_for_id_properties(
-        cls, v: Optional[dict[str, Any]]
-    ) -> Optional[dict[str, Any]]:
-        if v and "id" in v.keys():
-            raise TypeError("'id' as a property name is not allowed")
-        return v
-
     @property
     def token(self) -> str:
         return self.label