From abe280f0a3eab3db8717c6417bd94543fc56a756 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Wed, 16 Jul 2025 13:56:02 +0400 Subject: [PATCH] refac/fix: reranking function --- backend/open_webui/routers/retrieval.py | 20 ++++++++++++++------ backend/open_webui/utils/middleware.py | 8 ++++---- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/backend/open_webui/routers/retrieval.py b/backend/open_webui/routers/retrieval.py index 97db0a72f7..18e2b8c09a 100644 --- a/backend/open_webui/routers/retrieval.py +++ b/backend/open_webui/routers/retrieval.py @@ -2050,11 +2050,13 @@ def query_doc_handler( ), k=form_data.k if form_data.k else request.app.state.config.TOP_K, reranking_function=( - lambda sentences: ( - request.app.state.RERANKING_FUNCTION(sentences, user=user) - if request.app.state.RERANKING_FUNCTION - else None + ( + lambda sentences: request.app.state.RERANKING_FUNCTION( + sentences, user=user + ) ) + if request.app.state.RERANKING_FUNCTION + else None ), k_reranker=form_data.k_reranker or request.app.state.config.TOP_K_RERANKER, @@ -2112,8 +2114,14 @@ 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=lambda sentences: request.app.state.RERANKING_FUNCTION( - sentences, user=user + reranking_function=( + ( + lambda sentences: request.app.state.RERANKING_FUNCTION( + sentences, user=user + ) + ) + if request.app.state.RERANKING_FUNCTION + else None ), k_reranker=form_data.k_reranker or request.app.state.config.TOP_K_RERANKER, diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index fc21543457..1d02717530 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -653,13 +653,13 @@ async def chat_completion_files_handler( ), k=request.app.state.config.TOP_K, reranking_function=( - lambda sentences: ( - request.app.state.RERANKING_FUNCTION( + ( + lambda sentences: request.app.state.RERANKING_FUNCTION( sentences, user=user ) - if request.app.state.RERANKING_FUNCTION - else None ) + if request.app.state.RERANKING_FUNCTION + else None ), k_reranker=request.app.state.config.TOP_K_RERANKER, r=request.app.state.config.RELEVANCE_THRESHOLD,