From 9f77dfb761e9256ad9938e84fe50a67107fe442d Mon Sep 17 00:00:00 2001 From: KaifAhmad1 Date: Sun, 18 Jan 2026 22:10:02 +0530 Subject: [PATCH] Fix GraphBuilder external relationships; refs #208 #206 --- semantica/kg/graph_builder.py | 56 +++++- tests/kg/test_graph_builder_external.py | 232 ++++++++++++++++++++++++ tests/kg/test_kg.py | 68 +++++++ 3 files changed, 351 insertions(+), 5 deletions(-) create mode 100644 tests/kg/test_graph_builder_external.py diff --git a/semantica/kg/graph_builder.py b/semantica/kg/graph_builder.py index cb41ec01..6708b141 100644 --- a/semantica/kg/graph_builder.py +++ b/semantica/kg/graph_builder.py @@ -161,7 +161,14 @@ class GraphBuilder: } all_relationships.append(rel_dict) elif isinstance(item, dict): - # Detect and normalize Entity objects inside dict + if "source_id" in item and "source" not in item: + item["source"] = item["source_id"] + if "target_id" in item and "target" not in item: + item["target"] = item["target_id"] + if "subject" in item and "source" not in item: + item["source"] = item["subject"] + if "object" in item and "target" not in item: + item["target"] = item["object"] if "source" in item and not isinstance(item["source"], str): src = item["source"] item["source"] = getattr(src, "id", getattr(src, "text", str(src))) @@ -347,6 +354,21 @@ class GraphBuilder: elif not isinstance(sources, list): sources = [sources] + # Count input relationships for warning if all are dropped + input_relationships_count = 0 + if isinstance(source_dict, dict): + rels = source_dict.get("relationships", []) + if isinstance(rels, list): + input_relationships_count += len(rels) + elif rels is not None: + input_relationships_count += 1 + if explicit_relationships: + for rel_item in explicit_relationships: + if isinstance(rel_item, list): + input_relationships_count += len(rel_item) + else: + input_relationships_count += 1 + # Track graph building build_start_time = time.time() @@ -468,11 +490,12 @@ class GraphBuilder: pipeline_id=pipeline_id, ) - # Check if relationships are already in dictionary format sample_rel = relationships_list[0] if relationships_list else None is_dict_format = isinstance(sample_rel, dict) and ( - "source" in sample_rel and "target" in sample_rel - ) and not hasattr(sample_rel, "__dict__") # Ensure it's not a class instance + ("source" in sample_rel and "target" in sample_rel) + or ("source_id" in sample_rel and "target_id" in sample_rel) + or ("subject" in sample_rel and "object" in sample_rel) + ) and not hasattr(sample_rel, "__dict__") if is_dict_format: # Fast path: directly append dictionaries after normalizing source/target @@ -481,8 +504,15 @@ class GraphBuilder: batch = relationships_list[i:i + batch_size] for item in batch: if isinstance(item, dict): - # Normalize source/target if they are objects rel_dict = item.copy() + if "source_id" in rel_dict and "source" not in rel_dict: + rel_dict["source"] = rel_dict["source_id"] + if "target_id" in rel_dict and "target" not in rel_dict: + rel_dict["target"] = rel_dict["target_id"] + if "subject" in rel_dict and "source" not in rel_dict: + rel_dict["source"] = rel_dict["subject"] + if "object" in rel_dict and "target" not in rel_dict: + rel_dict["target"] = rel_dict["object"] if "source" in rel_dict and not isinstance(rel_dict["source"], str): src = rel_dict["source"] rel_dict["source"] = getattr(src, "id", getattr(src, "text", str(src))) @@ -571,6 +601,14 @@ class GraphBuilder: f"Entity resolution complete: {len(all_entities)} -> {len(resolved_entities)} unique entities" ) + if input_relationships_count > 0 and len(all_relationships) == 0: + warning_msg = ( + f"All relationships were dropped during graph building: " + f"{input_relationships_count} input relationships, 0 in final graph" + ) + self.logger.warning(warning_msg) + print(f"Warning: {warning_msg}") + # Build graph structure print("Building graph structure...") structure_start = time.time() @@ -682,6 +720,14 @@ class GraphBuilder: ) raise + def build_single_source( + self, + kg_data: Dict[str, Any], + pipeline_id: Optional[str] = None, + **options, + ) -> Dict[str, Any]: + return self.build(kg_data, pipeline_id=pipeline_id, **options) + def add_temporal_edge( self, graph, diff --git a/tests/kg/test_graph_builder_external.py b/tests/kg/test_graph_builder_external.py new file mode 100644 index 00000000..72a918f6 --- /dev/null +++ b/tests/kg/test_graph_builder_external.py @@ -0,0 +1,232 @@ +import unittest +from unittest.mock import MagicMock, patch + +from semantica.kg.graph_builder import GraphBuilder + + +class TestGraphBuilderExternal(unittest.TestCase): + def setUp(self): + self.mock_tracker_patcher = patch("semantica.utils.progress_tracker.get_progress_tracker") + self.mock_get_tracker = self.mock_tracker_patcher.start() + self.mock_tracker = MagicMock() + self.mock_get_tracker.return_value = self.mock_tracker + + self.mock_resolver_patcher = patch("semantica.kg.entity_resolver.EntityResolver") + self.mock_resolver_cls = self.mock_resolver_patcher.start() + + self.mock_conflict_patcher = patch("semantica.conflicts.conflict_detector.ConflictDetector") + self.mock_conflict_cls = self.mock_conflict_patcher.start() + + def tearDown(self): + self.mock_tracker_patcher.stop() + self.mock_resolver_patcher.stop() + self.mock_conflict_patcher.stop() + + def test_single_source_dict_with_source_id_target_id(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + + entities = [ + {"id": "drug:1", "name": "Aspirin", "type": "Drug"}, + {"id": "disease:1", "name": "Myocardial infarction", "type": "Disease"}, + ] + relationships = [ + {"source_id": "drug:1", "target_id": "disease:1", "type": "TREATS"}, + ] + + source = {"entities": entities, "relationships": relationships} + + kg = builder.build(source) + + self.assertEqual(len(kg["entities"]), 2) + self.assertEqual(len(kg["relationships"]), 1) + rel = kg["relationships"][0] + self.assertEqual(rel.get("source"), "drug:1") + self.assertEqual(rel.get("target"), "disease:1") + self.assertEqual(kg["metadata"]["num_relationships"], 1) + + def test_sources_list_merge_with_external_relationships(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + + source1 = { + "entities": [{"id": "1", "name": "A"}], + "relationships": [{"source_id": "1", "target_id": "2", "type": "REL_1"}], + } + source2 = { + "entities": [{"id": "2", "name": "B"}], + "relationships": [{"source_id": "2", "target_id": "1", "type": "REL_2"}], + } + + kg = builder.build([source1, source2]) + + self.assertEqual(len(kg["entities"]), 2) + self.assertEqual(len(kg["relationships"]), 2) + sources = {r["source"] for r in kg["relationships"]} + targets = {r["target"] for r in kg["relationships"]} + self.assertEqual(sources, {"1", "2"}) + self.assertEqual(targets, {"1", "2"}) + + def test_build_with_explicit_relationships_argument_external_ids(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + + entities = [ + {"id": "1", "name": "A"}, + {"id": "2", "name": "B"}, + ] + relationships = [ + {"source_id": "1", "target_id": "2", "type": "REL"}, + ] + + kg = builder.build(entities, relationships=relationships) + + self.assertEqual(len(kg["entities"]), 2) + self.assertEqual(len(kg["relationships"]), 1) + rel = kg["relationships"][0] + self.assertEqual(rel.get("source"), "1") + self.assertEqual(rel.get("target"), "2") + + def test_build_single_source_external_graph(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + + source = { + "entities": [{"id": "1", "name": "A"}], + "relationships": [{"source_id": "1", "target_id": "1", "type": "SELF"}], + } + + kg = builder.build_single_source(source) + + self.assertEqual(len(kg["entities"]), 1) + self.assertEqual(len(kg["relationships"]), 1) + rel = kg["relationships"][0] + self.assertEqual(rel.get("source"), "1") + self.assertEqual(rel.get("target"), "1") + + def test_relationship_key_variants_normalized(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + + entities = [ + {"id": "1", "name": "A"}, + {"id": "2", "name": "B"}, + {"id": "3", "name": "C"}, + {"id": "4", "name": "D"}, + ] + relationships = [ + {"source_id": "1", "target_id": "2", "type": "R1"}, + {"source": "2", "target": "3", "type": "R2"}, + {"subject": "3", "object": "4", "type": "R3"}, + ] + + kg = builder.build({"entities": entities, "relationships": relationships}) + + self.assertEqual(len(kg["relationships"]), 3) + ids = {(r["source"], r["target"]) for r in kg["relationships"]} + self.assertIn(("1", "2"), ids) + self.assertIn(("2", "3"), ids) + self.assertIn(("3", "4"), ids) + + def test_warning_when_all_relationships_dropped(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + + source = { + "entities": [], + "relationships": [{"foo": "x"}, {"bar": "y"}], + } + + with patch.object(builder.logger, "warning") as mock_warning: + kg = builder.build(source) + + self.assertEqual(len(kg["relationships"]), 0) + mock_warning.assert_called() + args, _ = mock_warning.call_args + self.assertIn("All relationships were dropped", args[0]) + + def test_no_warning_when_some_relationships_kept(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + + source = { + "entities": [{"id": "1"}, {"id": "2"}], + "relationships": [ + {"source_id": "1", "target_id": "2", "type": "REL"}, + {"foo": "x"}, + ], + } + + with patch.object(builder.logger, "warning") as mock_warning: + kg = builder.build(source) + + self.assertEqual(len(kg["relationships"]), 2) + mock_warning.assert_not_called() + + def test_issue_208_minimal_reproduction_shape(self): + builder = GraphBuilder( + merge_entities=False, + entity_resolution_strategy="none", + resolve_conflicts=False, + ) + + entities = [ + {"id": "e1", "name": "Entity 1"}, + {"id": "e2", "name": "Entity 2"}, + {"id": "e3", "name": "Entity 3"}, + ] + relationships = [ + {"source_id": "e1", "target_id": "e2", "type": "REL_1"}, + {"source_id": "e2", "target_id": "e3", "type": "REL_2"}, + ] + + entity_ids = {e["id"] for e in entities} + for r in relationships: + self.assertIn(r["source_id"], entity_ids) + self.assertIn(r["target_id"], entity_ids) + + kg = builder.build( + sources=[{"entities": entities, "relationships": relationships}], + merge_entities=False, + ) + + self.assertEqual(len(kg["entities"]), 3) + self.assertEqual(len(kg["relationships"]), 2) + pairs = {(r["source"], r["target"]) for r in kg["relationships"]} + self.assertIn(("e1", "e2"), pairs) + self.assertIn(("e2", "e3"), pairs) + + def test_issue_206_earnings_call_shape(self): + builder = GraphBuilder( + merge_entities=False, + entity_resolution_strategy="none", + resolve_conflicts=False, + ) + + entities = [ + { + "id": "entity_446_MDA Space Ltd.", + "name": "MDA Space Ltd.", + "type": "ORGANIZATION", + }, + { + "id": "entity_500_$409.8 million", + "name": "$409.8 million", + "type": "MONEY", + }, + ] + + relationships = [ + { + "id": None, + "source_id": "MDA Space Ltd.", + "target_id": "$409.8 million", + "type": "HAS_REVENUE", + "confidence": 0.975, + "metadata": {}, + } + ] + + kg = builder.build( + sources=[{"entities": entities, "relationships": relationships}], + merge_entities=False, + ) + + self.assertEqual(len(kg["entities"]), 2) + self.assertEqual(len(kg["relationships"]), 1) + rel = kg["relationships"][0] + self.assertEqual(rel.get("source"), "MDA Space Ltd.") + self.assertEqual(rel.get("target"), "$409.8 million") diff --git a/tests/kg/test_kg.py b/tests/kg/test_kg.py index 5f727537..1588e38b 100644 --- a/tests/kg/test_kg.py +++ b/tests/kg/test_kg.py @@ -91,6 +91,30 @@ class TestGraphBuilder(unittest.TestCase): graph2 = builder.build(source_list) self.assertEqual(len(graph2["entities"]), 2) + def test_build_with_external_relationship_ids(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + + entities = [ + {"id": "1", "name": "A"}, + {"id": "2", "name": "B"}, + ] + relationships = [ + {"source_id": "1", "target_id": "2", "type": "rel"}, + ] + + source = { + "entities": entities, + "relationships": relationships, + } + + graph = builder.build(source) + + self.assertEqual(len(graph["entities"]), 2) + self.assertEqual(len(graph["relationships"]), 1) + rel = graph["relationships"][0] + self.assertEqual(rel.get("source"), "1") + self.assertEqual(rel.get("target"), "2") + def test_build_with_conflict_resolution(self): """Test building with conflict resolution enabled""" builder = GraphBuilder(resolve_conflicts=True) @@ -106,6 +130,50 @@ class TestGraphBuilder(unittest.TestCase): self.mock_conflict_cls.return_value.detect_conflicts.assert_called_once() self.mock_conflict_cls.return_value.resolve_conflicts.assert_called_once() + def test_build_single_source(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + source = { + "entities": [{"id": "1", "name": "A"}], + "relationships": [{"source_id": "1", "target_id": "1", "type": "self"}], + } + graph = builder.build_single_source(source) + self.assertEqual(len(graph["entities"]), 1) + self.assertEqual(len(graph["relationships"]), 1) + + def test_build_with_explicit_relationships_argument(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + + entities = [ + {"id": "1", "name": "A"}, + {"id": "2", "name": "B"}, + ] + relationships = [ + {"source_id": "1", "target_id": "2", "type": "rel"}, + ] + + graph = builder.build(entities, relationships=relationships) + + self.assertEqual(len(graph["entities"]), 2) + self.assertEqual(len(graph["relationships"]), 1) + rel = graph["relationships"][0] + self.assertEqual(rel.get("source"), "1") + self.assertEqual(rel.get("target"), "2") + + def test_build_warns_when_all_relationships_dropped(self): + builder = GraphBuilder(merge_entities=False, resolve_conflicts=False) + source = { + "entities": [], + "relationships": [{"foo": "x"}, {"bar": "y"}], + } + + with patch.object(builder.logger, "warning") as mock_warning: + graph = builder.build(source) + + self.assertEqual(len(graph["relationships"]), 0) + mock_warning.assert_called() + args, _ = mock_warning.call_args + self.assertIn("All relationships were dropped", args[0]) + class TestGraphAnalyzer(unittest.TestCase): def setUp(self): self.mock_tracker_patcher = patch("semantica.kg.graph_analyzer.get_progress_tracker")