diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index 1162fde22..b05cb9641 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -1685,6 +1685,11 @@ BYPASS_EMBEDDING_AND_RETRIEVAL = PersistentConfig( RAG_TOP_K = PersistentConfig( "RAG_TOP_K", "rag.top_k", int(os.environ.get("RAG_TOP_K", "3")) ) +RAG_TOP_K_RERANKER = PersistentConfig( + "RAG_TOP_K_RERANKER", + "rag.top_k_reranker", + int(os.environ.get("RAG_TOP_K_RERANKER", "3")) +) RAG_RELEVANCE_THRESHOLD = PersistentConfig( "RAG_RELEVANCE_THRESHOLD", "rag.relevance_threshold", diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 674926055..3fbeb6c84 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -191,6 +191,7 @@ from open_webui.config import ( DOCUMENT_INTELLIGENCE_ENDPOINT, DOCUMENT_INTELLIGENCE_KEY, RAG_TOP_K, + RAG_TOP_K_RERANKER, RAG_TEXT_SPLITTER, TIKTOKEN_ENCODING_NAME, PDF_EXTRACT_IMAGES, @@ -552,6 +553,7 @@ app.state.FUNCTIONS = {} app.state.config.TOP_K = RAG_TOP_K +app.state.config.TOP_K_RERANKER = RAG_TOP_K_RERANKER app.state.config.RELEVANCE_THRESHOLD = RAG_RELEVANCE_THRESHOLD app.state.config.FILE_MAX_SIZE = RAG_FILE_MAX_SIZE app.state.config.FILE_MAX_COUNT = RAG_FILE_MAX_COUNT diff --git a/backend/open_webui/retrieval/utils.py b/backend/open_webui/retrieval/utils.py index b05057b28..3d6e739f3 100644 --- a/backend/open_webui/retrieval/utils.py +++ b/backend/open_webui/retrieval/utils.py @@ -106,6 +106,7 @@ def query_doc_with_hybrid_search( embedding_function, k: int, reranking_function, + k_reranker: int, r: float, ) -> dict: try: @@ -128,7 +129,7 @@ def query_doc_with_hybrid_search( ) compressor = RerankCompressor( embedding_function=embedding_function, - top_n=k, + top_n=k_reranker, reranking_function=reranking_function, r_score=r, ) @@ -138,10 +139,20 @@ def query_doc_with_hybrid_search( ) result = compression_retriever.invoke(query) + + distances = [d.metadata.get("score") for d in result] + documents = [d.page_content for d in result] + metadatas = [d.metadata for d in result] + + # retrieve only min(k, k_reranker) items, sort and cut by distance if k < k_reranker + if k < k_reranker: + sorted_items = sorted(zip(distances, metadatas, documents), key=lambda x: x[0], reverse=True) + sorted_items = sorted_items[:k] + distances, documents, metadatas = map(list, zip(*sorted_items)) result = { - "distances": [[d.metadata.get("score") for d in result]], - "documents": [[d.page_content for d in result]], - "metadatas": [[d.metadata for d in result]], + "distances": [distances], + "documents": [documents], + "metadatas": [metadatas], } log.info( @@ -281,6 +292,7 @@ def query_collection_with_hybrid_search( embedding_function, k: int, reranking_function, + k_reranker: int, r: float, ) -> dict: results = [] @@ -294,6 +306,7 @@ def query_collection_with_hybrid_search( embedding_function=embedding_function, k=k, reranking_function=reranking_function, + k_reranker=k_reranker, r=r, ) results.append(result) @@ -307,7 +320,7 @@ def query_collection_with_hybrid_search( raise Exception( "Hybrid search failed for all collections. Using Non hybrid search as fallback." ) - + return merge_and_sort_query_results(results, k=k, reverse=True) @@ -354,6 +367,7 @@ def get_sources_from_files( embedding_function, k, reranking_function, + k_reranker, r, hybrid_search, full_context=False, @@ -470,6 +484,7 @@ def get_sources_from_files( embedding_function=embedding_function, k=k, reranking_function=reranking_function, + k_reranker=k_reranker, r=r, ) except Exception as e: diff --git a/backend/open_webui/routers/retrieval.py b/backend/open_webui/routers/retrieval.py index 85ffdde74..3f5cb72cf 100644 --- a/backend/open_webui/routers/retrieval.py +++ b/backend/open_webui/routers/retrieval.py @@ -719,6 +719,7 @@ async def get_query_settings(request: Request, user=Depends(get_admin_user)): "status": True, "template": request.app.state.config.RAG_TEMPLATE, "k": request.app.state.config.TOP_K, + "k_reranker": request.app.state.config.TOP_K_RERANKER, "r": request.app.state.config.RELEVANCE_THRESHOLD, "hybrid": request.app.state.config.ENABLE_RAG_HYBRID_SEARCH, } @@ -726,6 +727,7 @@ async def get_query_settings(request: Request, user=Depends(get_admin_user)): class QuerySettingsForm(BaseModel): k: Optional[int] = None + k_reranker: Optional[int] = None r: Optional[float] = None template: Optional[str] = None hybrid: Optional[bool] = None @@ -737,6 +739,7 @@ async def update_query_settings( ): request.app.state.config.RAG_TEMPLATE = form_data.template request.app.state.config.TOP_K = form_data.k if form_data.k else 4 + request.app.state.config.TOP_K_RERANKER = form_data.k_reranker or 4 request.app.state.config.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0 request.app.state.config.ENABLE_RAG_HYBRID_SEARCH = ( @@ -747,6 +750,7 @@ async def update_query_settings( "status": True, "template": request.app.state.config.RAG_TEMPLATE, "k": request.app.state.config.TOP_K, + "k_reranker": request.app.state.config.TOP_K_RERANKER, "r": request.app.state.config.RELEVANCE_THRESHOLD, "hybrid": request.app.state.config.ENABLE_RAG_HYBRID_SEARCH, } @@ -1495,6 +1499,7 @@ class QueryDocForm(BaseModel): collection_name: str query: str k: Optional[int] = None + k_reranker: Optional[int] = None r: Optional[float] = None hybrid: Optional[bool] = None @@ -1515,6 +1520,7 @@ def query_doc_handler( ), k=form_data.k if form_data.k else request.app.state.config.TOP_K, reranking_function=request.app.state.rf, + k_reranker=form_data.k_reranker or request.app.state.config.TOP_K_RERANKER, r=( form_data.r if form_data.r @@ -1543,6 +1549,7 @@ class QueryCollectionsForm(BaseModel): collection_names: list[str] query: str k: Optional[int] = None + k_reranker: Optional[int] = None r: Optional[float] = None hybrid: Optional[bool] = None @@ -1563,6 +1570,7 @@ def query_collection_handler( ), k=form_data.k if form_data.k else request.app.state.config.TOP_K, reranking_function=request.app.state.rf, + k_reranker=form_data.k_reranker or request.app.state.config.TOP_K_RERANKER, r=( form_data.r if form_data.r diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index f2fbf4143..2ba892b9f 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -584,6 +584,7 @@ async def chat_completion_files_handler( ), k=request.app.state.config.TOP_K, reranking_function=request.app.state.rf, + k_reranker=request.app.state.config.TOP_K_RERANKER, r=request.app.state.config.RELEVANCE_THRESHOLD, hybrid_search=request.app.state.config.ENABLE_RAG_HYBRID_SEARCH, full_context=request.app.state.config.RAG_FULL_CONTEXT, diff --git a/src/lib/components/admin/Settings/Documents.svelte b/src/lib/components/admin/Settings/Documents.svelte index 21b8515f1..3edfa1b1d 100644 --- a/src/lib/components/admin/Settings/Documents.svelte +++ b/src/lib/components/admin/Settings/Documents.svelte @@ -76,6 +76,7 @@ template: '', r: 0.0, k: 4, + k_reranker: 4, hybrid: false }; @@ -765,6 +766,23 @@ + {#if querySettings.hybrid === true} +
+
{$i18n.t('Top K Reranker')}
+
+ +
+
+ {/if} + + {#if querySettings.hybrid === true}
diff --git a/src/lib/i18n/locales/de-DE/translation.json b/src/lib/i18n/locales/de-DE/translation.json index 5fdd4e147..172bc6b8f 100644 --- a/src/lib/i18n/locales/de-DE/translation.json +++ b/src/lib/i18n/locales/de-DE/translation.json @@ -437,6 +437,7 @@ "Enter timeout in seconds": "Geben Sie den Timeout in Sekunden ein", "Enter to Send": "'Enter' zum Senden", "Enter Top K": "Geben Sie Top K ein", + "Enter Top K Reranker": "Geben Sie Top K für Reranker ein", "Enter URL (e.g. http://127.0.0.1:7860/)": "Geben Sie die URL ein (z. B. http://127.0.0.1:7860/)", "Enter URL (e.g. http://localhost:11434)": "Geben Sie die URL ein (z. B. http://localhost:11434)", "Enter your current password": "Geben Sie Ihr aktuelles Passwort ein", diff --git a/src/lib/i18n/locales/en-GB/translation.json b/src/lib/i18n/locales/en-GB/translation.json index 19581081f..fdde49907 100644 --- a/src/lib/i18n/locales/en-GB/translation.json +++ b/src/lib/i18n/locales/en-GB/translation.json @@ -437,6 +437,7 @@ "Enter timeout in seconds": "", "Enter to Send": "", "Enter Top K": "", + "Enter Top K Reranker": "", "Enter URL (e.g. http://127.0.0.1:7860/)": "", "Enter URL (e.g. http://localhost:11434)": "", "Enter your current password": "", diff --git a/src/lib/i18n/locales/en-US/translation.json b/src/lib/i18n/locales/en-US/translation.json index 19581081f..fdde49907 100644 --- a/src/lib/i18n/locales/en-US/translation.json +++ b/src/lib/i18n/locales/en-US/translation.json @@ -437,6 +437,7 @@ "Enter timeout in seconds": "", "Enter to Send": "", "Enter Top K": "", + "Enter Top K Reranker": "", "Enter URL (e.g. http://127.0.0.1:7860/)": "", "Enter URL (e.g. http://localhost:11434)": "", "Enter your current password": "",