revert

backend/open_webui/retrieval/utils.py

@@ -1,8 +1,9 @@
 import logging
 import os
-import heapq
+import uuid
 from typing import Optional, Union
 
+import asyncio
 import requests
 
 from huggingface_hub import snapshot_download
@@ -33,6 +34,8 @@ 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,
@@ -44,12 +47,15 @@ class VectorSearchRetriever(BaseRetriever):
         metadatas = result.metadatas[0]
         documents = result.documents[0]
 
-        return [
+        results = []
+        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(
@@ -58,11 +64,13 @@ def query_doc(
     k: int,
 ):
     try:
-        if result := VECTOR_DB_CLIENT.search(
+        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
@@ -127,39 +135,45 @@ def query_doc_with_hybrid_search(
 def merge_and_sort_query_results(
     query_results: list[dict], k: int, reverse: bool = False
 ) -> list[dict]:
-    if not query_results:
+    # Initialize lists to store combined data
-        return {
+    combined_distances = []
-            "distances": [[]],
+    combined_documents = []
-            "documents": [[]],
+    combined_metadatas = []
-            "metadatas": [[]],
-        }
 
-    combined = (
+    for data in query_results:
-        (data.get("distances", [float('inf')])[0],
+        combined_distances.extend(data["distances"][0])
-         data.get("documents", [None])[0],
+        combined_documents.extend(data["documents"][0])
-         data.get("metadatas", [{}])[0])
+        combined_metadatas.extend(data["metadatas"][0])
-        for data in query_results
-    )
 
-    if reverse:
+    # Create a list of tuples (distance, document, metadata)
-        top_k = heapq.nlargest(k, combined, key=lambda x: x[0])
+    combined = list(zip(combined_distances, combined_documents, combined_metadatas))
+
+    # 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:
-        top_k = heapq.nsmallest(k, combined, key=lambda x: x[0])
+        # Unzip the sorted list
+        sorted_distances, sorted_documents, sorted_metadatas = zip(*combined)
 
-    if not top_k:
+        # Slicing the lists to include only k elements
-        return {
+        sorted_distances = list(sorted_distances)[:k]
-            "distances": [[]],
+        sorted_documents = list(sorted_documents)[:k]
-            "documents": [[]],
+        sorted_metadatas = list(sorted_metadatas)[:k]
-            "metadatas": [[]],
+
-        }
+    # Create the output dictionary
-    else:
+    result = {
-        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],
@@ -171,18 +185,19 @@ def query_collection(
     for query in queries:
         query_embedding = embedding_function(query)
         for collection_name in collection_names:
-            if not collection_name:
+            if collection_name:
-                continue
-            
                 try:
-                if result := query_doc(
+                    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)
 
@@ -198,8 +213,8 @@ def query_collection_with_hybrid_search(
     results = []
     error = False
     for collection_name in collection_names:
-        for query in queries:
         try:
+            for query in queries:
                 result = query_doc_with_hybrid_search(
                     collection_name=collection_name,
                     query=query,
@@ -244,10 +259,10 @@ def get_embedding_function(
 
         def generate_multiple(query, func):
             if isinstance(query, list):
-                return [
+                embeddings = []
-                    func(query[i : i + embedding_batch_size])
+                for i in range(0, len(query), embedding_batch_size):
-                    for i in range(0, len(query), embedding_batch_size)
+                    embeddings.extend(func(query[i : i + embedding_batch_size]))
-                ]
+                return embeddings
             else:
                 return func(query)
 
@@ -421,6 +436,7 @@ 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={
@@ -429,19 +445,17 @@ 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", "")