feat: add query endpoint for knowledge graph and enhance RAG service with async processing

Author: Dallas98

runtime/datamate-python/app/module/rag/interface/rag_interface.py

@@ -4,6 +4,7 @@
 from app.db.session import get_db
 from app.module.rag.service.rag_service import RAGService
 from app.module.shared.schema import StandardResponse
+from ..schema.rag_schema import QueryRequest
 
 router = APIRouter(prefix="/rag", tags=["rag"])
 
@@ -22,3 +23,16 @@ async def process_knowledge_base(knowledge_base_id: str, db: AsyncSession = Depe
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+@router.post("/query")
+async def query_knowledge_graph(payload: QueryRequest, db: AsyncSession = Depends(get_db)):
+    """
+    Query the knowledge graph with the given query text and knowledge base ID.
+    """
+    try:
+        rag_service = RAGService(db)
+        result = await rag_service.query_rag(payload.query, payload.knowledge_base_id)
+        return StandardResponse(code=200, message="success", data=result)
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

runtime/datamate-python/app/module/rag/schema/rag_schema.py

@@ -3,3 +3,6 @@
 class ProcessRequest(BaseModel):
     knowledge_base_id: str
 
+class QueryRequest(BaseModel):
+    knowledge_base_id: str
+    query: str

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

@@ -1,4 +1,5 @@
 import os
+import asyncio
 from typing import Optional, Sequence
 
 from fastapi import BackgroundTasks, Depends
@@ -61,7 +62,7 @@ async def init_graph_rag(self, knowledge_base_id: str):
         if self.background_tasks is not None:
             self.background_tasks.add_task(self._process_pending_files, knowledge_base_id)
         else:
-            await self._process_pending_files(knowledge_base_id)
+            asyncio.create_task(self._process_pending_files(knowledge_base_id))
 
         return {"status": "initialized", "knowledge_base_id": knowledge_base_id}
 
@@ -125,3 +126,9 @@ async def _get_model_config(self, model_id: Optional[str]):
         if not model:
             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)
+        return await self.rag.get_knowledge_graph(query)