Feat: Adjusted to handle individual rag config - load all the embedders and rerankers that are currently used

This commit is contained in:
weberm1 2025-05-23 10:38:35 +02:00
parent 904e17c470
commit 79bf347dfe

View File

@ -224,6 +224,11 @@ from open_webui.config import (
PDF_EXTRACT_IMAGES, PDF_EXTRACT_IMAGES,
YOUTUBE_LOADER_LANGUAGE, YOUTUBE_LOADER_LANGUAGE,
YOUTUBE_LOADER_PROXY_URL, YOUTUBE_LOADER_PROXY_URL,
DEFAULT_RAG_SETTINGS,
DOWNLOADED_EMBEDDING_MODELS,
DOWNLOADED_RERANKING_MODELS,
LOADED_EMBEDDING_MODELS,
LOADED_RERANKING_MODELS,
# Retrieval (Web Search) # Retrieval (Web Search)
ENABLE_WEB_SEARCH, ENABLE_WEB_SEARCH,
WEB_SEARCH_ENGINE, WEB_SEARCH_ENGINE,
@ -741,6 +746,11 @@ app.state.config.EXTERNAL_WEB_SEARCH_API_KEY = EXTERNAL_WEB_SEARCH_API_KEY
app.state.config.EXTERNAL_WEB_LOADER_URL = EXTERNAL_WEB_LOADER_URL app.state.config.EXTERNAL_WEB_LOADER_URL = EXTERNAL_WEB_LOADER_URL
app.state.config.EXTERNAL_WEB_LOADER_API_KEY = EXTERNAL_WEB_LOADER_API_KEY app.state.config.EXTERNAL_WEB_LOADER_API_KEY = EXTERNAL_WEB_LOADER_API_KEY
app.state.config.DEFAULT_RAG_SETTINGS = DEFAULT_RAG_SETTINGS
app.state.config.DOWNLOADED_EMBEDDING_MODELS = DOWNLOADED_EMBEDDING_MODELS
app.state.config.DOWNLOADED_RERANKING_MODELS = DOWNLOADED_RERANKING_MODELS
app.state.config.LOADED_EMBEDDING_MODELS = LOADED_EMBEDDING_MODELS
app.state.config.LOADED_RERANKING_MODELS = LOADED_RERANKING_MODELS
app.state.config.PLAYWRIGHT_WS_URL = PLAYWRIGHT_WS_URL app.state.config.PLAYWRIGHT_WS_URL = PLAYWRIGHT_WS_URL
app.state.config.PLAYWRIGHT_TIMEOUT = PLAYWRIGHT_TIMEOUT app.state.config.PLAYWRIGHT_TIMEOUT = PLAYWRIGHT_TIMEOUT
@ -748,49 +758,53 @@ app.state.config.FIRECRAWL_API_BASE_URL = FIRECRAWL_API_BASE_URL
app.state.config.FIRECRAWL_API_KEY = FIRECRAWL_API_KEY app.state.config.FIRECRAWL_API_KEY = FIRECRAWL_API_KEY
app.state.config.TAVILY_EXTRACT_DEPTH = TAVILY_EXTRACT_DEPTH app.state.config.TAVILY_EXTRACT_DEPTH = TAVILY_EXTRACT_DEPTH
app.state.EMBEDDING_FUNCTION = None app.state.EMBEDDING_FUNCTION = {}
app.state.ef = None app.state.ef = {}
app.state.rf = None app.state.rf = {}
app.state.YOUTUBE_LOADER_TRANSLATION = None app.state.YOUTUBE_LOADER_TRANSLATION = None
try: try:
app.state.ef = get_ef( # Load all embedding models that are currently in use
app.state.config.RAG_EMBEDDING_ENGINE, for engine, model_list in app.state.config.LOADED_EMBEDDING_MODELS.items():
app.state.config.RAG_EMBEDDING_MODEL, for model in model_list:
RAG_EMBEDDING_MODEL_AUTO_UPDATE, app.state.ef[model] = get_ef(
) engine,
model,
RAG_EMBEDDING_MODEL_AUTO_UPDATE,
)
app.state.EMBEDDING_FUNCTION[model] = get_embedding_function(
engine,
model,
app.state.ef[model],
(
app.state.config.RAG_OPENAI_API_BASE_URL
if engine == "openai"
else app.state.config.RAG_OLLAMA_BASE_URL
),
(
app.state.config.RAG_OPENAI_API_KEY
if engine == "openai"
else app.state.config.RAG_OLLAMA_API_KEY
),
app.state.config.RAG_EMBEDDING_BATCH_SIZE,
)
app.state.rf = get_rf( # Load all reranking models that are currently in use
app.state.config.RAG_RERANKING_ENGINE, for engine, model_list in app.state.config.LOADED_RERANKING_MODELS.items():
app.state.config.RAG_RERANKING_MODEL, for model in model_list:
app.state.config.RAG_EXTERNAL_RERANKER_URL, app.state.rf[model["RAG_RERANKING_MODEL"]] = get_rf(
app.state.config.RAG_EXTERNAL_RERANKER_API_KEY, engine,
RAG_RERANKING_MODEL_AUTO_UPDATE, model["RAG_RERANKING_MODEL"],
) model["RAG_EXTERNAL_RERANKER_URL"],
model["RAG_EXTERNAL_RERANKER_API_KEY"],
)
except Exception as e: except Exception as e:
log.error(f"Error updating models: {e}") log.error(f"Error updating models: {e}")
pass pass
app.state.EMBEDDING_FUNCTION = get_embedding_function(
app.state.config.RAG_EMBEDDING_ENGINE,
app.state.config.RAG_EMBEDDING_MODEL,
app.state.ef,
(
app.state.config.RAG_OPENAI_API_BASE_URL
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
else app.state.config.RAG_OLLAMA_BASE_URL
),
(
app.state.config.RAG_OPENAI_API_KEY
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
else app.state.config.RAG_OLLAMA_API_KEY
),
app.state.config.RAG_EMBEDDING_BATCH_SIZE,
)
######################################## ########################################
# #
# CODE EXECUTION # CODE EXECUTION