feat: enhance RAG service with async file processing and improve logging

Author: Dallas98

runtime/datamate-python/app/module/rag/service/graph_rag.py

@@ -8,7 +8,7 @@
 from lightrag.llm.openai import openai_embed, openai_complete_if_cache
 from lightrag.utils import setup_logger, EmbeddingFunc, get_env_value
 
-setup_logger("lightrag", level="DEBUG")
+setup_logger("lightrag", level="INFO")
 DEFAULT_WORKING_DIR = os.path.join(os.getcwd(), "rag_storage")
 
 

runtime/datamate-python/app/module/rag/service/rag_service.py

@@ -10,7 +10,7 @@
 from app.db.models.dataset_management import DatasetFiles
 from app.db.models.knowledge_gen import RagFile, RagKnowledgeBase
 from app.db.models.model_config import ModelConfig
-from app.db.session import get_db
+from app.db.session import get_db, AsyncSessionLocal
 from app.module.shared.common.document_loaders import load_documents
 from .graph_rag import (
     DEFAULT_WORKING_DIR,
@@ -59,12 +59,22 @@ async def init_graph_rag(self, knowledge_base_id: str):
         kb_working_dir = os.path.join(DEFAULT_WORKING_DIR, kb.name)
         self.rag = await initialize_rag(llm_callable, embedding_callable, kb_working_dir)
 
+        await self._schedule_file_processing(knowledge_base_id)
+
+        return {"status": "initialized", "knowledge_base_id": knowledge_base_id}
+
+    async def _schedule_file_processing(self, knowledge_base_id: str):
         if self.background_tasks is not None:
-            self.background_tasks.add_task(self._process_pending_files, knowledge_base_id)
+            self.background_tasks.add_task(self._process_with_fresh_session, knowledge_base_id, self.rag)
         else:
-            asyncio.create_task(self._process_pending_files(knowledge_base_id))
+            asyncio.create_task(self._process_with_fresh_session(knowledge_base_id, self.rag))
 
-        return {"status": "initialized", "knowledge_base_id": knowledge_base_id}
+    @staticmethod
+    async def _process_with_fresh_session(knowledge_base_id: str, rag_instance):
+        async with AsyncSessionLocal() as session:
+            service = RAGService(session)
+            service.rag = rag_instance
+            await service._process_pending_files(knowledge_base_id)
 
     async def _process_pending_files(self, knowledge_base_id: str):
         rag_files = await self.get_unprocessed_files(knowledge_base_id)
@@ -77,16 +87,17 @@ async def _process_pending_files(self, knowledge_base_id: str):
 
     async def _process_single_file(self, rag_file: RagFile):
         try:
+            await self._mark_file_status(rag_file, "PROCESSING")
             dataset_file = await self._get_dataset_file(rag_file.file_id)
             documents = load_documents(dataset_file.file_path)
             for doc in documents:
                 logger.info(f"Processing document {doc.page_content}")
                 await self.rag.ainsert(input=doc.page_content, file_paths=[dataset_file.file_path])
         except Exception:  # noqa: BLE001
             logger.exception("Failed to process rag file %s", rag_file.id)
-            await self._mark_file_failed(rag_file)
+            await self._mark_file_status(rag_file, "PROCESS_FAILED")
             return
-        await self._mark_file_processed(rag_file)
+        await self._mark_file_status(rag_file, "PROCESSED")
 
     async def _get_dataset_file(self, file_id: str) -> DatasetFiles:
         result = await self.db.execute(
@@ -97,14 +108,8 @@ async def _get_dataset_file(self, file_id: str) -> DatasetFiles:
             raise ValueError(f"Dataset file with ID {file_id} not found.")
         return dataset_file
 
-    async def _mark_file_processed(self, rag_file: RagFile):
-        rag_file.status = "PROCESSED"
-        self.db.add(rag_file)
-        await self.db.commit()
-        await self.db.refresh(rag_file)
-
-    async def _mark_file_failed(self, rag_file: RagFile):
-        rag_file.status = "PROCESS_FAILED"
+    async def _mark_file_status(self, rag_file: RagFile, status: str):
+        rag_file.status = status
         self.db.add(rag_file)
         await self.db.commit()
         await self.db.refresh(rag_file)
@@ -127,7 +132,6 @@ async def _get_model_config(self, model_id: Optional[str]):
             raise ValueError(f"Model config with ID {model_id} not found.")
         return model
 
-
     async def query_rag(self, query: str, knowledge_base_id: str) -> str:
         if not self.rag:
             await self.init_graph_rag(knowledge_base_id)