fix: fix some lint errors

Author: ChenZiHong-Gavin

graphgen/graphgen.py

@@ -8,9 +8,9 @@
 from tqdm.asyncio import tqdm as tqdm_async
 
 from models import Chunk, JsonKVStorage, OpenAIModel, NetworkXStorage, WikiSearch, Tokenizer, TraverseStrategy
-from utils import create_event_loop, logger, compute_content_hash
 from models.storage.base_storage import StorageNameSpace
-from .operators import *
+from utils import create_event_loop, logger, compute_content_hash
+from .operators import extract_kg, search_wikipedia, quiz_relations, judge_relations, traverse_graph_by_edge
 
 
 sys_path = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
@@ -145,7 +145,7 @@ async def async_insert(self, data: Union[List[list], List[dict]], data_type: str
 
         logger.info(f"[Wiki Search] is {'enabled' if self.if_web_search else 'disabled'}")
         if self.if_web_search:
-            logger.info(f"[Wiki Search]...")
+            logger.info("[Wiki Search]...")
             _add_wiki_data = await search_wikipedia(
                 llm_client= self.teacher_llm_client,
                 wiki_search_client=self.wiki_client,

models/storage/json_storage.py

@@ -12,7 +12,7 @@ class JsonKVStorage(BaseKVStorage):
     def __post_init__(self):
         self._file_name = os.path.join(self.working_dir, f"{self.namespace}.json")
         self._data = load_json(self._file_name) or {}
-        logger.info(f"Load KV {self.namespace} with {len(self._data)} data")
+        logger.info("Load KV %s with %d data", self.namespace, len(self._data))
 
     async def all_keys(self) -> list[str]:
         return list(self._data.keys())

models/storage/networkx_storage.py

@@ -1,12 +1,11 @@
 import os
 import html
-import networkx as nx
-
 from typing import Any, Union, cast, Optional
 from dataclasses import dataclass
+import networkx as nx
 
-from .base_storage import BaseGraphStorage
 from utils import logger
+from .base_storage import BaseGraphStorage
 
 @dataclass
 class NetworkXStorage(BaseGraphStorage):
@@ -120,6 +119,9 @@ async def get_node_edges(self, source_node_id: str) -> Union[list[tuple[str, str
             return list(self._graph.edges(source_node_id, data=True))
         return None
 
+    async def get_graph(self) -> nx.Graph:
+        return self._graph
+
     async def upsert_node(self, node_id: str, node_data: dict[str, str]):
         self._graph.add_node(node_id, **node_data)