Add required property support to SchemaFromTextExtractor

Author: AmirLayegh

src/neo4j_graphrag/experimental/components/schema.py

@@ -666,6 +666,66 @@ def _filter_invalid_constraints(
             filtered_constraints.append(constraint)
         return filtered_constraints
 
+    def _filter_properties_required_field(
+        self, node_types: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Sanitize the 'required' field in node type properties. Ensures 'required' is a valid boolean.
+           converts known string values (true, yes, 1, false, no, 0) to booleans and removes unrecognized values.
+        """
+        for node_type in node_types:
+            properties = node_type.get("properties", [])
+            if not properties:
+                continue
+            for prop in properties:
+                if not isinstance(prop, dict):
+                    continue
+
+                required_value = prop.get("required")
+
+                #  Not provided - will use Pydantic default (false)
+                if required_value is None:
+                    continue
+
+                # already a valid boolean
+                if isinstance(required_value, bool):
+                    continue
+
+                prop_name = prop.get("name", "unknown")
+                node_label = node_type.get("label", "unknown")
+
+                if isinstance(required_value, str):
+                    if required_value.lower() in ("true", "yes", "1"):
+                        prop["required"] = True
+                        logging.info(
+                            f"Converted 'required' value {required_value} to True "
+                            f"for property '{prop_name}' on node '{node_label}'"
+                        )
+                    elif required_value.lower() in ("false", "no", "0"):
+                        prop["required"] = False
+                        logging.info(
+                            f"Converted 'required' value '{required_value}' to False "
+                            f"for property '{prop_name}' on node '{node_label}' "
+                        )
+                    # Unknown string values
+                    else:
+                        logging.info(
+                            f"Removing unrecognized 'required' value '{required_value}' "
+                            f"for property '{prop_name}' on node '{node_label}'. "
+                            f"Using default (False) " # TODO: Not sure if we have to convert it to the default value - double check!
+                        )
+                        prop.pop("required", None)
+                else:
+                    # Non-string, non-boolean - remove
+                    logging.info(
+                        f"Removing invalid 'required' value '{required_value}' (type: {type(required_value).__name__}) "
+                        f"for property '{prop_name}' on node '{node_label}'. "
+                        f"Using default (False). "
+                    )
+                    prop.pop("required", None)
+
+        return node_types
+
+
     def _clean_json_content(self, content: str) -> str:
         content = content.strip()
 
@@ -746,6 +806,10 @@ async def run(self, text: str, examples: str = "", **kwargs: Any) -> GraphSchema
                 extracted_relationship_types
             )
 
+        extracted_node_types = self._filter_properties_required_field(
+            extracted_node_types
+        )
+
         # Filter out invalid patterns before validation
         if extracted_patterns:
             extracted_patterns = self._filter_invalid_patterns(

src/neo4j_graphrag/generation/prompts.py

@@ -223,7 +223,12 @@ class SchemaExtractionTemplate(PromptTemplate):
 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.
-
+9. REQUIRED PROPERTIES:
+9.1 Mark a property as "required": true if every instance of that node/relationship type MUST have this property (non-nullable).
+9.2 Mark a property as "required": false if the property is optional and may be absent on some instances.
+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.
 
 Accepted property types are: BOOLEAN, DATE, DURATION, FLOAT, INTEGER, LIST,
 LOCAL_DATETIME, LOCAL_TIME, POINT, STRING, ZONED_DATETIME, ZONED_TIME.
@@ -236,7 +241,13 @@ class SchemaExtractionTemplate(PromptTemplate):
       "properties": [
         {{
           "name": "name",
-          "type": "STRING"
+          "type": "STRING",
+          "required": true
+        }},
+        {{
+          "name": "email",
+          "type": "STRING",
+          "required": false
         }}
       ]
     }}