feat: add ConstraintType to GraphSchema for constraint extraction

Author: AmirLayegh

src/neo4j_graphrag/experimental/components/schema.py

@@ -160,12 +160,15 @@ def property_type_from_name(self, name: str) -> Optional[PropertyType]:
                 return prop
         return None
 
+
 class ConstraintType(BaseModel):
     """
     Represents a constraint on a node in the graph.
     """
 
-    type: Literal["UNIQUENESS"]  #TODO: add other constraint types ["propertyExistence", "propertyType", "key"]
+    type: Literal[
+        "UNIQUENESS"
+    ]  # TODO: add other constraint types ["propertyExistence", "propertyType", "key"]
     node_type: str
     property_name: str
 
@@ -258,11 +261,11 @@ def validate_constraints_against_node_types(self) -> Self:
         if not self.constraints:
             return self
         for constraint in self.constraints:
-            if not constraint.get('property_name'):
+            if not constraint.get("property_name"):
                 raise SchemaValidationError(
                     f"Constraint has no property name: {constraint}. Property name is required."
                 )
-            if constraint.get('node_type') not in self._node_type_index:
+            if constraint.get("node_type") not in self._node_type_index:
                 raise SchemaValidationError(
                     f"Constraint references undefined node type: {constraint.get('node_type')}"
                 )
@@ -587,35 +590,34 @@ def _filter_relationships_without_labels(
         return self._filter_items_without_labels(
             relationship_types, "relationship type"
         )
-    
+
     def _filter_invalid_constraints(
-        self, 
-        constraints: List[Dict[str, Any]], node_types: List[Dict[str, Any]]
+        self, constraints: List[Dict[str, Any]], node_types: List[Dict[str, Any]]
     ) -> List[Dict[str, Any]]:
         """Filter out constraints that reference undefined node types or have no property name."""
         if not constraints:
             return []
-        
+
         if not node_types:
             logging.info(
                 "Filtering out all constraints because no node types are defined. "
                 "Constraints reference node types that must be defined."
             )
             return []
-        
-        valid_node_labels = {node_type.get('label') for node_type in node_types}
+
+        valid_node_labels = {node_type.get("label") for node_type in node_types}
 
         filtered_constraints = []
         for constraint in constraints:
             # check if the property_name is provided
-            if not constraint.get('property_name'):
+            if not constraint.get("property_name"):
                 logging.info(
                     f"Filtering out constraint: {constraint}. "
                     f"Property name is not provided."
                 )
                 continue
             # check if the node_type is valid
-            if constraint.get('node_type') not in valid_node_labels:
+            if constraint.get("node_type") not in valid_node_labels:
                 logging.info(
                     f"Filtering out constraint: {constraint}. "
                     f"Node type '{constraint.get('node_type')}' is not valid. Valid node types: {valid_node_labels}"
@@ -709,7 +711,7 @@ async def run(self, text: str, examples: str = "", **kwargs: Any) -> GraphSchema
             extracted_patterns = self._filter_invalid_patterns(
                 extracted_patterns, extracted_node_types, extracted_relationship_types
             )
-        
+
         # Filter out invalid constraints
         if extracted_constraints:
             extracted_constraints = self._filter_invalid_constraints(