feat: make build_kg return list of nodes & edges

Author: ChenZiHong-Gavin

graphgen/models/kg_builder/light_rag_kg_builder.py

@@ -99,7 +99,7 @@ async def merge_nodes(
         self,
         node_data: tuple[str, List[dict]],
         kg_instance: BaseGraphStorage,
-    ) -> None:
+    ) -> dict:
         entity_name, node_data = node_data
         entity_types = []
         source_ids = []
@@ -131,16 +131,18 @@ async def merge_nodes(
 
         node_data = {
             "entity_type": entity_type,
+            "entity_name": entity_name,
             "description": description,
             "source_id": source_id,
         }
         kg_instance.upsert_node(entity_name, node_data=node_data)
+        return node_data
 
     async def merge_edges(
         self,
         edges_data: tuple[Tuple[str, str], List[dict]],
         kg_instance: BaseGraphStorage,
-    ) -> None:
+    ) -> dict:
         (src_id, tgt_id), edge_data = edges_data
 
         source_ids = []
@@ -175,11 +177,19 @@ async def merge_edges(
             f"({src_id}, {tgt_id})", description
         )
 
+        edge_data = {
+            "src_id": src_id,
+            "tgt_id": tgt_id,
+            "description": description,
+            "source_id": source_id,  # for traceability
+        }
+
         kg_instance.upsert_edge(
             src_id,
             tgt_id,
             edge_data={"source_id": source_id, "description": description},
         )
+        return edge_data
 
     async def _handle_kg_summary(
         self,

graphgen/operators/build_kg/build_kg_service.py

@@ -29,7 +29,10 @@ def process(self, batch: pd.DataFrame) -> pd.DataFrame:
 
         # consume the chunks and build kg
         nodes, edges = self.build_kg(docs)
-        return pd.DataFrame([{"nodes": nodes, "edges": edges}])
+        return pd.DataFrame(
+            [{"node": node, "edge": []} for node in nodes]
+            + [{"node": [], "edge": edge} for edge in edges]
+        )
 
     def build_kg(self, chunks: List[Chunk]) -> tuple:
         """
@@ -42,8 +45,8 @@ def build_kg(self, chunks: List[Chunk]) -> tuple:
             if chunk.type in ("image", "video", "table", "formula")
         ]
 
-        nodes = {}
-        edges = {}
+        nodes = []
+        edges = []
 
         if len(text_chunks) == 0:
             logger.info("All text chunks are already in the storage")
@@ -55,8 +58,8 @@ def build_kg(self, chunks: List[Chunk]) -> tuple:
                 chunks=text_chunks,
                 max_loop=self.max_loop,
             )
-            nodes.update(text_nodes)
-            edges.update(text_edges)
+            nodes += text_nodes
+            edges += text_edges
         if len(mm_chunks) == 0:
             logger.info("All multi-modal chunks are already in the storage")
         else:
@@ -66,8 +69,8 @@ def build_kg(self, chunks: List[Chunk]) -> tuple:
                 kg_instance=self.graph_storage,
                 chunks=mm_chunks,
             )
-            nodes.update(mm_nodes)
-            edges.update(mm_edges)
+            nodes += mm_nodes
+            edges += mm_edges
 
         self.graph_storage.index_done_callback()
         logger.info("Knowledge graph building completed.")

graphgen/operators/build_kg/build_mm_kg.py

@@ -37,13 +37,13 @@ def build_mm_kg(
         for k, v in e.items():
             edges[tuple(sorted(k))].extend(v)
 
-    run_concurrent(
+    nodes = run_concurrent(
         lambda kv: mm_builder.merge_nodes(kv, kg_instance=kg_instance),
         list(nodes.items()),
         desc="Inserting entities into storage",
     )
 
-    run_concurrent(
+    edges = run_concurrent(
         lambda kv: mm_builder.merge_edges(kv, kg_instance=kg_instance),
         list(edges.items()),
         desc="Inserting relationships into storage",

graphgen/operators/build_kg/build_text_kg.py

@@ -39,13 +39,13 @@ def build_text_kg(
         for k, v in e.items():
             edges[tuple(sorted(k))].extend(v)
 
-    run_concurrent(
+    nodes = run_concurrent(
         lambda kv: kg_builder.merge_nodes(kv, kg_instance=kg_instance),
         list(nodes.items()),
         desc="Inserting entities into storage",
     )
 
-    run_concurrent(
+    edges = run_concurrent(
         lambda kv: kg_builder.merge_edges(kv, kg_instance=kg_instance),
         list(edges.items()),
         desc="Inserting relationships into storage",

graphgen/operators/quiz/quiz_service.py

@@ -1,5 +1,3 @@
-from collections.abc import Iterable
-
 import pandas as pd
 
 from graphgen.bases import BaseGraphStorage, BaseKVStorage, BaseLLMWrapper, BaseOperator
@@ -28,7 +26,7 @@ def __init__(
         )
         self.generator = QuizGenerator(self.llm_client)
 
-    def process(self, batch: pd.DataFrame) -> Iterable[pd.DataFrame]:
+    def process(self, batch: pd.DataFrame) -> pd.DataFrame:
         data = batch.to_dict(orient="records")
         self.graph_storage.reload()
         return self.quiz(data)
@@ -62,22 +60,22 @@ async def _process_single_quiz(self, item: tuple) -> dict | None:
             logger.error("Error when quizzing description %s: %s", item, e)
             return None
 
-    def quiz(self, batch) -> Iterable[pd.DataFrame]:
+    def quiz(self, batch) -> pd.DataFrame:
         """
         Get all nodes and edges and quiz their descriptions using QuizGenerator.
         """
         items = []
 
         for item in batch:
-            nodes = item.get("nodes", [])
-            edges = item.get("edges", [])
+            node_data = item.get("node", [])
+            edge_data = item.get("edge", [])
 
-            for node_id, node_data in nodes.items():
-                node_data = node_data[0]
+            if node_data:
+                node_id = node_data["entity_name"]
                 desc = node_data["description"]
                 items.append((node_id, desc))
-            for edge_key, edge_data in edges.items():
-                edge_data = edge_data[0]
+            if edge_data:
+                edge_key = (edge_data["src_id"], edge_data["tgt_id"])
                 desc = edge_data["description"]
                 items.append((edge_key, desc))