AI/open-webui
1
0
Fork 1
mirror of https://git.mirrors.martin98.com/https://github.com/open-webui/open-webui synced 2025-08-19 06:35:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

Use k_reranker also for result merge, and add special sorting use case for ChromaDB

Browse Source
This commit is contained in:
Marko Henning 2025-03-18 16:25:24 +01:00
parent f13948d805
commit ba676b7ed6
1 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
backend/open_webui/retrieval/utils.py
View File

@ -146,6 +146,9 @@ def query_doc_with_hybrid_search(
# retrieve only min(k, k_reranker) items, sort and cut by distance if k < k_reranker # retrieve only min(k, k_reranker) items, sort and cut by distance if k < k_reranker
if k < k_reranker: if k < k_reranker:
if VECTOR_DB == "chroma":
sorted_items = sorted(zip(distances, metadatas, documents), key=lambda x: x[0], reverse=False)
else:
sorted_items = sorted(zip(distances, metadatas, documents), key=lambda x: x[0], reverse=True) sorted_items = sorted(zip(distances, metadatas, documents), key=lambda x: x[0], reverse=True)
sorted_items = sorted_items[:k] sorted_items = sorted_items[:k]
distances, documents, metadatas = map(list, zip(*sorted_items)) distances, documents, metadatas = map(list, zip(*sorted_items))
@ -310,9 +313,9 @@ def query_collection_with_hybrid_search(
if VECTOR_DB == "chroma": if VECTOR_DB == "chroma":
# Chroma uses unconventional cosine similarity, so we don't need to reverse the results # Chroma uses unconventional cosine similarity, so we don't need to reverse the results
# https://docs.trychroma.com/docs/collections/configure#configuring-chroma-collections # https://docs.trychroma.com/docs/collections/configure#configuring-chroma-collections
return merge_and_sort_query_results(results, k=k, reverse=False) return merge_and_sort_query_results(results, k=k_reranker, reverse=False)
else: else:
return merge_and_sort_query_results(results, k=k, reverse=True) return merge_and_sort_query_results(results, k=k_reranker, reverse=True)
def get_embedding_function( def get_embedding_function(