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:

backend/open_webui/retrieval/vector/async_client.py

@@ -82,6 +82,10 @@ class AsyncVectorDBClient:
         (e.g. already inside a worker thread)."""
         return self._sync
 
+    @property
+    def supports_hybrid_search(self) -> bool:
+        return type(self._sync).hybrid_search is not VectorDBBase.hybrid_search
+
     async def has_collection(self, collection_name: str) -> bool:
         return await asyncio.to_thread(self._sync.has_collection, collection_name)
 
@@ -103,6 +107,25 @@ class AsyncVectorDBClient:
     ) -> Optional[SearchResult]:
         return await asyncio.to_thread(self._sync.search, collection_name, vectors, filter, limit)
 
+    async def hybrid_search(
+        self,
+        collection_name: str,
+        query: str,
+        vectors: List[List[Union[float, int]]],
+        filter: Optional[Dict] = None,
+        limit: int = 10,
+        hybrid_bm25_weight: float = 0.5,
+    ) -> Optional[SearchResult]:
+        return await asyncio.to_thread(
+            self._sync.hybrid_search,
+            collection_name,
+            query,
+            vectors,
+            filter,
+            limit,
+            hybrid_bm25_weight,
+        )
+
     async def query(
         self,
         collection_name: str,

backend/open_webui/retrieval/vector/dbs/pgvector.py

@@ -24,7 +24,7 @@ from open_webui.retrieval.vector.main import (
     VectorDBBase,
     VectorItem,
 )
-from open_webui.retrieval.vector.utils import process_metadata
+from open_webui.retrieval.vector.utils import merge_hybrid_search_results, process_metadata
 from open_webui.utils.misc import sanitize_text_for_db
 from pgvector.sqlalchemy import HALFVEC, Vector
 from sqlalchemy import (
@@ -153,6 +153,7 @@ class PgvectorClient(VectorDBBase):
 
             index_method, index_options = self._vector_index_configuration()
             self._ensure_vector_index(index_method, index_options)
+            self._ensure_text_search_index()
 
             self.session.execute(
                 text(
@@ -236,6 +237,19 @@ class PgvectorClient(VectorDBBase):
                 f' {index_options}' if index_options else '',
             )
 
+    def _ensure_text_search_index(self) -> None:
+        if PGVECTOR_PGCRYPTO:
+            return
+
+        self.session.execute(
+            text("""
+                CREATE INDEX IF NOT EXISTS idx_document_chunk_text_search
+                ON document_chunk
+                USING GIN (to_tsvector('simple', coalesce(text, '')));
+                """)
+        )
+        log.info("Ensured text search index 'idx_document_chunk_text_search'.")
+
     def check_vector_length(self) -> None:
         """
         Check if the VECTOR_LENGTH matches the existing vector column dimension in the database.
@@ -521,6 +535,71 @@ class PgvectorClient(VectorDBBase):
             log.exception(f'Error during search: {e}')
             return None
 
+    def hybrid_search(
+        self,
+        collection_name: str,
+        query: str,
+        vectors: List[List[float]],
+        filter: Optional[Dict[str, Any]] = None,
+        limit: int = 10,
+        hybrid_bm25_weight: float = 0.5,
+    ) -> Optional[SearchResult]:
+        if PGVECTOR_PGCRYPTO or filter:
+            return None
+
+        try:
+            limit = max(1, limit)
+            vectors = [self.adjust_vector_length(vector) for vector in vectors] if vectors else []
+            num_queries = len(vectors) if vectors else 1
+            bm25_weight = min(max(hybrid_bm25_weight, 0.0), 1.0)
+            vector_weight = 1.0 - bm25_weight
+
+            vector_result = None
+            if vector_weight > 0 and vectors:
+                vector_result = self.search(collection_name=collection_name, vectors=vectors, limit=limit)
+
+            fts_results = []
+            if bm25_weight > 0 and query and query.strip():
+                fts_rows = self.session.execute(
+                    text("""
+                        WITH fts_query AS (
+                            SELECT plainto_tsquery('simple', :query) AS query
+                        )
+                        SELECT
+                            document_chunk.id AS id,
+                            document_chunk.text AS text,
+                            document_chunk.vmetadata AS vmetadata,
+                            ts_rank_cd(
+                                to_tsvector('simple', coalesce(document_chunk.text, '')),
+                                fts_query.query
+                            ) AS rank
+                        FROM document_chunk, fts_query
+                        WHERE document_chunk.collection_name = :collection_name
+                          AND to_tsvector('simple', coalesce(document_chunk.text, '')) @@ fts_query.query
+                        ORDER BY rank DESC
+                        LIMIT :limit
+                    """),
+                    {
+                        'collection_name': collection_name,
+                        'query': query,
+                        'limit': limit,
+                    },
+                )
+                fts_results = [dict(row) for row in fts_rows.mappings().all()]
+                self.session.rollback()
+
+            return merge_hybrid_search_results(
+                vector_result=vector_result,
+                fts_results=fts_results,
+                num_queries=num_queries,
+                limit=limit,
+                hybrid_bm25_weight=hybrid_bm25_weight,
+            )
+        except Exception as e:
+            self.session.rollback()
+            log.exception(f'Error during hybrid search: {e}')
+            return None
+
     def query(self, collection_name: str, filter: Dict[str, Any], limit: Optional[int] = None) -> Optional[GetResult]:
         try:
             if PGVECTOR_PGCRYPTO:

backend/open_webui/retrieval/vector/main.py

@@ -63,6 +63,18 @@ class VectorDBBase(ABC):
         """Search for similar vectors in a collection."""
         pass
 
+    def hybrid_search(
+        self,
+        collection_name: str,
+        query: str,
+        vectors: List[List[Union[float, int]]],
+        filter: Optional[Dict] = None,
+        limit: int = 10,
+        hybrid_bm25_weight: float = 0.5,
+    ) -> Optional[SearchResult]:
+        """Search using a backend-native hybrid keyword/vector implementation when available."""
+        return None
+
     @abstractmethod
     def query(self, collection_name: str, filter: Dict, limit: Optional[int] = None) -> Optional[GetResult]:
         """Query vectors from a collection using metadata filter."""

backend/open_webui/retrieval/vector/utils.py

@@ -1,5 +1,7 @@
 from datetime import datetime
+from typing import Any, Optional
 
+from open_webui.retrieval.vector.main import SearchResult
 from open_webui.utils.misc import sanitize_text_for_db
 
 KEYS_TO_EXCLUDE = ['content', 'pages', 'tables', 'paragraphs', 'sections', 'figures']
@@ -27,3 +29,63 @@ def process_metadata(
         else:
             result[key] = sanitize_text_for_db(value)
     return result
+
+
+def merge_hybrid_search_results(
+    vector_result: Optional[SearchResult],
+    fts_results: list[dict[str, Any]],
+    num_queries: int,
+    limit: int,
+    hybrid_bm25_weight: float,
+) -> SearchResult:
+    rank_constant = 60.0
+    bm25_weight = min(max(hybrid_bm25_weight, 0.0), 1.0)
+    vector_weight = 1.0 - bm25_weight
+
+    ids = [[] for _ in range(num_queries)]
+    distances = [[] for _ in range(num_queries)]
+    documents = [[] for _ in range(num_queries)]
+    metadatas = [[] for _ in range(num_queries)]
+
+    for qid in range(num_queries):
+        candidates: dict[str, dict[str, Any]] = {}
+
+        if vector_result and vector_result.ids and qid < len(vector_result.ids):
+            for rank, item_id in enumerate(vector_result.ids[qid] or [], start=1):
+                score = vector_weight / (rank_constant + rank) if vector_weight > 0 else 0
+                if score <= 0:
+                    continue
+
+                candidate = candidates.setdefault(
+                    item_id,
+                    {
+                        'score': 0.0,
+                        'document': vector_result.documents[qid][rank - 1],
+                        'metadata': vector_result.metadatas[qid][rank - 1],
+                    },
+                )
+                candidate['score'] += score
+
+        for rank, row in enumerate(fts_results, start=1):
+            score = bm25_weight / (rank_constant + rank) if bm25_weight > 0 else 0
+            if score <= 0:
+                continue
+
+            item_id = row['id']
+            candidate = candidates.setdefault(
+                item_id,
+                {
+                    'score': 0.0,
+                    'document': row['text'],
+                    'metadata': row['vmetadata'],
+                },
+            )
+            candidate['score'] += score
+
+        ranked = sorted(candidates.items(), key=lambda item: item[1]['score'], reverse=True)[:limit]
+        ids[qid] = [item_id for item_id, _ in ranked]
+        distances[qid] = [candidate['score'] for _, candidate in ranked]
+        documents[qid] = [candidate['document'] for _, candidate in ranked]
+        metadatas[qid] = [candidate['metadata'] for _, candidate in ranked]
+
+    return SearchResult(ids=ids, distances=distances, documents=documents, metadatas=metadatas)

backend/open_webui/routers/retrieval.py

@@ -2580,6 +2580,7 @@ class QueryDocForm(BaseModel):
     k_reranker: int | None = None
     r: float | None = None
     hybrid: bool | None = None
+    hybrid_bm25_weight: float | None = None
 
 
 @router.post('/query/doc')
@@ -2593,13 +2594,9 @@ async def query_doc_handler(
 
     try:
         if config.ENABLE_RAG_HYBRID_SEARCH and (form_data.hybrid is None or form_data.hybrid):
-            collection_results = {}
-            collection_results[form_data.collection_name] = await ASYNC_VECTOR_DB_CLIENT.get(
-                collection_name=form_data.collection_name
-            )
             return await query_doc_with_hybrid_search(
                 collection_name=form_data.collection_name,
-                collection_result=collection_results[form_data.collection_name],
+                collection_result=None,
                 query=form_data.query,
                 embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
                     query, prefix=prefix, user=user
@@ -2614,10 +2611,9 @@ async def query_doc_handler(
                 r=(form_data.r if form_data.r else config.RELEVANCE_THRESHOLD),
                 hybrid_bm25_weight=(
                     form_data.hybrid_bm25_weight
-                    if form_data.hybrid_bm25_weight
+                    if form_data.hybrid_bm25_weight is not None
                     else config.HYBRID_BM25_WEIGHT
                 ),
-                user=user,
             )
         else:
             query_embedding = await request.app.state.EMBEDDING_FUNCTION(
@@ -2678,7 +2674,7 @@ async def query_collection_handler(
                 r=(form_data.r if form_data.r else config.RELEVANCE_THRESHOLD),
                 hybrid_bm25_weight=(
                     form_data.hybrid_bm25_weight
-                    if form_data.hybrid_bm25_weight
+                    if form_data.hybrid_bm25_weight is not None
                     else config.HYBRID_BM25_WEIGHT
                 ),
                 enable_enriched_texts=(