Make RAG preparation and retrieval non-blocking

api/rag.py

@@ -1,6 +1,6 @@
 import logging
-import weakref
-import re
+import asyncio
+import os
 from dataclasses import dataclass
 from typing import Any, List, Tuple, Dict
 from uuid import uuid4
@@ -48,6 +48,20 @@ def append_dialog_turn(self, dialog_turn):
 # Maximum token limit for embedding models
 MAX_INPUT_TOKENS = 7500  # Safe threshold below 8192 token limit
 
+# Maximum concurrent RAG preparing count
+_RAG_PREPARE_SEMAPHORE: asyncio.Semaphore | None = None
+
+
+def _get_rag_semaphore() -> asyncio.Semaphore:
+    global _RAG_PREPARE_SEMAPHORE
+    if _RAG_PREPARE_SEMAPHORE is None:
+        _RAG_PREPARE_SEMAPHORE = asyncio.Semaphore(
+            int(os.environ.get("DEEPWIKI_MAX_CONCURRENT_RAG", "4"))
+        )
+    assert isinstance(_RAG_PREPARE_SEMAPHORE, asyncio.Semaphore)
+    return _RAG_PREPARE_SEMAPHORE
+
+
 class Memory(adal.core.component.DataComponent):
     """Simple conversation management with a list of dialog turns."""
 
@@ -396,6 +410,44 @@ def prepare_retriever(self, repo_url_or_path: str, type: str = "github", access_
                 logger.error(f"Sample embedding sizes: {', '.join(sizes)}")
             raise
 
+    async def aprepare_retriever(
+            self,
+            repo_url_or_path: str,
+            type: str = "github",
+            access_token: str | None = None,
+            excluded_dirs: list[str] | None = None,
+            excluded_files: list[str] | None = None,
+            included_dirs: list[str] | None = None,
+            included_files: list[str] | None = None,
+    ):
+        """Async version of the original `prepare_retriever`.
+
+        Reuse the synchronous `prepare_retriever` implementation, but runs it in
+        a worker thread via `asyncio.to_thread` so that blocking operations (such
+        as git.clone, file io, embedding calls) do not stall the outer event loop.
+        Concurrency is bounded by a module-level semaphore, set by system variable
+        'DEEPWIKI_MAX_CONCURRENT_RAG'.
+
+        Args:
+            repo_url_or_path: URL or local path to the repository
+            access_token: Optional access token for private repositories
+            excluded_dirs: Optional list of directories to exclude from processing
+            excluded_files: Optional list of file patterns to exclude from processing
+            included_dirs: Optional list of directories to include exclusively
+            included_files: Optional list of file patterns to include exclusively
+        """
+        async with _get_rag_semaphore():
+            return await asyncio.to_thread(
+                self.prepare_retriever,
+                repo_url_or_path,
+                type=type,
+                access_token=access_token,
+                excluded_dirs=excluded_dirs,
+                excluded_files=excluded_files,
+                included_dirs=included_dirs,
+                included_files=included_files,
+            )
+
     def call(self, query: str, language: str = "en") -> Tuple[List]:
         """
         Process a query using RAG.
@@ -426,3 +478,8 @@ def call(self, query: str, language: str = "en") -> Tuple[List]:
                 answer=f"I apologize, but I encountered an error while processing your question. Please try again or rephrase your question."
             )
             return error_response, []
+
+    async def acall(self, query: str, language: str = "en") -> tuple[list]:
+        """Async version of the original `call` method.
+        """
+        return await asyncio.to_thread(self.call, query, language)

api/simple_chat.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from typing import Callable
 from functools import partial
@@ -58,7 +59,12 @@ async def chat_completions_stream(request: ChatCompletionRequest):
 
         # Create a new RAG instance for this request
         try:
-            request_rag = RAG(provider=request.provider, model=request.model)
+            # async do RAG initialization
+            request_rag = await asyncio.to_thread(
+                RAG,
+                provider=request.provider,
+                model=request.model,
+            )
 
             # Extract custom file filter parameters if provided
             if request.excluded_dirs:
@@ -70,7 +76,7 @@ async def chat_completions_stream(request: ChatCompletionRequest):
             if request.included_files:
                 logger.info(f"Using custom included files: {request.included_files}")
 
-            request_rag.prepare_retriever(
+            await request_rag.aprepare_retriever(
                 request.repo_url,
                 request.type,
                 request.token,
@@ -167,7 +173,7 @@ async def chat_completions_stream(request: ChatCompletionRequest):
                 # Try to perform RAG retrieval
                 try:
                     # This will use the actual RAG implementation
-                    retrieved_documents = request_rag(rag_query, language=request.language)
+                    retrieved_documents = await request_rag.acall(rag_query, language=request.language)
 
                     if retrieved_documents and retrieved_documents[0].documents:
                         # Format context for the prompt in a more structured way
@@ -258,7 +264,13 @@ async def chat_completions_stream(request: ChatCompletionRequest):
         file_content = ""
         if request.filePath:
             try:
-                file_content = get_file_content(request.repo_url, request.filePath, request.type, request.token)
+                file_content = await asyncio.to_thread(
+                    get_file_content,
+                    repo_url=request.repo_url,
+                    file_path=request.filePath,
+                    repo_type=request.type,
+                    access_token=request.token,
+                )
                 logger.info(f"Successfully retrieved content for file: {request.filePath}")
             except Exception as e:
                 logger.error(f"Error retrieving file content: {str(e)}")

api/websocket_wiki.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from collections.abc import AsyncIterator, Callable
 from functools import partial
@@ -51,7 +52,11 @@ async def handle_websocket_chat(websocket: WebSocket):
 
         # Create a new RAG instance for this request
         try:
-            request_rag = RAG(provider=request.provider, model=request.model)
+            request_rag = await asyncio.to_thread(
+                RAG,
+                provider=request.provider,
+                model=request.model,
+            )
 
             # Extract custom file filter parameters if provided
             if request.excluded_dirs:
@@ -63,7 +68,7 @@ async def handle_websocket_chat(websocket: WebSocket):
             if request.included_files:
                 logger.info(f"Using custom included files: {request.included_files}")
 
-            request_rag.prepare_retriever(
+            await request_rag.aprepare_retriever(
                 request.repo_url,
                 request.type,
                 request.token,
@@ -170,7 +175,7 @@ async def handle_websocket_chat(websocket: WebSocket):
                 # Try to perform RAG retrieval
                 try:
                     # This will use the actual RAG implementation
-                    retrieved_documents = request_rag(rag_query, language=request.language)
+                    retrieved_documents = await request_rag.acall(rag_query, language=request.language)
 
                     if retrieved_documents and retrieved_documents[0].documents:
                         # Format context for the prompt in a more structured way
@@ -261,7 +266,13 @@ async def handle_websocket_chat(websocket: WebSocket):
         file_content = ""
         if request.filePath:
             try:
-                file_content = get_file_content(request.repo_url, request.filePath, request.type, request.token)
+                file_content = await asyncio.to_thread(
+                    get_file_content,
+                    repo_url=request.repo_url,
+                    file_path=request.filePath,
+                    repo_type=request.type,
+                    access_token=request.token,
+                )
                 logger.info(f"Successfully retrieved content for file: {request.filePath}")
             except Exception as e:
                 logger.error(f"Error retrieving file content: {str(e)}")