refac

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)