From 2e7c7d635d86376fc3b5095cee3663fce289ba05 Mon Sep 17 00:00:00 2001 From: Classic298 <27028174+Classic298@users.noreply.github.com> Date: Sat, 20 Dec 2025 14:43:40 +0100 Subject: [PATCH] fix: prevent ExternalReranker from blocking event loop during RAG queries (#20049) * fix: prevent ExternalReranker from blocking event loop during RAG queries (#120) Co-authored-by: Tim Baek Co-authored-by: Claude Fixes #19900 * Merge pull request open-webui#19030 from open-webui/dev (#122) Co-authored-by: Tim Baek Co-authored-by: Claude Fixes #19900 --------- Co-authored-by: Tim Baek Co-authored-by: Claude --- backend/open_webui/config.py | 6 ++++++ backend/open_webui/main.py | 3 +++ backend/open_webui/retrieval/models/external.py | 3 +++ backend/open_webui/retrieval/utils.py | 2 +- backend/open_webui/routers/retrieval.py | 14 ++++++++++++++ 5 files changed, 27 insertions(+), 1 deletion(-) diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index 1e17b9b1d1..cf2983d63f 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -2802,6 +2802,12 @@ RAG_EXTERNAL_RERANKER_API_KEY = PersistentConfig( os.environ.get("RAG_EXTERNAL_RERANKER_API_KEY", ""), ) +RAG_EXTERNAL_RERANKER_TIMEOUT = PersistentConfig( + "RAG_EXTERNAL_RERANKER_TIMEOUT", + "rag.external_reranker_timeout", + os.environ.get("RAG_EXTERNAL_RERANKER_TIMEOUT", ""), +) + RAG_TEXT_SPLITTER = PersistentConfig( "RAG_TEXT_SPLITTER", diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index ea8e20fe48..b3dac0aed8 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -227,6 +227,7 @@ from open_webui.config import ( RAG_RERANKING_MODEL, RAG_EXTERNAL_RERANKER_URL, RAG_EXTERNAL_RERANKER_API_KEY, + RAG_EXTERNAL_RERANKER_TIMEOUT, RAG_RERANKING_MODEL_AUTO_UPDATE, RAG_RERANKING_MODEL_TRUST_REMOTE_CODE, RAG_EMBEDDING_ENGINE, @@ -898,6 +899,7 @@ app.state.config.RAG_RERANKING_ENGINE = RAG_RERANKING_ENGINE app.state.config.RAG_RERANKING_MODEL = RAG_RERANKING_MODEL app.state.config.RAG_EXTERNAL_RERANKER_URL = RAG_EXTERNAL_RERANKER_URL app.state.config.RAG_EXTERNAL_RERANKER_API_KEY = RAG_EXTERNAL_RERANKER_API_KEY +app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT = RAG_EXTERNAL_RERANKER_TIMEOUT app.state.config.RAG_TEMPLATE = RAG_TEMPLATE @@ -1000,6 +1002,7 @@ try: app.state.config.RAG_RERANKING_MODEL, app.state.config.RAG_EXTERNAL_RERANKER_URL, app.state.config.RAG_EXTERNAL_RERANKER_API_KEY, + app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT, ) else: app.state.rf = None diff --git a/backend/open_webui/retrieval/models/external.py b/backend/open_webui/retrieval/models/external.py index d5b020602d..d567bf4fe5 100644 --- a/backend/open_webui/retrieval/models/external.py +++ b/backend/open_webui/retrieval/models/external.py @@ -18,10 +18,12 @@ class ExternalReranker(BaseReranker): api_key: str, url: str = "http://localhost:8080/v1/rerank", model: str = "reranker", + timeout: Optional[int] = None, ): self.api_key = api_key self.url = url self.model = model + self.timeout = timeout def predict( self, sentences: List[Tuple[str, str]], user=None @@ -52,6 +54,7 @@ class ExternalReranker(BaseReranker): f"{self.url}", headers=headers, json=payload, + timeout=self.timeout, ) r.raise_for_status() diff --git a/backend/open_webui/retrieval/utils.py b/backend/open_webui/retrieval/utils.py index 9f4ce210a9..88fa51f2dd 100644 --- a/backend/open_webui/retrieval/utils.py +++ b/backend/open_webui/retrieval/utils.py @@ -1279,7 +1279,7 @@ class RerankCompressor(BaseDocumentCompressor): scores = None if reranking: - scores = self.reranking_function(query, documents) + scores = await asyncio.to_thread(self.reranking_function, query, documents) else: from sentence_transformers import util diff --git a/backend/open_webui/routers/retrieval.py b/backend/open_webui/routers/retrieval.py index 3f7b647169..2bc0b00649 100644 --- a/backend/open_webui/routers/retrieval.py +++ b/backend/open_webui/routers/retrieval.py @@ -146,9 +146,12 @@ def get_rf( reranking_model: Optional[str] = None, external_reranker_url: str = "", external_reranker_api_key: str = "", + external_reranker_timeout: str = "", auto_update: bool = RAG_RERANKING_MODEL_AUTO_UPDATE, ): rf = None + # Convert timeout string to int or None (system default) + timeout_value = int(external_reranker_timeout) if external_reranker_timeout else None if reranking_model: if any(model in reranking_model for model in ["jinaai/jina-colbert-v2"]): try: @@ -171,6 +174,7 @@ def get_rf( url=external_reranker_url, api_key=external_reranker_api_key, model=reranking_model, + timeout=timeout_value, ) except Exception as e: log.error(f"ExternalReranking: {e}") @@ -480,6 +484,7 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)): "RAG_RERANKING_ENGINE": request.app.state.config.RAG_RERANKING_ENGINE, "RAG_EXTERNAL_RERANKER_URL": request.app.state.config.RAG_EXTERNAL_RERANKER_URL, "RAG_EXTERNAL_RERANKER_API_KEY": request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY, + "RAG_EXTERNAL_RERANKER_TIMEOUT": request.app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT, # Chunking settings "TEXT_SPLITTER": request.app.state.config.TEXT_SPLITTER, "CHUNK_SIZE": request.app.state.config.CHUNK_SIZE, @@ -667,6 +672,7 @@ class ConfigForm(BaseModel): RAG_RERANKING_ENGINE: Optional[str] = None RAG_EXTERNAL_RERANKER_URL: Optional[str] = None RAG_EXTERNAL_RERANKER_API_KEY: Optional[str] = None + RAG_EXTERNAL_RERANKER_TIMEOUT: Optional[str] = None # Chunking settings TEXT_SPLITTER: Optional[str] = None @@ -923,6 +929,12 @@ async def update_rag_config( else request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY ) + request.app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT = ( + form_data.RAG_EXTERNAL_RERANKER_TIMEOUT + if form_data.RAG_EXTERNAL_RERANKER_TIMEOUT is not None + else request.app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT + ) + log.info( f"Updating reranking model: {request.app.state.config.RAG_RERANKING_MODEL} to {form_data.RAG_RERANKING_MODEL}" ) @@ -943,6 +955,7 @@ async def update_rag_config( request.app.state.config.RAG_RERANKING_MODEL, request.app.state.config.RAG_EXTERNAL_RERANKER_URL, request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY, + request.app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT, ) request.app.state.RERANKING_FUNCTION = get_reranking_function( @@ -1164,6 +1177,7 @@ async def update_rag_config( "RAG_RERANKING_ENGINE": request.app.state.config.RAG_RERANKING_ENGINE, "RAG_EXTERNAL_RERANKER_URL": request.app.state.config.RAG_EXTERNAL_RERANKER_URL, "RAG_EXTERNAL_RERANKER_API_KEY": request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY, + "RAG_EXTERNAL_RERANKER_TIMEOUT": request.app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT, # Chunking settings "TEXT_SPLITTER": request.app.state.config.TEXT_SPLITTER, "CHUNK_SIZE": request.app.state.config.CHUNK_SIZE,