Optimization logic

lightrag/lightrag.py

@@ -12,9 +12,8 @@ from .llm import (
 from .operate import (
     chunking_by_token_size,
     extract_entities,
-    local_query,
-    global_query,
-    hybrid_query,
+    # local_query,global_query,hybrid_query,
+    kg_query,
     naive_query,
 )
 
@@ -309,28 +308,8 @@ class LightRAG:
         return loop.run_until_complete(self.aquery(query, param))
 
     async def aquery(self, query: str, param: QueryParam = QueryParam()):
-        if param.mode == "local":
-            response = await local_query(
-                query,
-                self.chunk_entity_relation_graph,
-                self.entities_vdb,
-                self.relationships_vdb,
-                self.text_chunks,
-                param,
-                asdict(self),
-            )
-        elif param.mode == "global":
-            response = await global_query(
-                query,
-                self.chunk_entity_relation_graph,
-                self.entities_vdb,
-                self.relationships_vdb,
-                self.text_chunks,
-                param,
-                asdict(self),
-            )
-        elif param.mode == "hybrid":
-            response = await hybrid_query(
+        if param.mode in ["local", "global", "hybrid"]:
+            response = await kg_query(
                 query,
                 self.chunk_entity_relation_graph,
                 self.entities_vdb,