Merge remote-tracking branch 'origin/feature' into personal/yfei/fast…

…_bm25

src/pai_rag/app/api/query.py

@@ -41,7 +41,7 @@ async def aupdate(new_config: Any = Body(None)):
 
 
 @router.post("/upload_data")
-async def load_data(input: DataInput, background_tasks: BackgroundTasks):
+def load_data(input: DataInput, background_tasks: BackgroundTasks):
     task_id = uuid.uuid4().hex
     background_tasks.add_task(
         rag_service.add_knowledge_async,

src/pai_rag/app/web/rag_client.py

@@ -4,6 +4,7 @@
 import requests
 import html
 import markdown
+import httpx
 
 cache_config = None
 
@@ -95,11 +96,12 @@ def add_knowledge(self, file_dir: str, enable_qa_extraction: bool):
         response = dotdict(json.loads(r.text))
         return response
 
-    def get_knowledge_state(self, task_id: str):
-        r = requests.get(self.get_load_state_url, params={"task_id": task_id})
-        r.raise_for_status()
-        response = dotdict(json.loads(r.text))
-        return response
+    async def get_knowledge_state(self, task_id: str):
+        async with httpx.AsyncClient() as client:
+            r = await client.get(self.get_load_state_url, params={"task_id": task_id})
+            r.raise_for_status()
+            response = dotdict(json.loads(r.text))
+            return response
 
     def reload_config(self, config: Any):
         global cache_config

src/pai_rag/app/web/tabs/upload_tab.py

@@ -6,6 +6,7 @@
 from pai_rag.app.web.view_model import view_model
 from pai_rag.utils.file_utils import MyUploadFile
 import pandas as pd
+import asyncio
 
 
 def upload_knowledge(upload_files, chunk_size, chunk_overlap, enable_qa_extraction):
@@ -15,7 +16,13 @@ def upload_knowledge(upload_files, chunk_size, chunk_overlap, enable_qa_extracti
     rag_client.reload_config(new_config)
 
     if not upload_files:
-        return "No file selected. Please choose at least one file."
+        yield [
+            gr.update(visible=False),
+            gr.update(
+                visible=True,
+                value="No file selected. Please choose at least one file.",
+            ),
+        ]
 
     my_upload_files = []
     for file in upload_files:
@@ -28,7 +35,7 @@ def upload_knowledge(upload_files, chunk_size, chunk_overlap, enable_qa_extracti
     result = {"Info": ["StartTime", "EndTime", "Duration(s)", "Status"]}
     while not all(file.finished is True for file in my_upload_files):
         for file in my_upload_files:
-            response = rag_client.get_knowledge_state(str(file.task_id))
+            response = asyncio.run(rag_client.get_knowledge_state(str(file.task_id)))
             file.update_state(response["status"])
             file.update_process_duration()
             result[file.file_name] = file.__info__()

src/pai_rag/core/rag_application.py

@@ -34,20 +34,31 @@ def reload(self, config):
         self.logger.info("RagApplication reloaded successfully.")
 
     # TODO: 大量文件上传实现异步添加
-    async def load_knowledge(self, file_dir, enable_qa_extraction=False):
+    async def aload_knowledge(self, file_dir, enable_qa_extraction=False):
         data_loader = module_registry.get_module_with_config(
             "DataLoaderModule", self.config
         )
         await data_loader.aload(file_dir, enable_qa_extraction)
 
+    def load_knowledge(self, file_dir, enable_qa_extraction=False):
+        data_loader = module_registry.get_module_with_config(
+            "DataLoaderModule", self.config
+        )
+        data_loader.load(file_dir, enable_qa_extraction)
+
     async def aquery_retrieval(self, query: RetrievalQuery) -> RetrievalResponse:
         if not query.question:
             return RetrievalResponse(docs=[])
 
+        sessioned_config = self.config
+        if query.vector_db and query.vector_db.faiss_path:
+            sessioned_config = self.config.copy()
+            sessioned_config.index.update({"persist_path": query.vector_db.faiss_path})
+
         query_bundle = QueryBundle(query.question)
 
         query_engine = module_registry.get_module_with_config(
-            "QueryEngineModule", self.config
+            "QueryEngineModule", sessioned_config
         )
         node_results = await query_engine.aretrieve(query_bundle)
 

src/pai_rag/core/rag_service.py

@@ -51,12 +51,12 @@ def reload(self, new_config: Any):
         self.rag.reload(self.rag_configuration.get_value())
         self.rag_configuration.persist()
 
-    async def add_knowledge_async(
+    def add_knowledge_async(
         self, task_id: str, file_dir: str, enable_qa_extraction: bool = False
     ):
         self.tasks_status[task_id] = "processing"
         try:
-            await self.rag.load_knowledge(file_dir, enable_qa_extraction)
+            self.rag.load_knowledge(file_dir, enable_qa_extraction)
             self.tasks_status[task_id] = "completed"
         except Exception as ex:
             logger.error(f"Upload failed: {ex}")

tests/core/test_rag_application.py

@@ -28,9 +28,9 @@ def rag_app():
 
 
 # Test load knowledge file
-async def test_add_knowledge_file(rag_app: RagApplication):
+def test_add_knowledge_file(rag_app: RagApplication):
     data_dir = os.path.join(BASE_DIR, "tests/testdata/paul_graham")
-    await rag_app.load_knowledge(data_dir)
+    rag_app.load_knowledge(data_dir)
 
 
 # Test rag query