WIP: apply constraints from DB

Author: stellasia

examples/customize/build_graph/components/schema_builders/schema.py

@@ -12,48 +12,59 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+import asyncio
+
+import neo4j
+
 from neo4j_graphrag.experimental.components.schema import (
     SchemaBuilder,
 )
-from neo4j_graphrag.experimental.components.types import (
-    NodeType,
-    PropertyType,
-    RelationshipType,
-)
 
 
 async def main() -> None:
-    schema_builder = SchemaBuilder()
+    with neo4j.GraphDatabase.driver(
+        "bolt://localhost:7687",
+        auth=("neo4j", "password"),
+    ) as driver:
+        schema_builder = SchemaBuilder(driver)
+
+        schema = await schema_builder.run(
+            node_types=[
+                {
+                    "label": "Person",
+                    "properties": [
+                        {"name": "name", "type": "STRING"},
+                        {"name": "place_of_birth", "type": "STRING"},
+                        {"name": "date_of_birth", "type": "DATE"},
+                    ],
+                },
+                {
+                    "label": "Organization",
+                    "properties": [
+                        {"name": "name", "type": "STRING"},
+                        {"name": "country", "type": "STRING"},
+                    ],
+                },
+                {
+                    "label": "Field",
+                    "properties": [
+                        {"name": "name", "type": "STRING"},
+                    ],
+                },
+            ],
+            relationship_types=[
+                "WORKED_ON",
+                {
+                    "label": "WORKED_FOR",
+                },
+            ],
+            patterns=[
+                ("Person", "WORKED_ON", "Field"),
+                ("Person", "WORKED_FOR", "Organization"),
+            ],
+        )
+        print(schema)
+
 
-    result = await schema_builder.run(
-        node_types=[
-            NodeType(
-                label="Person",
-                properties=[
-                    PropertyType(name="name", type="STRING"),
-                    PropertyType(name="place_of_birth", type="STRING"),
-                    PropertyType(name="date_of_birth", type="DATE"),
-                ],
-            ),
-            NodeType(
-                label="Organization",
-                properties=[
-                    PropertyType(name="name", type="STRING"),
-                    PropertyType(name="country", type="STRING"),
-                ],
-            ),
-        ],
-        relationship_types=[
-            RelationshipType(
-                label="WORKED_ON",
-            ),
-            RelationshipType(
-                label="WORKED_FOR",
-            ),
-        ],
-        patterns=[
-            ("Person", "WORKED_ON", "Field"),
-            ("Person", "WORKED_FOR", "Organization"),
-        ],
-    )
-    print(result)
+if __name__ == "__main__":
+    asyncio.run(main())

src/neo4j_graphrag/experimental/components/entity_relation_extractor.py

@@ -25,10 +25,10 @@
 
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.experimental.components.lexical_graph import LexicalGraphBuilder
-from neo4j_graphrag.experimental.components.schema import GraphSchema
 from neo4j_graphrag.experimental.components.types import (
     DocumentInfo,
     LexicalGraphConfig,
+    GraphSchema,
     Neo4jGraph,
     TextChunk,
     TextChunks,

src/neo4j_graphrag/experimental/components/schema.py

@@ -14,10 +14,12 @@
 #  limitations under the License.
 from __future__ import annotations
 
+import copy
 import json
 import logging
-from typing import Any, Dict, List, Optional, Tuple, Sequence
+from typing import Any, Dict, List, Optional, Tuple, Sequence, Literal
 
+import neo4j
 from pydantic import (
     validate_call,
     ValidationError,
@@ -29,13 +31,22 @@
     SchemaExtractionError,
 )
 from neo4j_graphrag.experimental.pipeline.component import Component
+from neo4j_graphrag.experimental.pipeline.types.schema import (
+    EntityInputType,
+    RelationInputType,
+)
 from neo4j_graphrag.generation import SchemaExtractionTemplate, PromptTemplate
 from neo4j_graphrag.llm import LLMInterface
 from neo4j_graphrag.experimental.components.types import (
-    NodeType,
     RelationshipType,
     GraphSchema,
+    SchemaConstraint,
+    ConstraintTypeEnum,
+    Neo4jConstraintTypeEnum,
+    GraphEntityType,
+    Neo4jPropertyType,
 )
+from neo4j_graphrag.schema import get_constraints
 
 
 class SchemaBuilder(Component):
@@ -97,9 +108,83 @@ class SchemaBuilder(Component):
         pipe.run(pipe_inputs)
     """
 
+    def __init__(
+        self, driver: neo4j.Driver, neo4j_database: Optional[str] = None
+    ) -> None:
+        self.driver = driver
+        self.neo4j_database = neo4j_database
+
+    def _get_constraints_from_db(self) -> list[dict[str, Any]]:
+        constraints = get_constraints(
+            self.driver, database=self.neo4j_database, sanitize=False
+        )
+        return constraints
+
+    def _apply_all_constraints_from_db(
+        self,
+        node_or_relationship_type: Literal["NODE", "RELATIONSHIP"],
+        constraints: list[dict[str, Any]],
+        entities: tuple[GraphEntityType, ...],
+    ) -> list[GraphEntityType]:
+        constrained_entity_types = []
+        for entity_type in entities:
+            new_entity_type = copy.deepcopy(entity_type)
+            # find constraints related to this node type
+            for constraint in constraints:
+                if constraint["entityType"] != node_or_relationship_type:
+                    continue
+                if constraint["labelsOrTypes"][0] != entity_type.label:
+                    continue
+                # now we can add the constraint to this node type
+                self._apply_constraint_from_db(new_entity_type, constraint)
+            constrained_entity_types.append(new_entity_type)
+        return constrained_entity_types
+
     @staticmethod
-    def create_schema_model(
-        node_types: Sequence[NodeType],
+    def _parse_property_type(property_type: str) -> Neo4jPropertyType | None:
+        if not property_type:
+            return None
+        prop = None
+        for prop_str in property_type.split("|"):
+            p = prop_str.strip()
+            try:
+                prop = Neo4jPropertyType(p)
+            except ValueError:
+                pass
+        return prop
+
+    def _apply_constraint_from_db(
+        self, entity_type: GraphEntityType, constraint: dict[str, Any]
+    ) -> None:
+        neo4j_constraint_type = Neo4jConstraintTypeEnum(constraint["type"])
+        # TODO: detect potential conflict and raise ValueError if any
+        # existing_schema_constraints_on_property = node_type.get_constraints_on_properties(constraint["properties"])
+        constraint_properties = constraint["properties"]
+        for p in constraint_properties:
+            if entity_type.get_property_by_name(p) is None:
+                raise ValueError(
+                    f"Can not add constraint {constraint} on non existing property"
+                )
+        constraint_type = neo4j_constraint_type.to_constraint_type()
+        entity_type.constraints.append(
+            SchemaConstraint(
+                type=constraint_type,
+                properties=constraint["properties"],
+                property_type=self._parse_property_type(constraint["propertyType"]),
+                name=constraint["name"],
+            )
+        )
+        # if property required constraint, make sure the flag is set properly on
+        # the PropertyType
+        if constraint_type == ConstraintTypeEnum.PROPERTY_EXISTENCE:
+            prop = entity_type.get_property_by_name(constraint["properties"][0])
+            if prop:
+                prop.required = True
+        return None
+
+    def _create_schema_model(
+        self,
+        node_types: Sequence[EntityInputType],
         relationship_types: Optional[Sequence[RelationshipType]] = None,
         patterns: Optional[Sequence[Tuple[str, str, str]]] = None,
         **kwargs: Any,
@@ -118,7 +203,7 @@ def create_schema_model(
             GraphSchema: A configured schema object.
         """
         try:
-            return GraphSchema.model_validate(
+            schema = GraphSchema.model_validate(
                 dict(
                     node_types=node_types,
                     relationship_types=relationship_types or (),
@@ -129,11 +214,39 @@ def create_schema_model(
         except ValidationError as e:
             raise SchemaValidationError() from e
 
+        constraints = self._get_constraints_from_db()
+        # apply constraints
+        constrained_node_types = self._apply_all_constraints_from_db(
+            "NODE",
+            constraints,
+            schema.node_types,
+        )
+        constrained_relationship_types = self._apply_all_constraints_from_db(
+            "RELATIONSHIP",
+            constraints,
+            schema.relationship_types,
+        )
+
+        try:
+            constrained_schema = GraphSchema.model_validate(
+                dict(
+                    node_types=constrained_node_types,
+                    relationship_types=constrained_relationship_types,
+                    patterns=patterns,
+                    **kwargs,
+                )
+            )
+        except ValidationError as e:
+            raise SchemaValidationError(
+                "Error when applying constraints from database"
+            ) from e
+        return constrained_schema
+
     @validate_call
     async def run(
         self,
-        node_types: Sequence[NodeType],
-        relationship_types: Optional[Sequence[RelationshipType]] = None,
+        node_types: Sequence[EntityInputType],
+        relationship_types: Optional[Sequence[RelationInputType]] = None,
         patterns: Optional[Sequence[Tuple[str, str, str]]] = None,
         **kwargs: Any,
     ) -> GraphSchema:
@@ -148,7 +261,7 @@ async def run(
         Returns:
             GraphSchema: A configured schema object, constructed asynchronously.
         """
-        return self.create_schema_model(
+        return self._create_schema_model(
             node_types,
             relationship_types,
             patterns,
@@ -164,10 +277,12 @@ class SchemaFromTextExtractor(Component):
 
     def __init__(
         self,
+        driver: neo4j.Driver,
         llm: LLMInterface,
         prompt_template: Optional[PromptTemplate] = None,
         llm_params: Optional[Dict[str, Any]] = None,
     ) -> None:
+        self.driver = driver
         self._llm: LLMInterface = llm
         self._prompt_template: PromptTemplate = (
             prompt_template or SchemaExtractionTemplate()

src/neo4j_graphrag/experimental/components/types.py

@@ -232,10 +232,6 @@ class PropertyType(BaseModel):
     description: str = ""
     required: bool = False
 
-    model_config = ConfigDict(
-        frozen=True,
-    )
-
 
 class ConstraintTypeEnum(str, enum.Enum):
     # see: https://neo4j.com/docs/cypher-manual/current/constraints/
@@ -245,6 +241,44 @@ class ConstraintTypeEnum(str, enum.Enum):
     PROPERTY_TYPE = "PROPERTY_TYPE"
 
 
+class Neo4jConstraintTypeEnum(str, enum.Enum):
+    NODE_KEY = "NODE_KEY"
+    UNIQUENESS = "UNIQUENESS"
+    NODE_PROPERTY_EXISTENCE = "NODE_PROPERTY_EXISTENCE"
+    NODE_PROPERTY_UNIQUENESS = "NODE_PROPERTY_UNIQUENESS"
+    NODE_PROPERTY_TYPE = "NODE_PROPERTY_TYPE"
+    RELATIONSHIP_KEY = "RELATIONSHIP_KEY"
+    RELATIONSHIP_UNIQUENESS = "RELATIONSHIP_UNIQUENESS"
+    RELATIONSHIP_PROPERTY_EXISTENCE = "RELATIONSHIP_PROPERTY_EXISTENCE"
+    RELATIONSHIP_PROPERTY_UNIQUENESS = "RELATIONSHIP_PROPERTY_UNIQUENESS"
+    RELATIONSHIP_PROPERTY_TYPE = "RELATIONSHIP_PROPERTY_TYPE"
+
+    def to_constraint_type(self) -> ConstraintTypeEnum:
+        if self in (
+            Neo4jConstraintTypeEnum.NODE_KEY,
+            Neo4jConstraintTypeEnum.RELATIONSHIP_KEY,
+        ):
+            return ConstraintTypeEnum.KEY
+        if self in (
+            Neo4jConstraintTypeEnum.UNIQUENESS,
+            Neo4jConstraintTypeEnum.NODE_PROPERTY_UNIQUENESS,
+            Neo4jConstraintTypeEnum.RELATIONSHIP_UNIQUENESS,
+            Neo4jConstraintTypeEnum.RELATIONSHIP_PROPERTY_UNIQUENESS,
+        ):
+            return ConstraintTypeEnum.PROPERTY_UNIQUENESS
+        if self in (
+            Neo4jConstraintTypeEnum.NODE_PROPERTY_EXISTENCE,
+            Neo4jConstraintTypeEnum.RELATIONSHIP_PROPERTY_EXISTENCE,
+        ):
+            return ConstraintTypeEnum.PROPERTY_EXISTENCE
+        if self in (
+            Neo4jConstraintTypeEnum.NODE_PROPERTY_TYPE,
+            Neo4jConstraintTypeEnum.RELATIONSHIP_PROPERTY_TYPE,
+        ):
+            return ConstraintTypeEnum.PROPERTY_TYPE
+        raise ValueError(f"Can't convert {self} to ConstraintTypeEnum")
+
+
 class SchemaConstraint(BaseModel):
     """Constraints that can be applied either on a node or relationship property."""
 
@@ -290,6 +324,21 @@ def validate_constraint_on_properties(self) -> Self:
                     )
         return self
 
+    def get_property_by_name(self, name: str) -> PropertyType | None:
+        for prop in self.properties:
+            if prop.name == name:
+                return prop
+        return None
+
+    def get_constraints_on_properties(
+        self, prop_names: list[str]
+    ) -> list[SchemaConstraint]:
+        constraints = []
+        for constraint in self.constraints:
+            if set(prop_names) == set(constraint.properties):
+                constraints.append(constraint)
+        return constraints
+
     def get_unique_properties(self) -> list[str]:
         for c in self.constraints:
             if c.type in (