Merge pull request #209 from Hawksight-AI/kg

Fix GraphBuilder External Relationships (#208, #206)

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,

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")

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")