feat(graphgen): add clear func

Author: ChenZiHong-Gavin

graphgen/graphgen.py

@@ -1,4 +1,4 @@
-# https://github.com/HKUDS/LightRAG
+# Adapt from https://github.com/HKUDS/LightRAG
 
 import os
 import asyncio
@@ -17,7 +17,6 @@
 
 sys_path = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
 
-
 @dataclass
 class GraphGen:
     unique_id: int = int(time.time())
@@ -207,3 +206,17 @@ async def async_traverse(self):
                                                    self.graph_storage, self.traverse_strategy, self.text_chunks_storage)
         await self.qa_storage.upsert(results)
         await self.qa_storage.index_done_callback()
+
+    def clear(self):
+        loop = create_event_loop()
+        loop.run_until_complete(self.async_clear())
+
+    async def async_clear(self):
+        await self.full_docs_storage.clear()
+        await self.text_chunks_storage.clear()
+        await self.wiki_storage.clear()
+        await self.graph_storage.clear()
+        await self.rephrase_storage.clear()
+        await self.qa_storage.clear()
+
+        logger.info("All caches are cleared")

graphgen/operators/extract_kg.py

@@ -45,6 +45,7 @@ async def _process_single_content(chunk: Chunk, max_loop: int = 3):
             hint_prompt = KG_EXTRACTION_PROMPT[language]["TEMPLATE"].format(
                 **KG_EXTRACTION_PROMPT["FORMAT"], input_text=content
             )
+            print(hint_prompt, len(hint_prompt))
 
             final_result = await llm_client.generate_answer(hint_prompt)
             logger.info('First result: %s', final_result)

models/storage/json_storage.py

@@ -45,3 +45,7 @@ async def upsert(self, data: dict):
 
     async def drop(self):
         self._data = {}
+
+    async def clear(self):
+        if os.path.exists(self._file_name):
+            os.remove(self._file_name)

models/storage/networkx_storage.py

@@ -150,3 +150,10 @@ async def delete_node(self, node_id: str):
             logger.info("Node %s deleted from the graph.", node_id)
         else:
             logger.warning("Node %s not found in the graph for deletion.", node_id)
+
+    async def clear(self):
+        """
+        Clear the graph by removing all nodes and edges.
+        """
+        self._graph.clear()
+        logger.info("Graph %s cleared.", self.namespace)