Merge pull request #8212 from ashm-dev/main

feat: Small optimization

backend/open_webui/retrieval/utils.py

@@ -1,9 +1,8 @@
 import logging
 import os
-import uuid
+import heapq
 from typing import Optional, Union
 
-import asyncio
 import requests
 
 from huggingface_hub import snapshot_download
@@ -34,8 +33,6 @@ class VectorSearchRetriever(BaseRetriever):
     def _get_relevant_documents(
         self,
         query: str,
-        *,
-        run_manager: CallbackManagerForRetrieverRun,
     ) -> list[Document]:
         result = VECTOR_DB_CLIENT.search(
             collection_name=self.collection_name,
@@ -47,15 +44,12 @@ class VectorSearchRetriever(BaseRetriever):
         metadatas = result.metadatas[0]
         documents = result.documents[0]
 
-        results = []
+        return [
-        for idx in range(len(ids)):
-            results.append(
             Document(
                 metadata=metadatas[idx],
                 page_content=documents[idx],
-                )
+            ) for idx in range(len(ids))
-            )
+        ]
-        return results
 
 
 def query_doc(
@@ -64,13 +58,11 @@ def query_doc(
     k: int,
 ):
     try:
-        result = VECTOR_DB_CLIENT.search(
+        if result := VECTOR_DB_CLIENT.search(
             collection_name=collection_name,
             vectors=[query_embedding],
             limit=k,
-        )
+        ):
-
-        if result:
             log.info(f"query_doc:result {result.ids} {result.metadatas}")
 
             return result
@@ -135,45 +127,39 @@ def query_doc_with_hybrid_search(
 def merge_and_sort_query_results(
     query_results: list[dict], k: int, reverse: bool = False
 ) -> list[dict]:
-    # Initialize lists to store combined data
+    if not query_results:
-    combined_distances = []
+        return {
-    combined_documents = []
+            "distances": [[]],
-    combined_metadatas = []
+            "documents": [[]],
+            "metadatas": [[]],
+        }
         
-    for data in query_results:
+    combined = (
-        combined_distances.extend(data["distances"][0])
+        (data.get("distances", [float('inf')])[0],
-        combined_documents.extend(data["documents"][0])
+         data.get("documents", [None])[0],
-        combined_metadatas.extend(data["metadatas"][0])
+         data.get("metadatas", [{}])[0])
+        for data in query_results
+    )
     
-    # Create a list of tuples (distance, document, metadata)
+    if reverse:
-    combined = list(zip(combined_distances, combined_documents, combined_metadatas))
+        top_k = heapq.nlargest(k, combined, key=lambda x: x[0])
-
-    # Sort the list based on distances
-    combined.sort(key=lambda x: x[0], reverse=reverse)
-
-    # We don't have anything :-(
-    if not combined:
-        sorted_distances = []
-        sorted_documents = []
-        sorted_metadatas = []
     else:
-        # Unzip the sorted list
+        top_k = heapq.nsmallest(k, combined, key=lambda x: x[0])
-        sorted_distances, sorted_documents, sorted_metadatas = zip(*combined)
     
-        # Slicing the lists to include only k elements
+    if not top_k:
-        sorted_distances = list(sorted_distances)[:k]
+        return {
-        sorted_documents = list(sorted_documents)[:k]
+            "distances": [[]],
-        sorted_metadatas = list(sorted_metadatas)[:k]
+            "documents": [[]],
-
+            "metadatas": [[]],
-    # Create the output dictionary
+        }
-    result = {
+    else:
+        sorted_distances, sorted_documents, sorted_metadatas = zip(*top_k)
+        return {
             "distances": [sorted_distances],
             "documents": [sorted_documents],
             "metadatas": [sorted_metadatas],
         }
 
-    return result
-
 
 def query_collection(
     collection_names: list[str],
@@ -185,19 +171,18 @@ def query_collection(
     for query in queries:
         query_embedding = embedding_function(query)
         for collection_name in collection_names:
-            if collection_name:
+            if not collection_name:
+                continue
+            
             try:
-                    result = query_doc(
+                if result := query_doc(
                     collection_name=collection_name,
                     k=k,
                     query_embedding=query_embedding,
-                    )
+                ):
-                    if result is not None:
                     results.append(result.model_dump())
             except Exception as e:
                 log.exception(f"Error when querying the collection: {e}")
-            else:
-                pass
 
     return merge_and_sort_query_results(results, k=k)
 
@@ -213,8 +198,8 @@ def query_collection_with_hybrid_search(
     results = []
     error = False
     for collection_name in collection_names:
-        try:
         for query in queries:
+            try:
                 result = query_doc_with_hybrid_search(
                     collection_name=collection_name,
                     query=query,
@@ -259,10 +244,10 @@ def get_embedding_function(
 
         def generate_multiple(query, func):
             if isinstance(query, list):
-                embeddings = []
+                return [
-                for i in range(0, len(query), embedding_batch_size):
+                    func(query[i : i + embedding_batch_size])
-                    embeddings.extend(func(query[i : i + embedding_batch_size]))
+                    for i in range(0, len(query), embedding_batch_size)
-                return embeddings
+                ]
             else:
                 return func(query)
 
@@ -436,7 +421,6 @@ def generate_openai_batch_embeddings(
 def generate_ollama_batch_embeddings(
     model: str, texts: list[str], url: str, key: str = ""
 ) -> Optional[list[list[float]]]:
-    try:
     r = requests.post(
         f"{url}/api/embed",
         headers={
@@ -445,17 +429,19 @@ def generate_ollama_batch_embeddings(
         },
         json={"input": texts, "model": model},
     )
+    try:
         r.raise_for_status()
-        data = r.json()
-
-        if "embeddings" in data:
-            return data["embeddings"]
-        else:
-            raise "Something went wrong :/"
     except Exception as e:
         print(e)
         return None
     
+    data = r.json()
+    
+    if 'embeddings' not in data:
+        raise "Something went wrong :/"
+    
+    return data['embeddings']
+
 
 def generate_embeddings(engine: str, model: str, text: Union[str, list[str]], **kwargs):
     url = kwargs.get("url", "")