add required support to SchemaFromTextExtractor

Author: AmirLayegh

src/neo4j_graphrag/experimental/components/schema.py

@@ -725,6 +725,40 @@ def _filter_properties_required_field(
 
         return node_types
 
+    def _enforce_required_for_constraint_properties(
+        self,
+        node_types: List[Dict[str, Any]],
+        constraints: List[Dict[str, Any]],
+    ) -> None:
+        """Ensure properties with UNIQUENESS constraints are marked as required."""
+        if not constraints:
+            return
+
+        # Build a lookup for property_names and constraints
+        constraint_props: Dict[str, set[str]] = {}
+        for c in constraints:
+            if c.get("type") == "UNIQUENESS":
+                label = c.get("node_type")
+                prop = c.get("property_name")
+                if label and prop:
+                    constraint_props.setdefault(label, set()).add(prop)
+
+        # Skop node_types without constraints
+        for node_type in node_types:
+            label = node_type.get("label")
+            if label not in constraint_props:
+                continue
+
+            props_to_fix = constraint_props[label]
+            for prop in node_type.get("properties", []):
+                if isinstance(prop, dict) and prop.get("name") in props_to_fix:
+                    if prop.get("required") is not True:
+                        logging.info(
+                            f"Auto-setting 'required' as True for property '{prop.get('name')}' "
+                            f"on node '{label}' (has UNIQUENESS constraint)."
+                        )
+                        prop["required"] = True
+
     def _clean_json_content(self, content: str) -> str:
         content = content.strip()
 
@@ -815,6 +849,12 @@ async def run(self, text: str, examples: str = "", **kwargs: Any) -> GraphSchema
                 extracted_patterns, extracted_node_types, extracted_relationship_types
             )
 
+        # Enforce required=true for properties with UNIQUENESS constraints
+        if extracted_constraints:
+            self._enforce_required_for_constraint_properties(
+                extracted_node_types, extracted_constraints
+            )
+
         # Filter out invalid constraints
         if extracted_constraints:
             extracted_constraints = self._filter_invalid_constraints(

src/neo4j_graphrag/generation/prompts.py

@@ -229,6 +229,7 @@ class SchemaExtractionTemplate(PromptTemplate):
 9.3 Properties that are identifiers, names, or essential characteristics are typically required.
 9.4 Properties that are supplementary information (phone numbers, descriptions, metadata) are typically optional.
 9.5 When uncertain, default to "required": false.
+9.6 If a property has a UNIQUENESS constraint, it MUST be marked as "required": true.
 
 Accepted property types are: BOOLEAN, DATE, DURATION, FLOAT, INTEGER, LIST,
 LOCAL_DATETIME, LOCAL_TIME, POINT, STRING, ZONED_DATETIME, ZONED_TIME.

tests/unit/experimental/components/test_schema.py

@@ -1567,6 +1567,127 @@ def test_filter_properties_required_field_missing(
     assert "required" not in result[0]["properties"][0]
 
 
+def test_enforce_required_for_constraint_properties_sets_required_true(
+    schema_from_text: SchemaFromTextExtractor,
+) -> None:
+    node_types: list[dict[str, Any]] = [
+        {
+            "label": "Person",
+            "properties": [
+                {"name": "name", "type": "STRING", "required": False},
+                {"name": "email", "type": "STRING", "required": False},
+            ],
+        }
+    ]
+    constraints = [
+        {"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"}
+    ]
+
+    schema_from_text._enforce_required_for_constraint_properties(
+        node_types, constraints
+    )
+
+    # name should now be required=true
+    assert node_types[0]["properties"][0]["required"] is True
+    # email should remain required=false
+    assert node_types[0]["properties"][1]["required"] is False
+
+
+def test_enforce_required_for_constraint_properties_already_true(
+    schema_from_text: SchemaFromTextExtractor,
+) -> None:
+    node_types: list[dict[str, Any]] = [
+        {
+            "label": "Person",
+            "properties": [
+                {"name": "name", "type": "STRING", "required": True},
+            ],
+        }
+    ]
+    constraints = [
+        {"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"}
+    ]
+
+    schema_from_text._enforce_required_for_constraint_properties(
+        node_types, constraints
+    )
+
+    assert node_types[0]["properties"][0]["required"] is True
+
+
+def test_enforce_required_for_constraint_properties_missing_required_field(
+    schema_from_text: SchemaFromTextExtractor,
+) -> None:
+    node_types: list[dict[str, Any]] = [
+        {
+            "label": "Person",
+            "properties": [
+                {"name": "name", "type": "STRING"},  # No required field
+            ],
+        }
+    ]
+    constraints = [
+        {"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"}
+    ]
+
+    schema_from_text._enforce_required_for_constraint_properties(
+        node_types, constraints
+    )
+
+    assert node_types[0]["properties"][0]["required"] is True
+
+
+def test_enforce_required_for_constraint_properties_no_constraints(
+    schema_from_text: SchemaFromTextExtractor,
+) -> None:
+    node_types: list[dict[str, Any]] = [
+        {
+            "label": "Person",
+            "properties": [
+                {"name": "name", "type": "STRING", "required": False},
+            ],
+        }
+    ]
+    constraints: list[dict[str, Any]] = []
+
+    schema_from_text._enforce_required_for_constraint_properties(
+        node_types, constraints
+    )
+
+    assert node_types[0]["properties"][0]["required"] is False
+
+
+def test_enforce_required_for_constraint_properties_skips_unconstrained_nodes(
+    schema_from_text: SchemaFromTextExtractor,
+) -> None:
+    node_types: list[dict[str, Any]] = [
+        {
+            "label": "Person",
+            "properties": [
+                {"name": "name", "type": "STRING", "required": False},
+            ],
+        },
+        {
+            "label": "Company",
+            "properties": [
+                {"name": "name", "type": "STRING", "required": False},
+            ],
+        },
+    ]
+    constraints = [
+        {"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"}
+    ]
+
+    schema_from_text._enforce_required_for_constraint_properties(
+        node_types, constraints
+    )
+
+    # Person.name should be required=true
+    assert node_types[0]["properties"][0]["required"] is True
+    # Company.name should remain required=false (no constraint on Company)
+    assert node_types[1]["properties"][0]["required"] is False
+
+
 @pytest.mark.asyncio
 async def test_schema_from_text_with_required_properties(
     schema_from_text: SchemaFromTextExtractor,
@@ -1638,6 +1759,45 @@ async def test_schema_from_text_handles_missing_required_field(
         assert prop.required is False
 
 
+@pytest.mark.asyncio
+async def test_schema_from_text_enforces_required_for_constrained_properties(
+    schema_from_text: SchemaFromTextExtractor,
+    mock_llm: AsyncMock,
+) -> None:
+    schema_json = """
+    {
+        "node_types": [
+            {
+                "label": "Person",
+                "properties": [
+                    {"name": "name", "type": "STRING", "required": false},
+                    {"name": "email", "type": "STRING", "required": false}
+                ]
+            }
+        ],
+        "relationship_types": [],
+        "patterns": [],
+        "constraints": [
+            {"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"}
+        ]
+    }
+    """
+    mock_llm.ainvoke.return_value = LLMResponse(content=schema_json)
+
+    schema = await schema_from_text.run(text="Sample text")
+
+    person = schema.node_type_from_label("Person")
+    assert person is not None
+
+    name_prop = next((p for p in person.properties if p.name == "name"), None)
+    email_prop = next((p for p in person.properties if p.name == "email"), None)
+
+    # name should be auto-fixed to required=true
+    assert name_prop is not None and name_prop.required is True
+    # email should remain required=false
+    assert email_prop is not None and email_prop.required is False
+
+
 @pytest.mark.asyncio
 @patch("neo4j_graphrag.experimental.components.schema.get_structured_schema")
 async def test_schema_from_existing_graph(mock_get_structured_schema: Mock) -> None: