fix: fix lint errors

Author: ChenZiHong-Gavin

graphgen/graphgen.py

@@ -108,16 +108,16 @@ async def insert(self, read_config: Dict, split_config: Dict):
             self.progress_bar,
         )
 
-        # _add_chunk_keys = await self.chunks_storage.filter_keys(
-        #     list(inserting_chunks.keys())
-        # )
-        # inserting_chunks = {
-        #     k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys
-        # }
-        #
-        # if len(inserting_chunks) == 0:
-        #     logger.warning("All chunks are already in the storage")
-        #     return
+        _add_chunk_keys = await self.chunks_storage.filter_keys(
+            list(inserting_chunks.keys())
+        )
+        inserting_chunks = {
+            k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys
+        }
+
+        if len(inserting_chunks) == 0:
+            logger.warning("All chunks are already in the storage")
+            return
 
         logger.info("[New Chunks] inserting %d chunks", len(inserting_chunks))
         await self.chunks_storage.upsert(inserting_chunks)

graphgen/operators/build_kg/build_kg.py

@@ -2,9 +2,9 @@
 
 import gradio as gr
 
+from graphgen.bases import BaseLLMWrapper
 from graphgen.bases.base_storage import BaseGraphStorage
 from graphgen.bases.datatypes import Chunk
-from graphgen.models import OpenAIClient
 from graphgen.utils import logger
 
 from .build_mm_kg import build_mm_kg
@@ -13,7 +13,7 @@
 
 
 async def build_kg(
-    llm_client: OpenAIClient,
+    llm_client: BaseLLMWrapper,
     kg_instance: BaseGraphStorage,
     chunks: List[Chunk],
     anchor_type: Optional[str] = None,

graphgen/operators/build_kg/build_mo_kg.py

@@ -1,25 +1,17 @@
 import json
-import re
 from typing import List
 
 import gradio as gr
 
+from graphgen.bases import BaseLLMWrapper
 from graphgen.bases.base_storage import BaseGraphStorage
 from graphgen.bases.datatypes import Chunk
-from graphgen.models import OpenAIClient
 from graphgen.templates import PROTEIN_ANCHOR_PROMPT, PROTEIN_KG_EXTRACTION_PROMPT
-from graphgen.utils import (
-    detect_main_language,
-    handle_single_entity_extraction,
-    handle_single_relationship_extraction,
-    logger,
-    run_concurrent,
-    split_string_by_multi_markers,
-)
+from graphgen.utils import detect_main_language, logger, run_concurrent
 
 
 async def build_mo_kg(
-    llm_client: OpenAIClient,
+    llm_client: BaseLLMWrapper,
     kg_instance: BaseGraphStorage,
     chunks: List[Chunk],
     progress_bar: gr.Progress = None,
@@ -73,48 +65,13 @@ async def extract_mo_info(chunk: Chunk):
     #     logger.warning("Failed to search for protein info: %s", e)
     #     search_results = {}
 
-    # 组织成文本
     mo_text = "\n".join([f"{k}: {v}" for k, v in merged.items()])
     lang = detect_main_language(mo_text)
     prompt = PROTEIN_KG_EXTRACTION_PROMPT[lang].format(
         input_text=mo_text,
         **PROTEIN_KG_EXTRACTION_PROMPT["FORMAT"],
     )
     kg_output = await llm_client.generate_answer(prompt)
-
-    logger.debug("Image chunk extraction result: %s", kg_output)
-
-    # parse the result
-    records = split_string_by_multi_markers(
-        kg_output,
-        [
-            PROTEIN_KG_EXTRACTION_PROMPT["FORMAT"]["record_delimiter"],
-            PROTEIN_KG_EXTRACTION_PROMPT["FORMAT"]["completion_delimiter"],
-        ],
-    )
-
-    print(records)
-    raise NotImplementedError
-
-    nodes = defaultdict(list)
-    edges = defaultdict(list)
-
-    for record in records:
-        match = re.search(r"\((.*)\)", record)
-        if not match:
-            continue
-        inner = match.group(1)
-
-        attributes = split_string_by_multi_markers(
-            inner, [PROTEIN_KG_EXTRACTION_PROMPT["FORMAT"]["tuple_delimiter"]]
-        )
-
-        entity = await handle_single_entity_extraction(attributes, "temp")
-        if entity is not None:
-            nodes[entity["entity_name"]].append(entity)
-            continue
-
-        relation = await handle_single_relationship_extraction(attributes, "temp")
-        if relation is not None:
-            key = (relation["src_id"], relation["tgt_id"])
-            edges[key].append(relation)
+    print(kg_output)
+    # TODO: parse kg_output and insert into kg_instance
+    return kg_instance