diff --git a/semantica/kg/graph_builder.py b/semantica/kg/graph_builder.py index 314b923d..78039a60 100644 --- a/semantica/kg/graph_builder.py +++ b/semantica/kg/graph_builder.py @@ -262,6 +262,77 @@ class GraphBuilder: self._extractor_cache[key] = extractor_cls(method=method, **self.config) return self._extractor_cache[key] + def _remap_relationship_endpoints( + self, + entities: List[Dict[str, Any]], + relationships: List[Dict[str, Any]], + ) -> int: + """Rewrite relationship endpoints after entity resolution. + + Entity merging keeps the canonical entity ID and records the IDs of all + merged inputs in ``merged_from``. Relationships are collected before + resolution, so without this remapping they can continue to reference an + entity that is no longer present in the graph. + + Returns: + The number of relationship endpoints that were remapped. + """ + endpoint_map: Dict[Any, Any] = {} + + for entity in entities: + if not isinstance(entity, dict): + continue + + canonical_id = entity.get("id") + if canonical_id is None: + canonical_id = entity.get("entity_id") + if canonical_id is None: + continue + + # Keep canonical IDs stable and map every source ID retained by the + # merge operation to the surviving entity. + try: + endpoint_map[canonical_id] = canonical_id + except TypeError: + # Invalid/unhashable IDs are left for graph validation to report + # rather than making graph construction fail here. + continue + + merged_from = entity.get("merged_from") or [] + if isinstance(merged_from, (list, tuple, set)): + for source_id in merged_from: + if source_id is not None: + try: + endpoint_map[source_id] = canonical_id + except TypeError: + # Skip invalid aliases while preserving valid ones. + continue + + remapped_count = 0 + for relationship in relationships: + if not isinstance(relationship, dict): + continue + + for endpoint in ("source", "target"): + endpoint_id = relationship.get(endpoint) + try: + canonical_id = endpoint_map.get(endpoint_id) + except TypeError: + # Invalid/unhashable endpoints are left for graph validation + # to report rather than making graph construction fail here. + continue + + if canonical_id is not None and canonical_id != endpoint_id: + relationship[endpoint] = canonical_id + remapped_count += 1 + + if remapped_count: + self.logger.info( + "Remapped %d relationship endpoint(s) after entity resolution", + remapped_count, + ) + return remapped_count + def _extract_from_text(self, text: str, all_entities: List[Any], all_relationships: List[Any], **options): """Helper to extract knowledge from text using configured methods.""" if not options.get("extract", True): @@ -685,6 +756,16 @@ class GraphBuilder: f"Entity resolution complete: {len(all_entities)} -> {len(resolved_entities)} unique entities" ) + # Relationships were collected before entity resolution. Rewrite + # endpoints only when resolution produced merged entity IDs. + if resolver_to_use: + has_merged_entities = any( + isinstance(entity, dict) and entity.get("merged_from") + for entity in resolved_entities + ) + if has_merged_entities: + self._remap_relationship_endpoints(resolved_entities, all_relationships) + if input_relationships_count > 0 and len(all_relationships) == 0: warning_msg = ( f"All relationships were dropped during graph building: " diff --git a/tests/kg/test_graph_builder_external.py b/tests/kg/test_graph_builder_external.py index 72a918f6..a878c99b 100644 --- a/tests/kg/test_graph_builder_external.py +++ b/tests/kg/test_graph_builder_external.py @@ -123,6 +123,96 @@ class TestGraphBuilderExternal(unittest.TestCase): self.assertIn(("2", "3"), ids) self.assertIn(("3", "4"), ids) + def test_relationship_endpoints_are_remapped_after_entity_resolution(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + resolver = MagicMock() + resolver.resolve_entities.return_value = [ + { + "id": "alice:1", + "name": "Alice Chen", + "type": "Person", + "merged_from": ["alice:1", "alice:2"], + }, + {"id": "org:1", "name": "Zyx Qqqq", "type": "Organization"}, + ] + + graph = builder.build( + { + "entities": [ + {"id": "alice:1", "name": "Alice Chen", "type": "Person"}, + {"id": "alice:2", "name": "Alice Chen", "type": "Person"}, + {"id": "org:1", "name": "Zyx Qqqq", "type": "Organization"}, + ], + "relationships": [ + { + "source": "alice:2", + "target": "org:1", + "type": "WORKS_FOR", + } + ], + }, + entity_resolver=resolver, + ) + + self.assertEqual( + graph["relationships"], + [{"source": "alice:1", "target": "org:1", "type": "WORKS_FOR"}], + ) + entity_ids = {entity["id"] for entity in graph["entities"]} + for relationship in graph["relationships"]: + self.assertIn(relationship["source"], entity_ids) + self.assertIn(relationship["target"], entity_ids) + + def test_unhashable_entity_ids_do_not_crash_remapping(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + resolver = MagicMock() + resolver.resolve_entities.return_value = [ + { + "id": ["invalid-canonical-id"], + "name": "Invalid ID", + "type": "Person", + "merged_from": ["invalid-canonical-id"], + }, + { + "id": "alice:1", + "name": "Alice Chen", + "type": "Person", + "merged_from": [["invalid-source-id"]], + }, + ] + + graph = builder.build( + { + "entities": [{"id": "alice:1", "name": "Alice Chen", "type": "Person"}], + "relationships": [], + }, + entity_resolver=resolver, + ) + + self.assertEqual(len(graph["entities"]), 2) + + def test_relationship_remapping_skips_unmerged_entities(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + resolver = MagicMock() + resolver.resolve_entities.return_value = [ + {"id": "alice:1", "name": "Alice Chen", "type": "Person"}, + {"id": "org:1", "name": "Zyx Qqqq", "type": "Organization"}, + ] + + with patch.object(builder, "_remap_relationship_endpoints") as remap: + builder.build( + { + "entities": [ + {"id": "alice:1", "name": "Alice Chen", "type": "Person"}, + {"id": "org:1", "name": "Zyx Qqqq", "type": "Organization"}, + ], + "relationships": [], + }, + entity_resolver=resolver, + ) + + remap.assert_not_called() + def test_warning_when_all_relationships_dropped(self): builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)