Add Constraint to GraphSchema - move types in the types.py file

Author: stellasia

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

@@ -3,13 +3,11 @@
 import asyncio
 
 from neo4j_graphrag.experimental.components.graph_pruning import GraphPruning
-from neo4j_graphrag.experimental.components.schema import (
+from neo4j_graphrag.experimental.components.types import (
     GraphSchema,
     NodeType,
     PropertyType,
     RelationshipType,
-)
-from neo4j_graphrag.experimental.components.types import (
     Neo4jGraph,
     Neo4jNode,
     Neo4jRelationship,

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

@@ -14,6 +14,8 @@
 #  limitations under the License.
 from neo4j_graphrag.experimental.components.schema import (
     SchemaBuilder,
+)
+from neo4j_graphrag.experimental.components.types import (
     NodeType,
     PropertyType,
     RelationshipType,

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

@@ -14,6 +14,8 @@
 
 from neo4j_graphrag.experimental.components.schema import (
     SchemaFromTextExtractor,
+)
+from neo4j_graphrag.experimental.components.types import (
     GraphSchema,
 )
 from neo4j_graphrag.llm import OpenAILLM

src/neo4j_graphrag/experimental/components/graph_pruning.py

@@ -18,7 +18,7 @@
 
 from pydantic import validate_call, BaseModel
 
-from neo4j_graphrag.experimental.components.schema import (
+from neo4j_graphrag.experimental.components.types import (
     GraphSchema,
     PropertyType,
     NodeType,

src/neo4j_graphrag/experimental/components/schema.py

@@ -16,258 +16,26 @@
 
 import json
 import logging
-import warnings
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union, Sequence
-from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Sequence
 
 from pydantic import (
-    BaseModel,
-    PrivateAttr,
-    model_validator,
     validate_call,
-    ConfigDict,
     ValidationError,
 )
-from typing_extensions import Self
 
 from neo4j_graphrag.exceptions import (
     SchemaValidationError,
     LLMGenerationError,
     SchemaExtractionError,
 )
-from neo4j_graphrag.experimental.pipeline.component import Component, DataModel
-from neo4j_graphrag.experimental.pipeline.types.schema import (
-    EntityInputType,
-    RelationInputType,
-)
+from neo4j_graphrag.experimental.pipeline.component import Component
 from neo4j_graphrag.generation import SchemaExtractionTemplate, PromptTemplate
 from neo4j_graphrag.llm import LLMInterface
-from neo4j_graphrag.utils.file_handler import FileHandler, FileFormat
-
-
-class PropertyType(BaseModel):
-    """
-    Represents a property on a node or relationship in the graph.
-    """
-
-    name: str
-    # See https://neo4j.com/docs/cypher-manual/current/values-and-types/property-structural-constructed/#property-types
-    type: Literal[
-        "BOOLEAN",
-        "DATE",
-        "DURATION",
-        "FLOAT",
-        "INTEGER",
-        "LIST",
-        "LOCAL_DATETIME",
-        "LOCAL_TIME",
-        "POINT",
-        "STRING",
-        "ZONED_DATETIME",
-        "ZONED_TIME",
-    ]
-    description: str = ""
-    required: bool = False
-
-    model_config = ConfigDict(
-        frozen=True,
-    )
-
-
-class NodeType(BaseModel):
-    """
-    Represents a possible node in the graph.
-    """
-
-    label: str
-    description: str = ""
-    properties: list[PropertyType] = []
-    additional_properties: bool = True
-
-    @model_validator(mode="before")
-    @classmethod
-    def validate_input_if_string(cls, data: EntityInputType) -> EntityInputType:
-        if isinstance(data, str):
-            return {"label": data}
-        return data
-
-    @model_validator(mode="after")
-    def validate_additional_properties(self) -> Self:
-        if len(self.properties) == 0 and not self.additional_properties:
-            raise ValueError(
-                "Using `additional_properties=False` with no defined "
-                "properties will cause the model to be pruned during graph cleaning.",
-            )
-        return self
-
-
-class RelationshipType(BaseModel):
-    """
-    Represents a possible relationship between nodes in the graph.
-    """
-
-    label: str
-    description: str = ""
-    properties: list[PropertyType] = []
-    additional_properties: bool = True
-
-    @model_validator(mode="before")
-    @classmethod
-    def validate_input_if_string(cls, data: RelationInputType) -> RelationInputType:
-        if isinstance(data, str):
-            return {"label": data}
-        return data
-
-    @model_validator(mode="after")
-    def validate_additional_properties(self) -> Self:
-        if len(self.properties) == 0 and not self.additional_properties:
-            raise ValueError(
-                "Using `additional_properties=False` with no defined "
-                "properties will cause the model to be pruned during graph cleaning.",
-            )
-        return self
-
-
-class GraphSchema(DataModel):
-    """This model represents the expected
-    node and relationship types in the graph.
-
-    It is used both for guiding the LLM in the entity and relation
-    extraction component, and for cleaning the extracted graph in a
-    post-processing step.
-
-    .. warning::
-
-        This model is immutable.
-    """
-
-    node_types: Tuple[NodeType, ...]
-    relationship_types: Tuple[RelationshipType, ...] = tuple()
-    patterns: Tuple[Tuple[str, str, str], ...] = tuple()
-
-    additional_node_types: bool = True
-    additional_relationship_types: bool = True
-    additional_patterns: bool = True
-
-    _node_type_index: dict[str, NodeType] = PrivateAttr()
-    _relationship_type_index: dict[str, RelationshipType] = PrivateAttr()
-
-    model_config = ConfigDict(
-        frozen=True,
-    )
-
-    @model_validator(mode="after")
-    def validate_patterns_against_node_and_rel_types(self) -> Self:
-        self._node_type_index = {node.label: node for node in self.node_types}
-        self._relationship_type_index = (
-            {r.label: r for r in self.relationship_types}
-            if self.relationship_types
-            else {}
-        )
-
-        relationship_types = self.relationship_types
-        patterns = self.patterns
-
-        if patterns:
-            if not relationship_types:
-                raise SchemaValidationError(
-                    "Relationship types must also be provided when using patterns."
-                )
-            for entity1, relation, entity2 in patterns:
-                if entity1 not in self._node_type_index:
-                    raise SchemaValidationError(
-                        f"Node type '{entity1}' is not defined in the provided node_types."
-                    )
-                if relation not in self._relationship_type_index:
-                    raise SchemaValidationError(
-                        f"Relationship type '{relation}' is not defined in the provided relationship_types."
-                    )
-                if entity2 not in self._node_type_index:
-                    raise ValueError(
-                        f"Node type '{entity2}' is not defined in the provided node_types."
-                    )
-
-        return self
-
-    @model_validator(mode="after")
-    def validate_additional_parameters(self) -> Self:
-        if (
-            self.additional_patterns is False
-            and self.additional_relationship_types is True
-        ):
-            raise ValueError(
-                "`additional_relationship_types` must be set to False when using `additional_patterns=False`"
-            )
-        return self
-
-    def node_type_from_label(self, label: str) -> Optional[NodeType]:
-        return self._node_type_index.get(label)
-
-    def relationship_type_from_label(self, label: str) -> Optional[RelationshipType]:
-        return self._relationship_type_index.get(label)
-
-    def save(
-        self,
-        file_path: Union[str, Path],
-        overwrite: bool = False,
-        format: Optional[FileFormat] = None,
-    ) -> None:
-        """
-        Save the schema configuration to file.
-
-        Args:
-            file_path (str): The path where the schema configuration will be saved.
-            overwrite (bool): If set to True, existing file will be overwritten. Default to False.
-            format (Optional[FileFormat]): The file format to save the schema configuration into. By default, it is inferred from file_path extension.
-        """
-        data = self.model_dump(mode="json")
-        file_handler = FileHandler()
-        file_handler.write(data, file_path, overwrite=overwrite, format=format)
-
-    def store_as_json(
-        self, file_path: Union[str, Path], overwrite: bool = False
-    ) -> None:
-        warnings.warn(
-            "Use .save(..., format=FileFormat.JSON) instead.", DeprecationWarning
-        )
-        return self.save(file_path, overwrite=overwrite, format=FileFormat.JSON)
-
-    def store_as_yaml(
-        self, file_path: Union[str, Path], overwrite: bool = False
-    ) -> None:
-        warnings.warn(
-            "Use .save(..., format=FileFormat.YAML) instead.", DeprecationWarning
-        )
-        return self.save(file_path, overwrite=overwrite, format=FileFormat.YAML)
-
-    @classmethod
-    def from_file(
-        cls, file_path: Union[str, Path], format: Optional[FileFormat] = None
-    ) -> Self:
-        """
-        Load a schema configuration from a file (either JSON or YAML).
-
-        The file format is automatically detected based on the file extension,
-        unless the format parameter is set.
-
-        Args:
-            file_path (Union[str, Path]): The path to the schema configuration file.
-            format (Optional[FileFormat]): The format of the schema configuration file (json or yaml).
-
-        Returns:
-            GraphSchema: The loaded schema configuration.
-        """
-        file_path = Path(file_path)
-        file_handler = FileHandler()
-        try:
-            data = file_handler.read(file_path, format=format)
-        except ValueError:
-            raise
-
-        try:
-            return cls.model_validate(data)
-        except ValidationError as e:
-            raise SchemaValidationError(str(e)) from e
+from neo4j_graphrag.experimental.components.types import (
+    NodeType,
+    RelationshipType,
+    GraphSchema,
+)
 
 
 class SchemaBuilder(Component):