fix(graphgen): fix lint errors

Author: ChenZiHong-Gavin

graphgen/operators/__init__.py

@@ -10,4 +10,4 @@
     "judge_statement",
     "search_wikipedia",
     "traverse_graph_by_edge"
-]
+]

graphgen/operators/extract_kg.py

@@ -3,14 +3,14 @@
 
 from typing import List
 from collections import defaultdict
+from tqdm.asyncio import tqdm as tqdm_async
 from models import Chunk, OpenAIModel, Tokenizer
 from models.storage.base_storage import BaseGraphStorage
 from templates import KG_EXTRACTION_PROMPT
-from tqdm.asyncio import tqdm as tqdm_async
 from utils import (logger, pack_history_conversations, split_string_by_multi_markers,
                    handle_single_entity_extraction, handle_single_relationship_extraction,
                    detect_if_chinese)
-from .merge_kg import merge_nodes, merge_edges
+from graphgen.operators.merge_kg import merge_nodes, merge_edges
 
 
 async def extract_kg(
@@ -51,12 +51,18 @@ async def _process_single_content(chunk: Chunk, max_loop: int = 3):
 
             history = pack_history_conversations(hint_prompt, final_result)
             for loop_index in range(max_loop):
-                if_loop_result = await llm_client.generate_answer(text=KG_EXTRACTION_PROMPT[language]["IF_LOOP"], history=history)
+                if_loop_result = await llm_client.generate_answer(
+                    text=KG_EXTRACTION_PROMPT[language]["IF_LOOP"],
+                    history=history
+                )
                 if_loop_result = if_loop_result.strip().strip('"').strip("'").lower()
                 if if_loop_result != "yes":
                     break
 
-                glean_result = await llm_client.generate_answer(text=KG_EXTRACTION_PROMPT[language]["CONTINUE"], history=history)
+                glean_result = await llm_client.generate_answer(
+                    text=KG_EXTRACTION_PROMPT[language]["CONTINUE"],
+                    history=history
+                )
                 logger.info(f"Loop {loop_index} glean: {glean_result}")
 
                 history += pack_history_conversations(KG_EXTRACTION_PROMPT[language]["CONTINUE"], glean_result)
@@ -101,7 +107,7 @@ async def _process_single_content(chunk: Chunk, max_loop: int = 3):
     ):
         try:
             results.append(await result)
-        except Exception as e:
+        except Exception as e: # pylint: disable=broad-except
             logger.error("Error occurred while extracting entities and relationships from chunks: %s", e)
 
     nodes = defaultdict(list)

graphgen/operators/merge_kg.py

@@ -1,11 +1,12 @@
 from collections import Counter
+import asyncio
+from tqdm.asyncio import tqdm as tqdm_async
+
 from utils.format import split_string_by_multi_markers
+from utils import logger, detect_main_language
 from models import TopkTokenModel, Tokenizer
 from models.storage.base_storage import BaseGraphStorage
 from templates import KG_SUMMARIZATION_PROMPT, KG_EXTRACTION_PROMPT
-from utils import logger, detect_main_language
-from tqdm.asyncio import tqdm as tqdm_async
-import asyncio
 
 async def _handle_kg_summary(
     entity_or_relation_name: str,
@@ -100,11 +101,11 @@ async def process_single_node(entity_name: str, node_data: list[dict]):
                 set([dp["source_id"] for dp in node_data] + source_ids)
             )
 
-            node_data = dict(
-                entity_type=entity_type,
-                description=description,
-                source_id=source_id
-            )
+            node_data = {
+                "entity_type": entity_type,
+                "description": description,
+                "source_id": source_id
+            }
             await kg_instance.upsert_node(
                 entity_name,
                 node_data=node_data
@@ -124,7 +125,7 @@ async def process_single_node(entity_name: str, node_data: list[dict]):
     ):
         try:
             entities_data.append(await result)
-        except Exception as e:
+        except Exception as e: # pylint: disable=broad-except
             logger.error("Error occurred while inserting entities into storage: %s", e)
 
 
@@ -168,7 +169,7 @@ async def process_single_edge(src_id: str, tgt_id: str, edge_data: list[dict]):
             )
 
             for insert_id in [src_id, tgt_id]:
-                if not (await kg_instance.has_node(insert_id)):
+                if not await kg_instance.has_node(insert_id):
                     await kg_instance.upsert_node(
                         insert_id,
                         node_data={
@@ -185,17 +186,17 @@ async def process_single_edge(src_id: str, tgt_id: str, edge_data: list[dict]):
             await kg_instance.upsert_edge(
                 src_id,
                 tgt_id,
-                edge_data=dict(
-                    description=description,
-                    source_id=source_id
-                )
+                edge_data={
+                    "source_id": source_id,
+                    "description": description
+                }
             )
 
-            edge_data = dict(
-                src_id=src_id,
-                tgt_id=tgt_id,
-                description=description,
-            )
+            edge_data = {
+                "src_id": src_id,
+                "tgt_id": tgt_id,
+                "description": description
+            }
             return edge_data
 
     logger.info("Inserting relationships into storage...")
@@ -210,5 +211,5 @@ async def process_single_edge(src_id: str, tgt_id: str, edge_data: list[dict]):
     ):
         try:
             relationships_data.append(await result)
-        except Exception as e:
+        except Exception as e: # pylint: disable=broad-except
             logger.error("Error occurred while inserting relationships into storage: %s", e)

models/storage/networkx_storage.py

@@ -126,7 +126,7 @@ async def update_node(self, node_id: str, node_data: dict[str, str]):
         if self._graph.has_node(node_id):
             self._graph.nodes[node_id].update(node_data)
         else:
-            logger.warning(f"Node {node_id} not found in the graph for update.")
+            logger.warning("Node %s not found in the graph for update.", node_id)
 
     async def upsert_edge(
         self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
@@ -147,6 +147,6 @@ async def delete_node(self, node_id: str):
         """
         if self._graph.has_node(node_id):
             self._graph.remove_node(node_id)
-            logger.info(f"Node {node_id} deleted from the graph.")
+            logger.info("Node %s deleted from the graph.", node_id)
         else:
-            logger.warning(f"Node {node_id} not found in the graph for deletion.")
+            logger.warning("Node %s not found in the graph for deletion.", node_id)