Test happy path

Author: stellasia

src/neo4j_graphrag/experimental/components/schema.py

@@ -49,6 +49,9 @@
 from neo4j_graphrag.schema import get_structured_schema
 
 
+logger = logging.getLogger(__name__)
+
+
 class PropertyType(BaseModel):
     """
     Represents a property on a node or relationship in the graph.
@@ -622,19 +625,19 @@ async def run(self, text: str, examples: str = "", **kwargs: Any) -> GraphSchema
 class SchemaFromExistingGraphExtractor(BaseSchemaBuilder):
     """A class to build a GraphSchema object from an existing graph.
 
-    Uses the get_structured_schema function to extract existing node labels,
-    relationship types, properties and existence constraints.
+     Uses the get_structured_schema function to extract existing node labels,
+     relationship types, properties and existence constraints.
 
-    By default, the built schema does not allow any additional item (property,
-    node label, relationship type or pattern).
+     By default, the built schema does not allow any additional item (property,
+     node label, relationship type or pattern).
 
-       Args:
-            driver (neo4j.Driver): connection to the neo4j database.
-            additional_properties (bool, default False): see GraphSchema
-            additional_node_types (bool, default False): see GraphSchema
-            additional_relationship_types (bool, default False): see GraphSchema:
-            additional_patterns (bool, default False): see GraphSchema:
-            neo4j_database (Optional | str): name of the neo4j database to use
+    Args:
+         driver (neo4j.Driver): connection to the neo4j database.
+         additional_properties (bool, default False): see GraphSchema
+         additional_node_types (bool, default False): see GraphSchema
+         additional_relationship_types (bool, default False): see GraphSchema:
+         additional_patterns (bool, default False): see GraphSchema:
+         neo4j_database (Optional | str): name of the neo4j database to use
     """
 
     def __init__(
@@ -672,7 +675,7 @@ def _extract_required_properties(
         """
         schema_metadata = structured_schema.get("metadata", {})
         existence_constraint = []  # list of (node label, property name)
-        for constraint in schema_metadata.get("constraints", []):
+        for constraint in schema_metadata.get("constraint", []):
             if constraint["type"] in (
                 "NODE_PROPERTY_EXISTENCE",
                 "NODE_KEY",
@@ -688,10 +691,11 @@ def _extract_required_properties(
                 existence_constraint.append((lab, prop))
         return existence_constraint
 
-    async def run(self) -> GraphSchema:
+    async def run(self, *args, **kwargs) -> GraphSchema:
         structured_schema = get_structured_schema(self.driver, database=self.database)
         existence_constraint = self._extract_required_properties(structured_schema)
 
+        # node label with properties
         node_labels = set(structured_schema["node_props"].keys())
         node_types = [
             {
@@ -708,6 +712,8 @@ async def run(self) -> GraphSchema:
             }
             for key, properties in structured_schema["node_props"].items()
         ]
+
+        # relationships with properties
         rel_labels = set(structured_schema["rel_props"].keys())
         relationship_types = [
             {
@@ -723,27 +729,41 @@ async def run(self) -> GraphSchema:
             }
             for key, properties in structured_schema["rel_props"].items()
         ]
+
         patterns = [
             (s["start"], s["type"], s["end"])
             for s in structured_schema["relationships"]
         ]
+
         # deal with nodes and relationships without properties
         for source, rel, target in patterns:
             if source not in node_labels:
+                if not self.additional_properties:
+                    logger.warning(
+                        f"SCHEMA: found node label {source} without property and additional_properties=False: this node label will always be pruned!"
+                    )
                 node_labels.add(source)
                 node_types.append(
                     {
                         "label": source,
                     }
                 )
             if target not in node_labels:
+                if not self.additional_properties:
+                    logger.warning(
+                        f"SCHEMA: found node label {target} without property and additional_properties=False: this node label will always be pruned!"
+                    )
                 node_labels.add(target)
                 node_types.append(
                     {
                         "label": target,
                     }
                 )
             if rel not in rel_labels:
+                if not self.additional_properties:
+                    logger.warning(
+                        f"SCHEMA: found relationship type {rel} without property and additional_properties=False: this relationship type will always be pruned!"
+                    )
                 rel_labels.add(rel)
                 relationship_types.append(
                     {

tests/unit/experimental/components/test_schema.py

@@ -16,7 +16,7 @@
 
 import json
 from typing import Tuple, Any
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, patch, Mock
 
 import pytest
 from pydantic import ValidationError
@@ -29,6 +29,7 @@
     RelationshipType,
     SchemaFromTextExtractor,
     GraphSchema,
+    SchemaFromExistingGraphExtractor,
 )
 import os
 import tempfile
@@ -957,3 +958,66 @@ async def test_schema_from_text_filters_relationships_without_labels(
     assert len(schema.patterns) == 2
     assert ("Person", "WORKS_FOR", "Organization") in schema.patterns
     assert ("Person", "MANAGES", "Organization") in schema.patterns
+
+
+@pytest.mark.asyncio
+@patch("neo4j_graphrag.experimental.components.schema.get_structured_schema")
+async def test_schema_from_existing_graph(mock_get_structured_schema: Mock):
+    mock_get_structured_schema.return_value = {
+        "node_props": {
+            "Person": [
+                {"property": "id", "type": "INTEGER"},
+                {"property": "name", "type": "STRING"},
+            ]
+        },
+        "rel_props": {"KNOWS": [{"property": "fromDate", "type": "DATE"}]},
+        "relationships": [
+            {"start": "Person", "type": "KNOWS", "end": "Person"},
+            {"start": "Person", "type": "LIVES_IN", "end": "City"},
+        ],
+        "metadata": {
+            "constraint": [
+                {
+                    "id": 7,
+                    "name": "person_id",
+                    "type": "NODE_PROPERTY_EXISTENCE",
+                    "entityType": "NODE",
+                    "labelsOrTypes": ["Person"],
+                    "properties": ["id"],
+                    "ownedIndex": "person_id",
+                    "propertyType": None,
+                },
+            ],
+            "index": [
+                {
+                    "label": "Person",
+                    "properties": ["name"],
+                    "size": 2,
+                    "type": "RANGE",
+                    "valuesSelectivity": 1.0,
+                    "distinctValues": 2.0,
+                },
+            ],
+        },
+    }
+    driver = Mock()
+    schema_builder = SchemaFromExistingGraphExtractor(
+        driver=driver,
+    )
+    schema = await schema_builder.run()
+    assert isinstance(schema, GraphSchema)
+    assert len(schema.node_types) == 2
+    person_node_type = schema.node_type_from_label("Person")
+    assert person_node_type is not None
+    id_person_property = [p for p in person_node_type.properties if p.name == "id"][0]
+    assert id_person_property.required is True
+
+    assert schema.node_type_from_label("City") is not None
+    assert len(schema.relationship_types) == 2
+    assert schema.relationship_type_from_label("KNOWS") is not None
+    assert schema.relationship_type_from_label("LIVES_IN") is not None
+
+    assert schema.patterns == (
+        ("Person", "KNOWS", "Person"),
+        ("Person", "LIVES_IN", "City"),
+    )