From 6bff5a4d09db086a0b2e200c73c3f0b52e168ebe Mon Sep 17 00:00:00 2001 From: Jannik Streidl Date: Mon, 30 Sep 2024 16:18:02 +0200 Subject: [PATCH] fix: do not embed the query every single time --- backend/open_webui/apps/retrieval/utils.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/backend/open_webui/apps/retrieval/utils.py b/backend/open_webui/apps/retrieval/utils.py index 12c30edbb..32a71fa84 100644 --- a/backend/open_webui/apps/retrieval/utils.py +++ b/backend/open_webui/apps/retrieval/utils.py @@ -65,14 +65,13 @@ class VectorSearchRetriever(BaseRetriever): def query_doc( collection_name: str, - query: str, - embedding_function, + query_embedding: list[float], k: int, ): try: result = VECTOR_DB_CLIENT.search( collection_name=collection_name, - vectors=[embedding_function(query)], + vectors=[query_embedding], limit=k, ) @@ -182,15 +181,17 @@ def query_collection( embedding_function, k: int, ) -> dict: + results = [] + query_embedding = embedding_function(query) + for collection_name in collection_names: if collection_name: try: result = query_doc( collection_name=collection_name, - query=query, k=k, - embedding_function=embedding_function, + query_embedding=query_embedding, ) results.append(result.model_dump()) except Exception as e: