Invert rel direction

Author: NathalieCharbel

docs/source/user_guide_kg_builder.rst

@@ -847,6 +847,7 @@ This behaviour can be changed by using the `enforce_schema` flag in the `LLMEnti
 
 In this scenario, any extracted node/relation/property that is not part of the provided schema will be pruned.
 Any relation whose start node or end node does not conform to the provided tuple in `potential_schema` will be pruned.
+If a relation start/end nodes are valid but the direction is incorrect, the latter will be inverted.
 If a node is left with no properties, it will be also pruned.
 
 .. warning::

src/neo4j_graphrag/experimental/components/entity_relation_extractor.py

@@ -435,6 +435,7 @@ def _enforce_relationships(
         Keep only those whose types are in schema, start/end node conform to schema,
         and start/end nodes are in filtered nodes (i.e., kept after node enforcement).
         For each valid relationship, filter out properties not present in the schema.
+        If a relationship direct is incorrect, invert it.
         """
         if self.enforce_schema != SchemaEnforcementMode.STRICT:
             return extracted_relationships
@@ -457,17 +458,22 @@ def _enforce_relationships(
             start_label = valid_nodes[rel.start_node_id]
             end_label = valid_nodes[rel.end_node_id]
 
-            if (potential_schema and
-                    (start_label, rel.type, end_label) not in potential_schema):
-                continue
+            tuple_valid = True
+            if potential_schema:
+                tuple_valid = (start_label, rel.type, end_label) in potential_schema
+                reverse_tuple_valid = ((end_label, rel.type, start_label) in
+                                       potential_schema)
+
+                if not tuple_valid and not reverse_tuple_valid:
+                    continue
 
             allowed_props = schema_relation.get("properties", [])
             filtered_props = self._enforce_properties(rel.properties, allowed_props)
 
             valid_rels.append(
                 Neo4jRelationship(
-                    start_node_id=rel.start_node_id,
-                    end_node_id=rel.end_node_id,
+                    start_node_id=rel.start_node_id if tuple_valid else rel.end_node_id,
+                    end_node_id=rel.end_node_id if tuple_valid else rel.start_node_id,
                     type=rel.type,
                     properties=filtered_props,
                     embedding_properties=rel.embedding_properties,

tests/unit/experimental/components/test_entity_relation_extractor.py

@@ -481,6 +481,38 @@ async def test_extractor_schema_enforcement_removed_relation_start_end_nodes():
     assert len(result.relationships) == 0
 
 
+@pytest.mark.asyncio
+async def test_extractor_schema_enforcement_inverted_relation_direction():
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(
+        content='{"nodes":[{"id":"1","label":"Person","properties":{"name":"Alice"}},'
+                '{"id":"2","label":"City","properties":{"name":"London"}}],'
+                '"relationships":[{"start_node_id":"2","end_node_id":"1",'
+                '"type":"LIVES_IN","properties":{}}]}'
+    )
+
+    extractor = LLMEntityRelationExtractor(llm=llm,
+                                           create_lexical_graph=False,
+                                           enforce_schema=SchemaEnforcementMode.STRICT)
+
+    schema = SchemaConfig(
+        entities={"Person": {"label": "Person",
+                             "properties": [{"name": "name", "type": "STRING"}]},
+                  "City": {"label": "City",
+                           "properties": [{"name": "name", "type": "STRING"}]}},
+        relations={"LIVES_IN": {"label": "LIVES_IN"}},
+        potential_schema=[("Person", "LIVES_IN", "City")])
+
+    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].start_node_id.split(":")[1] == "1"
+    assert result.relationships[0].end_node_id.split(":")[1] == "2"
+
+
 def test_fix_invalid_json_empty_result() -> None:
     json_string = "invalid json"