From 223f484ded01d092979693341dc03351a9fa17fa Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Mon, 22 Jun 2026 16:10:19 +0200 Subject: [PATCH] refac --- backend/open_webui/retrieval/utils.py | 136 +++++++++++++++++- .../retrieval/vector/async_client.py | 23 +++ .../retrieval/vector/dbs/pgvector.py | 81 ++++++++++- backend/open_webui/retrieval/vector/main.py | 12 ++ backend/open_webui/retrieval/vector/utils.py | 62 ++++++++ backend/open_webui/routers/retrieval.py | 12 +- 6 files changed, 315 insertions(+), 11 deletions(-) diff --git a/backend/open_webui/retrieval/utils.py b/backend/open_webui/retrieval/utils.py index 16f6dad18a..c8c603f3d8 100644 --- a/backend/open_webui/retrieval/utils.py +++ b/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: diff --git a/backend/open_webui/retrieval/vector/async_client.py b/backend/open_webui/retrieval/vector/async_client.py index 0bea6696a9..481f26e19f 100644 --- a/backend/open_webui/retrieval/vector/async_client.py +++ b/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, diff --git a/backend/open_webui/retrieval/vector/dbs/pgvector.py b/backend/open_webui/retrieval/vector/dbs/pgvector.py index 861d49bc1b..b37d774f72 100644 --- a/backend/open_webui/retrieval/vector/dbs/pgvector.py +++ b/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: diff --git a/backend/open_webui/retrieval/vector/main.py b/backend/open_webui/retrieval/vector/main.py index 38ea699514..dd284eae1a 100644 --- a/backend/open_webui/retrieval/vector/main.py +++ b/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.""" diff --git a/backend/open_webui/retrieval/vector/utils.py b/backend/open_webui/retrieval/vector/utils.py index 4915b024c3..44af9c2b61 100644 --- a/backend/open_webui/retrieval/vector/utils.py +++ b/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) diff --git a/backend/open_webui/routers/retrieval.py b/backend/open_webui/routers/retrieval.py index 3269ad0e3a..c49af3af1b 100644 --- a/backend/open_webui/routers/retrieval.py +++ b/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=(