Deal with missing relationships, add tests

Author: stellasia

src/neo4j_graphrag/experimental/components/entity_relation_extractor.py

@@ -403,7 +403,7 @@ def _enforce_nodes(
             schema_entity = schema.entities.get(node.label)
             if not schema_entity:
                 continue
-            allowed_props = schema_entity.get("properties", [])
+            allowed_props = schema_entity.get("properties")
             if allowed_props:
                 filtered_props = self._enforce_properties(
                     node.properties, allowed_props
@@ -439,16 +439,17 @@ def _enforce_relationships(
         if self.enforce_schema != SchemaEnforcementMode.STRICT:
             return extracted_relationships
 
+        if schema.relations is None:
+            return extracted_relationships
+
         valid_rels = []
 
         valid_nodes = {node.id: node.label for node in filtered_nodes}
 
         potential_schema = schema.potential_schema
 
         for rel in extracted_relationships:
-            schema_relation = (
-                schema.relations.get(rel.type) if schema.relations else None
-            )
+            schema_relation = schema.relations.get(rel.type)
             if not schema_relation:
                 continue
 
@@ -473,7 +474,7 @@ def _enforce_relationships(
                 if not tuple_valid and not reverse_tuple_valid:
                     continue
 
-            allowed_props = schema_relation.get("properties", [])
+            allowed_props = schema_relation.get("properties")
             if allowed_props:
                 filtered_props = self._enforce_properties(rel.properties, allowed_props)
             else:

src/neo4j_graphrag/experimental/components/schema.py

@@ -99,7 +99,7 @@ class SchemaConfig(DataModel):
     @model_validator(mode="before")
     def check_schema(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         entities = data.get("entities", {}).keys()
-        relations = data.get("relations", {}).keys()
+        relations = (data.get("relations") or {}).keys()
         potential_schema = data.get("potential_schema", [])
 
         if potential_schema:

tests/unit/experimental/components/test_entity_relation_extractor.py

@@ -564,6 +564,74 @@ async def test_extractor_schema_enforcement_inverted_relation_direction() -> Non
     assert result.relationships[0].end_node_id.split(":")[1] == "2"
 
 
+@pytest.mark.asyncio
+async def test_extractor_schema_enforcement_none_relationships_in_schema() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(
+        content='{"nodes":[{"id":"1","label":"Person","properties":'
+        '{"name":"Alice"}},{"id":"2","label":"Person","properties":'
+        '{"name":"Bob"}}],'
+        '"relationships":[{"start_node_id":"1","end_node_id":"2",'
+        '"type":"FRIENDS_WITH","properties":{}}]}'
+    )
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm, create_lexical_graph=False, enforce_schema=SchemaEnforcementMode.STRICT
+    )
+
+    schema = SchemaConfig(
+        entities={
+            "Person": {
+                "label": "Person",
+                "properties": [{"name": "name", "type": "STRING"}],
+            }
+        },
+        relations=None,
+        potential_schema=None,
+    )
+
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+
+    result: Neo4jGraph = await extractor.run(chunks, schema=schema)
+
+    assert len(result.nodes) == 2
+    assert len(result.relationships) == 1
+    assert result.relationships[0].type == "FRIENDS_WITH"
+
+
+@pytest.mark.asyncio
+async def test_extractor_schema_enforcement_empty_relationships_in_schema() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(
+        content='{"nodes":[{"id":"1","label":"Person","properties":'
+        '{"name":"Alice"}},{"id":"2","label":"Person","properties":'
+        '{"name":"Bob"}}],'
+        '"relationships":[{"start_node_id":"1","end_node_id":"2",'
+        '"type":"FRIENDS_WITH","properties":{}}]}'
+    )
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm, create_lexical_graph=False, enforce_schema=SchemaEnforcementMode.STRICT
+    )
+
+    schema = SchemaConfig(
+        entities={
+            "Person": {
+                "label": "Person",
+                "properties": [{"name": "name", "type": "STRING"}],
+            }
+        },
+        relations={},
+        potential_schema=None,
+    )
+
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+
+    result: Neo4jGraph = await extractor.run(chunks, schema=schema)
+
+    assert len(result.relationships) == 0
+
+
 def test_fix_invalid_json_empty_result() -> None:
     json_string = "invalid json"