Sfoglia il codice sorgente

Merge pull request #5829 from jannikstdl/query-embedding-perf-fix

fix:  performance issues on large collections
Timothy Jaeryang Baek 7 mesi fa
parent
commit
0876c9b5ef
1 ha cambiato i file con 6 aggiunte e 5 eliminazioni
  1. 6 5
      backend/open_webui/apps/retrieval/utils.py

+ 6 - 5
backend/open_webui/apps/retrieval/utils.py

@@ -65,14 +65,13 @@ class VectorSearchRetriever(BaseRetriever):
 
 
 def query_doc(
 def query_doc(
     collection_name: str,
     collection_name: str,
-    query: str,
-    embedding_function,
+    query_embedding: list[float],
     k: int,
     k: int,
 ):
 ):
     try:
     try:
         result = VECTOR_DB_CLIENT.search(
         result = VECTOR_DB_CLIENT.search(
             collection_name=collection_name,
             collection_name=collection_name,
-            vectors=[embedding_function(query)],
+            vectors=[query_embedding],
             limit=k,
             limit=k,
         )
         )
 
 
@@ -182,15 +181,17 @@ def query_collection(
     embedding_function,
     embedding_function,
     k: int,
     k: int,
 ) -> dict:
 ) -> dict:
+    
     results = []
     results = []
+    query_embedding = embedding_function(query)
+
     for collection_name in collection_names:
     for collection_name in collection_names:
         if collection_name:
         if collection_name:
             try:
             try:
                 result = query_doc(
                 result = query_doc(
                     collection_name=collection_name,
                     collection_name=collection_name,
-                    query=query,
                     k=k,
                     k=k,
-                    embedding_function=embedding_function,
+                    query_embedding=query_embedding,
                 )
                 )
                 results.append(result.model_dump())
                 results.append(result.model_dump())
             except Exception as e:
             except Exception as e: