feat: add tests for ConstraintType

Author: AmirLayegh

src/neo4j_graphrag/experimental/components/schema.py

@@ -76,7 +76,7 @@ class PropertyType(BaseModel):
     ]
     description: str = ""
     required: bool = False
-
+    # unique: bool = False
     model_config = ConfigDict(
         frozen=True,
     )
@@ -269,6 +269,16 @@ def validate_constraints_against_node_types(self) -> Self:
                 raise SchemaValidationError(
                     f"Constraint references undefined node type: {constraint.node_type}"
                 )
+            # Check if property_name exists on the node type (only if additional_properties is False)
+            node_type = self._node_type_index[constraint.node_type]
+            if not node_type.additional_properties:
+                valid_property_names = {p.name for p in node_type.properties}
+                if constraint.property_name not in valid_property_names:
+                    raise SchemaValidationError(
+                        f"Constraint references undefined property '{constraint.property_name}' "
+                        f"on node type '{constraint.node_type}'. "
+                        f"Valid properties: {valid_property_names}"
+                    )
         return self
 
     def node_type_from_label(self, label: str) -> Optional[NodeType]:
@@ -594,7 +604,8 @@ def _filter_relationships_without_labels(
     def _filter_invalid_constraints(
         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."""
+        """Filter out constraints that reference undefined node types, have no property name,
+        or reference a property that doesn't exist on the node type."""
         if not constraints:
             return []
 
@@ -605,7 +616,16 @@ def _filter_invalid_constraints(
             )
             return []
 
-        valid_node_labels = {node_type.get("label") for node_type in node_types}
+        # Build a mapping of node_type label -> set of property names
+        node_type_properties: Dict[str, set[str]] = {}
+        for node_type_dict in node_types:
+            label = node_type_dict.get("label")
+            if label:
+                properties = node_type_dict.get("properties", [])
+                property_names = {p.get("name") for p in properties if p.get("name")}
+                node_type_properties[label] = property_names
+
+        valid_node_labels = set(node_type_properties.keys())
 
         filtered_constraints = []
         for constraint in constraints:
@@ -617,10 +637,20 @@ def _filter_invalid_constraints(
                 )
                 continue
             # check if the node_type is valid
-            if constraint.get("node_type") not in valid_node_labels:
+            node_type = constraint.get("node_type")
+            if node_type not in valid_node_labels:
+                logging.info(
+                    f"Filtering out constraint: {constraint}. "
+                    f"Node type '{node_type}' is not valid. Valid node types: {valid_node_labels}"
+                )
+                continue
+            # check if the property_name exists on the node type
+            property_name = constraint.get("property_name")
+            if property_name not in node_type_properties.get(node_type, set()):
                 logging.info(
                     f"Filtering out constraint: {constraint}. "
-                    f"Node type '{constraint.get('node_type')}' is not valid. Valid node types: {valid_node_labels}"
+                    f"Property '{property_name}' does not exist on node type '{node_type}'. "
+                    f"Valid properties: {node_type_properties.get(node_type, set())}"
                 )
                 continue
             filtered_constraints.append(constraint)

src/neo4j_graphrag/generation/prompts.py

@@ -217,10 +217,12 @@ class SchemaExtractionTemplate(PromptTemplate):
 4. Include property definitions only when the type can be confidently inferred, otherwise omit them.
 5. When defining patterns, ensure that every node label and relationship label mentioned exists in your lists of node types and relationship types.
 6. Do not create node types that aren't clearly mentioned in the text.
-7. For each node type, identify a unique identifier property and add it as a UNIQUENESS constraint to the list of constraints.
-8. Constraints must reference a node_type label that exists in the list of node types.
-9. Each constraint must have a property_name having a name that indicates it is a unique identifier for the node type (e.g., person_id for Person, company_id for Company)
-10. Keep your schema minimal and focused on clearly identifiable patterns in the text.
+7. Keep your schema minimal and focused on clearly identifiable patterns in the text.
+8. UNIQUENESS CONSTRAINTS:
+8.1 UNIQUENESS is optional; each node_type may or may not have exactly one uniqueness constraint.
+8.2 Only use properties that seem to not have too many missing values in the sample.
+8.3 Constraints reference node_types by label and specify which property is unique.
+8.4 If a property appears in a uniqueness constraint it MUST also appear in the corresponding node_type as a property.
 
 
 Accepted property types are: BOOLEAN, DATE, DURATION, FLOAT, INTEGER, LIST,
@@ -254,7 +256,7 @@ class SchemaExtractionTemplate(PromptTemplate):
     {{
       "type": "UNIQUENESS",
       "node_type": "Person",
-      "property_name": "person_id"
+      "property_name": "name"
     }}
     ...
   ]