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

@@ -89,6 +89,7 @@ from open_webui.routers import (
from open_webui.routers.retrieval import (
get_embedding_function,
get_reranking_function,
get_ef,
get_rf,
)
@@ -878,6 +879,7 @@ app.state.config.FIRECRAWL_API_KEY = FIRECRAWL_API_KEY
app.state.config.TAVILY_EXTRACT_DEPTH = TAVILY_EXTRACT_DEPTH
app.state.EMBEDDING_FUNCTION = None
app.state.RERANKING_FUNCTION = None
app.state.ef = None
app.state.rf = None
@@ -906,8 +908,8 @@ except Exception as e:
app.state.EMBEDDING_FUNCTION = get_embedding_function(
app.state.config.RAG_EMBEDDING_ENGINE,
app.state.config.RAG_EMBEDDING_MODEL,
app.state.ef,
(
embedding_function=app.state.ef,
url=(
app.state.config.RAG_OPENAI_API_BASE_URL
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
else (
@@ -916,7 +918,7 @@ app.state.EMBEDDING_FUNCTION = get_embedding_function(
else app.state.config.RAG_AZURE_OPENAI_BASE_URL
)
),
(
key=(
app.state.config.RAG_OPENAI_API_KEY
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
else (
@@ -925,7 +927,7 @@ app.state.EMBEDDING_FUNCTION = get_embedding_function(
else app.state.config.RAG_AZURE_OPENAI_API_KEY
)
),
app.state.config.RAG_EMBEDDING_BATCH_SIZE,
embedding_batch_size=app.state.config.RAG_EMBEDDING_BATCH_SIZE,
azure_api_version=(
app.state.config.RAG_AZURE_OPENAI_API_VERSION
if app.state.config.RAG_EMBEDDING_ENGINE == "azure_openai"
@@ -933,6 +935,12 @@ app.state.EMBEDDING_FUNCTION = get_embedding_function(
),
)
app.state.RERANKING_FUNCTION = get_reranking_function(
app.state.config.RAG_RERANKING_ENGINE,
app.state.config.RAG_RERANKING_MODEL,
reranking_function=app.state.rf,
)
########################################
#
# CODE EXECUTION