refac

backend/open_webui/retrieval/utils.py

@@ -44,7 +44,7 @@ from open_webui.models.users import UserModel
 from open_webui.retrieval.loaders.youtube import YoutubeLoader
 from open_webui.retrieval.vector.async_client import ASYNC_VECTOR_DB_CLIENT
 from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT
-from open_webui.retrieval.vector.main import GetResult
+from open_webui.retrieval.vector.main import GetResult, SearchResult
 from open_webui.retrieval.web.utils import get_web_loader
 from open_webui.utils.access_control.files import has_access_to_file
 from open_webui.utils.headers import include_user_info_headers
@@ -360,9 +360,96 @@ def get_enriched_texts(collection_result: GetResult) -> list[str]:
     return enriched_texts
 
 
+def _search_result_to_documents(result: SearchResult) -> list[Document]:
+    ids = result.ids[0] if result and result.ids else []
+    metadatas = result.metadatas[0] if result and result.metadatas else []
+    documents = result.documents[0] if result and result.documents else []
+    distances = result.distances[0] if result and result.distances else []
+
+    docs = []
+    for idx in range(len(ids)):
+        document = documents[idx]
+        metadata = dict(metadatas[idx] or {})
+        metadata[CHUNK_HASH_KEY] = _content_hash(document)
+        if idx < len(distances):
+            metadata.setdefault('score', distances[idx])
+        docs.append(Document(metadata=metadata, page_content=document))
+    return docs
+
+
+def _supports_native_hybrid_search() -> bool:
+    supports_hybrid_search = getattr(ASYNC_VECTOR_DB_CLIENT, 'supports_hybrid_search', None)
+    if supports_hybrid_search is not None:
+        return bool(supports_hybrid_search)
+    return callable(getattr(ASYNC_VECTOR_DB_CLIENT, 'hybrid_search', None))
+
+
+async def query_doc_with_native_hybrid_search(
+    collection_name: str,
+    query: str,
+    embedding_function,
+    k: int,
+    reranking_function,
+    k_reranker: int,
+    r: float,
+    hybrid_bm25_weight: float,
+) -> Optional[dict]:
+    try:
+        if not _supports_native_hybrid_search():
+            return None
+
+        query_vectors = []
+        if hybrid_bm25_weight < 1:
+            query_vectors = [await embedding_function(query, RAG_EMBEDDING_QUERY_PREFIX)]
+
+        result = await ASYNC_VECTOR_DB_CLIENT.hybrid_search(
+            collection_name=collection_name,
+            query=query,
+            vectors=query_vectors,
+            limit=k,
+            hybrid_bm25_weight=hybrid_bm25_weight,
+        )
+        if result is None:
+            return None
+
+        documents = _search_result_to_documents(result)
+        if not documents:
+            return {'distances': [[]], 'documents': [[]], 'metadatas': [[]]}
+
+        compressor = RerankCompressor(
+            embedding_function=embedding_function,
+            top_n=k_reranker,
+            reranking_function=reranking_function,
+            r_score=r,
+        )
+        compressed = await compressor.acompress_documents(documents, query)
+
+        distances = [d.metadata.get('score') for d in compressed]
+        documents = [d.page_content for d in compressed]
+        metadatas = [d.metadata for d in compressed]
+
+        if k < k_reranker:
+            sorted_items = sorted(zip(distances, documents, metadatas), key=lambda x: x[0], reverse=True)
+            sorted_items = sorted_items[:k]
+
+            if sorted_items:
+                distances, documents, metadatas = map(list, zip(*sorted_items))
+            else:
+                distances, documents, metadatas = [], [], []
+
+        return {
+            'distances': [distances],
+            'documents': [documents],
+            'metadatas': [metadatas],
+        }
+    except Exception as e:
+        log.debug(f'Native hybrid search failed for {collection_name}, falling back to legacy hybrid search: {e}')
+        return None
+
+
 async def query_doc_with_hybrid_search(
     collection_name: str,
-    collection_result: GetResult,
+    collection_result: Optional[GetResult],
     query: str,
     embedding_function,
     k: int,
@@ -371,8 +458,26 @@ async def query_doc_with_hybrid_search(
     r: float,
     hybrid_bm25_weight: float,
     enable_enriched_texts: bool = False,
+    native_hybrid_search: bool = True,
 ) -> dict:
     try:
+        if native_hybrid_search and not enable_enriched_texts:
+            native_result = await query_doc_with_native_hybrid_search(
+                collection_name=collection_name,
+                query=query,
+                embedding_function=embedding_function,
+                k=k,
+                reranking_function=reranking_function,
+                k_reranker=k_reranker,
+                r=r,
+                hybrid_bm25_weight=hybrid_bm25_weight,
+            )
+            if native_result is not None:
+                return native_result
+
+        if collection_result is None:
+            collection_result = await ASYNC_VECTOR_DB_CLIENT.get(collection_name=collection_name)
+
         # First check if collection_result has the required attributes
         if (
             not collection_result
@@ -652,6 +757,32 @@ async def query_collection_with_hybrid_search(
 ) -> dict:
     results = []
     error = False
+
+    if not enable_enriched_texts:
+
+        async def process_native_query(collection_name, query):
+            result = await query_doc_with_native_hybrid_search(
+                collection_name=collection_name,
+                query=query,
+                embedding_function=embedding_function,
+                k=k,
+                reranking_function=reranking_function,
+                k_reranker=k_reranker,
+                r=r,
+                hybrid_bm25_weight=hybrid_bm25_weight,
+            )
+            return result
+
+        native_task_results = await asyncio.gather(
+            *[
+                process_native_query(collection_name, query)
+                for collection_name in collection_names
+                for query in queries
+            ]
+        )
+        if native_task_results and all(result is not None for result in native_task_results):
+            return merge_and_sort_query_results(native_task_results, k=k)
+
     # Fetch every collection's contents once up front so the
     # per-query/per-document loop below can reuse them. Each fetch
     # offloads to a worker thread, so run them concurrently with
@@ -686,6 +817,7 @@ async def query_collection_with_hybrid_search(
                 r=r,
                 hybrid_bm25_weight=hybrid_bm25_weight,
                 enable_enriched_texts=enable_enriched_texts,
+                native_hybrid_search=False,
             )
             return result, None
         except Exception as e: