refac/enh: forward user info header to reranker

This commit is contained in:
Timothy Jaeryang Baek
2025-07-14 13:59:10 +04:00
parent b4f04ff3a7
commit 0013f5c1fc
5 changed files with 54 additions and 10 deletions

View File

@@ -70,6 +70,7 @@ from open_webui.retrieval.web.external import search_external
from open_webui.retrieval.utils import (
get_embedding_function,
get_reranking_function,
get_model_path,
query_collection,
query_collection_with_hybrid_search,
@@ -824,6 +825,12 @@ async def update_rag_config(
request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY,
True,
)
request.app.state.RERANKING_FUNCTION = get_reranking_function(
request.app.state.config.RAG_RERANKING_ENGINE,
request.app.state.config.RAG_RERANKING_MODEL,
request.app.state.rf,
)
except Exception as e:
log.error(f"Error loading reranking model: {e}")
request.app.state.config.ENABLE_RAG_HYBRID_SEARCH = False
@@ -2042,7 +2049,9 @@ def query_doc_handler(
query, prefix=prefix, user=user
),
k=form_data.k if form_data.k else request.app.state.config.TOP_K,
reranking_function=request.app.state.rf,
reranking_function=lambda sentences: request.app.state.RERANKING_FUNCTION(
sentences, user=user
),
k_reranker=form_data.k_reranker
or request.app.state.config.TOP_K_RERANKER,
r=(
@@ -2099,7 +2108,9 @@ def query_collection_handler(
query, prefix=prefix, user=user
),
k=form_data.k if form_data.k else request.app.state.config.TOP_K,
reranking_function=request.app.state.rf,
reranking_function=lambda sentences: request.app.state.RERANKING_FUNCTION(
sentences, user=user
),
k_reranker=form_data.k_reranker
or request.app.state.config.TOP_K_RERANKER,
r=(