ti.sil/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2025-12-21 08:45:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

fix: prevent ExternalReranker from blocking event loop during RAG queries (#20049)

Browse source 
* fix: prevent ExternalReranker from blocking event loop during RAG queries (#120)

Co-authored-by: Tim Baek <tim@openwebui.com>
Co-authored-by: Claude <noreply@anthropic.com>
Fixes #19900

* Merge pull request open-webui#19030 from open-webui/dev (#122)

Co-authored-by: Tim Baek <tim@openwebui.com>
Co-authored-by: Claude <noreply@anthropic.com>
Fixes #19900

---------

Co-authored-by: Tim Baek <tim@openwebui.com>
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Classic298 2025-12-20 14:43:40 +01:00 committed by GitHub
parent b0d8372d31
commit 2e7c7d635d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 27 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
backend/open_webui/config.py
View file

@ -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",

3
backend/open_webui/main.py
View file

@ -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

3
backend/open_webui/retrieval/models/external.py
View file

@ -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()

2
backend/open_webui/retrieval/utils.py
View file

@ -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

14
backend/open_webui/routers/retrieval.py
View file

@ -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,