Fix cleanup ops for inferred schema (#403)

* Add filtering on inferred node types/rel types with empty string

* Update Unit tests for schema

* Update changelog

* Filter out invalid labels/rel types

Author: NathalieCharbel

CHANGELOG.md

@@ -6,6 +6,7 @@
 
 - Fixed documentation for PdfLoader
 - Fixed a bug where the `format` argument for `OllamaLLM` was not propagated to the client.
+- Fixed `AttributeError` in `SchemaFromTextExtractor` when filtering out node/relationship types with no labels.
 - Fixed an import error in `VertexAIEmbeddings`.
 
 ## 1.9.0

src/neo4j_graphrag/experimental/components/schema.py

@@ -498,49 +498,38 @@ def _filter_invalid_patterns(
 
         return filtered_patterns
 
-    def _filter_nodes_without_labels(
-        self, node_types: List[Dict[str, Any]]
+    def _filter_items_without_labels(
+        self, items: List[Dict[str, Any]], item_type: str
     ) -> List[Dict[str, Any]]:
-        """
-        Filter out node types that have no labels.
-
-        Args:
-            node_types: List of node type definitions.
-
-        Returns:
-            Filtered list of node types containing only those with valid labels.
-        """
-        filtered_nodes = []
-        for node_type in node_types:
-            if node_type.get("label"):
-                filtered_nodes.append(node_type)
+        """Filter out items that have no valid labels."""
+        filtered_items = []
+        for item in items:
+            if isinstance(item, str):
+                if item and " " not in item and not item.startswith("{"):
+                    filtered_items.append({"label": item})
+                elif item:
+                    logging.info(
+                        f"Filtering out {item_type} with invalid label: {item}"
+                    )
+            elif isinstance(item, dict) and item.get("label"):
+                filtered_items.append(item)
             else:
-                logging.info(f"Filtering out node type with missing label: {node_type}")
+                logging.info(f"Filtering out {item_type} with missing label: {item}")
+        return filtered_items
 
-        return filtered_nodes
+    def _filter_nodes_without_labels(
+        self, node_types: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Filter out node types that have no labels."""
+        return self._filter_items_without_labels(node_types, "node type")
 
     def _filter_relationships_without_labels(
         self, relationship_types: List[Dict[str, Any]]
     ) -> List[Dict[str, Any]]:
-        """
-        Filter out relationship types that have no labels.
-
-        Args:
-            relationship_types: List of relationship type definitions.
-
-        Returns:
-            Filtered list of relationship types containing only those with valid labels.
-        """
-        filtered_relationships = []
-        for rel_type in relationship_types:
-            if rel_type.get("label"):
-                filtered_relationships.append(rel_type)
-            else:
-                logging.info(
-                    f"Filtering out relationship type with missing label: {rel_type}"
-                )
-
-        return filtered_relationships
+        """Filter out relationship types that have no labels."""
+        return self._filter_items_without_labels(
+            relationship_types, "relationship type"
+        )
 
     @validate_call
     async def run(self, text: str, examples: str = "", **kwargs: Any) -> GraphSchema:

tests/unit/experimental/components/test_schema.py

@@ -794,12 +794,11 @@ def schema_json_with_nodes_without_labels() -> str:
                     {"name": "name", "type": "STRING"}
                 ]
             },
-            {
-                "label": "Organization",
-                "properties": [
-                    {"name": "name", "type": "STRING"}
-                ]
-            }
+            "Organization",
+            "",
+            "Company",
+            "Invalid description with spaces",
+            "{\\"invalid\\": \\"json object\\"}"
         ],
         "relationship_types": [
             {
@@ -852,12 +851,11 @@ def schema_json_with_relationships_without_labels() -> str:
                     {"name": "since", "type": "DATE"}
                 ]
             },
-            {
-                "label": "MANAGES",
-                "properties": [
-                    {"name": "since", "type": "DATE"}
-                ]
-            }
+            "MANAGES",
+            "",
+            "SUPERVISES",
+            "invalid relationship description",
+            "{\\"invalid\\": \\"json\\"}"
         ],
         "patterns": [
             ["Person", "WORKS_FOR", "Organization"],
@@ -921,10 +919,11 @@ async def test_schema_from_text_filters_nodes_without_labels(
     # run the schema extraction
     schema = await schema_from_text.run(text="Sample text for extraction")
 
-    # verify that nodes without labels were filtered out (2 out of 4 nodes should be removed)
-    assert len(schema.node_types) == 2
+    # verify that nodes without labels were filtered out (5 out of 8 nodes should be removed)
+    assert len(schema.node_types) == 3
     assert schema.node_type_from_label("Person") is not None
     assert schema.node_type_from_label("Organization") is not None
+    assert schema.node_type_from_label("Company") is not None
 
     # verify that the pattern is still valid with the remaining nodes
     assert schema.patterns is not None
@@ -946,11 +945,12 @@ async def test_schema_from_text_filters_relationships_without_labels(
     # run the schema extraction
     schema = await schema_from_text.run(text="Sample text for extraction")
 
-    # verify that relationships without labels were filtered out (2 out of 4 relationships should be removed)
+    # verify that relationships without labels were filtered out (5 out of 8 relationships should be removed)
     assert schema.relationship_types is not None
-    assert len(schema.relationship_types) == 2
+    assert len(schema.relationship_types) == 3
     assert schema.relationship_type_from_label("WORKS_FOR") is not None
     assert schema.relationship_type_from_label("MANAGES") is not None
+    assert schema.relationship_type_from_label("SUPERVISES") is not None
 
     # verify that the patterns are still valid with the remaining relationships
     assert schema.patterns is not None