diff --git a/backend/open_webui/retrieval/utils.py b/backend/open_webui/retrieval/utils.py index f7d2775c52..b4bd615ccf 100644 --- a/backend/open_webui/retrieval/utils.py +++ b/backend/open_webui/retrieval/utils.py @@ -20,6 +20,7 @@ from langchain_community.retrievers import BM25Retriever from langchain_core.documents import Document from open_webui.config import VECTOR_DB +from open_webui.retrieval.vector.async_client import ASYNC_VECTOR_DB_CLIENT from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT @@ -121,7 +122,7 @@ class VectorSearchRetriever(BaseRetriever): run_manager: CallbackManagerForRetrieverRun, ) -> list[Document]: embedding = await self.embedding_function(query, RAG_EMBEDDING_QUERY_PREFIX) - result = VECTOR_DB_CLIENT.search( + result = await ASYNC_VECTOR_DB_CLIENT.search( collection_name=self.collection_name, vectors=[embedding], limit=self.top_k, @@ -488,16 +489,26 @@ async def query_collection_with_hybrid_search( ) -> dict: results = [] error = False - # Fetch collection data once per collection sequentially - # Avoid fetching the same data multiple times later - collection_results = {} - for collection_name in collection_names: + # 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 + # `asyncio.gather` instead of awaiting them serially — otherwise + # latency scales linearly with `len(collection_names)`. + log.debug( + 'query_collection_with_hybrid_search: prefetching %d collections', + len(collection_names), + ) + + async def _fetch_collection(name: str): try: - log.debug(f'query_collection_with_hybrid_search:VECTOR_DB_CLIENT.get:collection {collection_name}') - collection_results[collection_name] = VECTOR_DB_CLIENT.get(collection_name=collection_name) + return name, await ASYNC_VECTOR_DB_CLIENT.get(collection_name=name) except Exception as e: - log.exception(f'Failed to fetch collection {collection_name}: {e}') - collection_results[collection_name] = None + log.exception(f'Failed to fetch collection {name}: {e}') + return name, None + + collection_results = dict( + await asyncio.gather(*(_fetch_collection(name) for name in collection_names)) + ) log.info(f'Starting hybrid search for {len(queries)} queries in {len(collection_names)} collections...') @@ -1140,7 +1151,11 @@ async def get_sources_from_items( try: if full_context: - query_result = get_all_items_from_collections(collection_names) + # Sync helper makes blocking VECTOR_DB_CLIENT calls; + # offload so the async caller's event loop stays free. + query_result = await asyncio.to_thread( + get_all_items_from_collections, collection_names + ) else: query_result = await query_collection( request, diff --git a/backend/open_webui/retrieval/vector/async_client.py b/backend/open_webui/retrieval/vector/async_client.py new file mode 100644 index 0000000000..bb49f003a7 --- /dev/null +++ b/backend/open_webui/retrieval/vector/async_client.py @@ -0,0 +1,135 @@ +""" +Async facade over the synchronous VECTOR_DB_CLIENT. + +The vector DB backends bundled with Open WebUI (Chroma, pgvector, Qdrant, +Milvus, OpenSearch, Pinecone, Weaviate, …) all expose a uniformly +synchronous API. Each method performs blocking network or disk I/O — and +some, like `insert`/`upsert`, can run for several seconds. + +When such a sync method is awaited from an async route handler, it blocks +the event loop for its entire duration, freezing every other in-flight +HTTP request, websocket message and background task. + +This module wraps the sync client in an `AsyncVectorDBClient` that +transparently dispatches each call to a worker thread via +`asyncio.to_thread`. Async callers can `await ASYNC_VECTOR_DB_CLIENT.x(...)` +in place of `VECTOR_DB_CLIENT.x(...)` and the loop stays responsive. + +The original `VECTOR_DB_CLIENT` is unchanged, so callers already running +inside `run_in_threadpool` (e.g. `save_docs_to_vector_db`) are not +affected. + +Thread-safety expectations +-------------------------- +Every async caller now invokes `VECTOR_DB_CLIENT` from a worker thread +rather than the event-loop thread, and many can run concurrently. The +sync client (and its underlying backend driver) is therefore expected +to be safe for concurrent use across threads, which is the standard +contract for the bundled drivers (chroma, pgvector via SQLAlchemy +pool, qdrant-client, opensearch-py, …). This is *not* a new exposure +introduced by this facade — `save_docs_to_vector_db` already called +the sync client from `run_in_threadpool`, so concurrent threaded +access has always been a requirement of the codebase. Adding a global +serialization lock here would defeat the responsiveness this facade +exists to provide; any backend that genuinely cannot tolerate +concurrent access should grow its own internal serialization. + +API surface +----------- +Method signatures mirror `VectorDBBase` exactly. This is deliberate: +permissive `*args/**kwargs` forwarding hides typos at the call site +(an earlier revision of this file shipped that, and a `metadata=` +typo silently broke an entire endpoint until explicit signatures +surfaced it). Callers that need a backend-specific parameter not on +`VectorDBBase` should reach for the `.sync` escape hatch and wrap +their own `asyncio.to_thread`, e.g. :: + + await asyncio.to_thread( + ASYNC_VECTOR_DB_CLIENT.sync.some_backend_specific_op, + collection_name, special_kwarg=value, + ) +""" + +from __future__ import annotations + +import asyncio +from typing import Dict, List, Optional, Union + +from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT +from open_webui.retrieval.vector.main import ( + GetResult, + SearchResult, + VectorDBBase, + VectorItem, +) + + +class AsyncVectorDBClient: + """Awaitable mirror of `VectorDBBase` that off-loads each call to a thread. + + Method signatures mirror `VectorDBBase` exactly so static analysis + catches bad kwargs at the call site instead of letting them surface + deep inside the worker thread (where the resulting ``TypeError`` is + typically swallowed by surrounding ``try/except``). + """ + + def __init__(self, sync_client: VectorDBBase) -> None: + self._sync = sync_client + + @property + def sync(self) -> VectorDBBase: + """Escape hatch for code that must call the sync client directly + (e.g. already inside a worker thread).""" + return self._sync + + async def has_collection(self, collection_name: str) -> bool: + return await asyncio.to_thread(self._sync.has_collection, collection_name) + + async def delete_collection(self, collection_name: str) -> None: + return await asyncio.to_thread(self._sync.delete_collection, collection_name) + + async def insert(self, collection_name: str, items: List[VectorItem]) -> None: + return await asyncio.to_thread(self._sync.insert, collection_name, items) + + async def upsert(self, collection_name: str, items: List[VectorItem]) -> None: + return await asyncio.to_thread(self._sync.upsert, collection_name, items) + + async def search( + self, + collection_name: str, + vectors: List[List[Union[float, int]]], + filter: Optional[Dict] = None, + limit: int = 10, + ) -> Optional[SearchResult]: + return await asyncio.to_thread( + self._sync.search, collection_name, vectors, filter, limit + ) + + async def query( + self, + collection_name: str, + filter: Dict, + limit: Optional[int] = None, + ) -> Optional[GetResult]: + return await asyncio.to_thread( + self._sync.query, collection_name, filter, limit + ) + + async def get(self, collection_name: str) -> Optional[GetResult]: + return await asyncio.to_thread(self._sync.get, collection_name) + + async def delete( + self, + collection_name: str, + ids: Optional[List[str]] = None, + filter: Optional[Dict] = None, + ) -> None: + return await asyncio.to_thread( + self._sync.delete, collection_name, ids, filter + ) + + async def reset(self) -> None: + return await asyncio.to_thread(self._sync.reset) + + +ASYNC_VECTOR_DB_CLIENT = AsyncVectorDBClient(VECTOR_DB_CLIENT) diff --git a/backend/open_webui/routers/files.py b/backend/open_webui/routers/files.py index 66d7539278..055e09c6bf 100644 --- a/backend/open_webui/routers/files.py +++ b/backend/open_webui/routers/files.py @@ -25,7 +25,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from open_webui.internal.db import get_async_session, get_async_db_context from open_webui.constants import ERROR_MESSAGES -from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT +from open_webui.retrieval.vector.async_client import ASYNC_VECTOR_DB_CLIENT from open_webui.models.channels import Channels from open_webui.models.users import Users @@ -407,7 +407,7 @@ async def delete_all_files(user=Depends(get_admin_user), db: AsyncSession = Depe if result: try: Storage.delete_all_files() - VECTOR_DB_CLIENT.reset() + await ASYNC_VECTOR_DB_CLIENT.reset() except Exception as e: log.exception(e) log.error('Error deleting files') @@ -577,7 +577,7 @@ async def update_file_data_content_by_id( for knowledge in knowledges: try: # Remove old embeddings for this file from the KB collection - VECTOR_DB_CLIENT.delete(collection_name=knowledge.id, filter={'file_id': id}) + await ASYNC_VECTOR_DB_CLIENT.delete(collection_name=knowledge.id, filter={'file_id': id}) # Re-add from the now-updated file-{file_id} collection await process_file( request, @@ -789,9 +789,9 @@ async def delete_file_by_id(id: str, user=Depends(get_verified_user), db: AsyncS await Knowledges.remove_file_from_knowledge_by_id(knowledge.id, id, db=db) # Clean KB embeddings (same logic as /knowledge/{id}/file/remove) try: - VECTOR_DB_CLIENT.delete(collection_name=knowledge.id, filter={'file_id': id}) + await ASYNC_VECTOR_DB_CLIENT.delete(collection_name=knowledge.id, filter={'file_id': id}) if file.hash: - VECTOR_DB_CLIENT.delete(collection_name=knowledge.id, filter={'hash': file.hash}) + await ASYNC_VECTOR_DB_CLIENT.delete(collection_name=knowledge.id, filter={'hash': file.hash}) except Exception as e: log.debug(f'KB embedding cleanup for {knowledge.id}: {e}') @@ -799,7 +799,7 @@ async def delete_file_by_id(id: str, user=Depends(get_verified_user), db: AsyncS if result: try: Storage.delete_file(file.path) - VECTOR_DB_CLIENT.delete(collection_name=f'file-{id}') + await ASYNC_VECTOR_DB_CLIENT.delete(collection_name=f'file-{id}') except Exception as e: log.exception(e) log.error('Error deleting files') diff --git a/backend/open_webui/routers/knowledge.py b/backend/open_webui/routers/knowledge.py index c6f8ce5ecd..77b72cacf0 100644 --- a/backend/open_webui/routers/knowledge.py +++ b/backend/open_webui/routers/knowledge.py @@ -19,7 +19,7 @@ from open_webui.models.knowledge import ( KnowledgeUserResponse, ) from open_webui.models.files import Files, FileModel, FileMetadataResponse -from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT +from open_webui.retrieval.vector.async_client import ASYNC_VECTOR_DB_CLIENT from open_webui.routers.retrieval import ( process_file, ProcessFileForm, @@ -66,7 +66,7 @@ async def embed_knowledge_base_metadata( try: content = f'{name}\n\n{description}' if description else name embedding = await request.app.state.EMBEDDING_FUNCTION(content) - VECTOR_DB_CLIENT.upsert( + await ASYNC_VECTOR_DB_CLIENT.upsert( collection_name=KNOWLEDGE_BASES_COLLECTION, items=[ { @@ -85,10 +85,10 @@ async def embed_knowledge_base_metadata( return False -def remove_knowledge_base_metadata_embedding(knowledge_base_id: str) -> bool: +async def remove_knowledge_base_metadata_embedding(knowledge_base_id: str) -> bool: """Remove knowledge base embedding.""" try: - VECTOR_DB_CLIENT.delete( + await ASYNC_VECTOR_DB_CLIENT.delete( collection_name=KNOWLEDGE_BASES_COLLECTION, ids=[knowledge_base_id], ) @@ -310,8 +310,8 @@ async def reindex_knowledge_files( try: files = await Knowledges.get_files_by_id(knowledge_base.id, db=db) try: - if VECTOR_DB_CLIENT.has_collection(collection_name=knowledge_base.id): - VECTOR_DB_CLIENT.delete_collection(collection_name=knowledge_base.id) + if await ASYNC_VECTOR_DB_CLIENT.has_collection(collection_name=knowledge_base.id): + await ASYNC_VECTOR_DB_CLIENT.delete_collection(collection_name=knowledge_base.id) except Exception as e: log.error(f'Error deleting collection {knowledge_base.id}: {str(e)}') continue # Skip, don't raise @@ -732,7 +732,7 @@ async def update_file_from_knowledge_by_id( ) # Remove content from the vector database - VECTOR_DB_CLIENT.delete(collection_name=knowledge.id, filter={'file_id': form_data.file_id}) + await ASYNC_VECTOR_DB_CLIENT.delete(collection_name=knowledge.id, filter={'file_id': form_data.file_id}) # Add content to the vector database try: @@ -814,11 +814,11 @@ async def remove_file_from_knowledge_by_id( # Remove content from the vector database try: - VECTOR_DB_CLIENT.delete( + await ASYNC_VECTOR_DB_CLIENT.delete( collection_name=knowledge.id, filter={'file_id': form_data.file_id} ) # Remove by file_id first - VECTOR_DB_CLIENT.delete( + await ASYNC_VECTOR_DB_CLIENT.delete( collection_name=knowledge.id, filter={'hash': file.hash} ) # Remove by hash as well in case of duplicates except Exception as e: @@ -830,8 +830,8 @@ async def remove_file_from_knowledge_by_id( try: # Remove the file's collection from vector database file_collection = f'file-{form_data.file_id}' - if VECTOR_DB_CLIENT.has_collection(collection_name=file_collection): - VECTOR_DB_CLIENT.delete_collection(collection_name=file_collection) + if await ASYNC_VECTOR_DB_CLIENT.has_collection(collection_name=file_collection): + await ASYNC_VECTOR_DB_CLIENT.delete_collection(collection_name=file_collection) except Exception as e: log.debug('This was most likely caused by bypassing embedding processing') log.debug(e) @@ -915,13 +915,13 @@ async def delete_knowledge_by_id( # Clean up vector DB try: - VECTOR_DB_CLIENT.delete_collection(collection_name=id) + await ASYNC_VECTOR_DB_CLIENT.delete_collection(collection_name=id) except Exception as e: log.debug(e) pass # Remove knowledge base embedding - remove_knowledge_base_metadata_embedding(id) + await remove_knowledge_base_metadata_embedding(id) result = await Knowledges.delete_knowledge_by_id(id=id, db=db) return result @@ -960,7 +960,7 @@ async def reset_knowledge_by_id( ) try: - VECTOR_DB_CLIENT.delete_collection(collection_name=id) + await ASYNC_VECTOR_DB_CLIENT.delete_collection(collection_name=id) except Exception as e: log.debug(e) pass diff --git a/backend/open_webui/routers/memories.py b/backend/open_webui/routers/memories.py index cfd8274812..3a42801d01 100644 --- a/backend/open_webui/routers/memories.py +++ b/backend/open_webui/routers/memories.py @@ -5,7 +5,7 @@ import asyncio from typing import Optional from open_webui.models.memories import Memories, MemoryModel -from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT +from open_webui.retrieval.vector.async_client import ASYNC_VECTOR_DB_CLIENT from open_webui.utils.auth import get_verified_user from open_webui.internal.db import get_async_session from sqlalchemy.ext.asyncio import AsyncSession @@ -85,7 +85,7 @@ async def add_memory( vector = await request.app.state.EMBEDDING_FUNCTION(memory.content, user=user) - VECTOR_DB_CLIENT.upsert( + await ASYNC_VECTOR_DB_CLIENT.upsert( collection_name=f'user-memory-{user.id}', items=[ { @@ -138,7 +138,7 @@ async def query_memory( vector = await request.app.state.EMBEDDING_FUNCTION(form_data.content, user=user) - results = VECTOR_DB_CLIENT.search( + results = await ASYNC_VECTOR_DB_CLIENT.search( collection_name=f'user-memory-{user.id}', vectors=[vector], limit=form_data.k, @@ -175,7 +175,7 @@ async def reset_memory_from_vector_db( detail=ERROR_MESSAGES.ACCESS_PROHIBITED, ) - VECTOR_DB_CLIENT.delete_collection(f'user-memory-{user.id}') + await ASYNC_VECTOR_DB_CLIENT.delete_collection(f'user-memory-{user.id}') memories = await Memories.get_memories_by_user_id(user.id) @@ -184,7 +184,7 @@ async def reset_memory_from_vector_db( *[request.app.state.EMBEDDING_FUNCTION(memory.content, user=user) for memory in memories] ) - VECTOR_DB_CLIENT.upsert( + await ASYNC_VECTOR_DB_CLIENT.upsert( collection_name=f'user-memory-{user.id}', items=[ { @@ -230,7 +230,7 @@ async def delete_memory_by_user_id( if result: try: - VECTOR_DB_CLIENT.delete_collection(f'user-memory-{user.id}') + await ASYNC_VECTOR_DB_CLIENT.delete_collection(f'user-memory-{user.id}') except Exception as e: log.error(e) return True @@ -273,7 +273,7 @@ async def update_memory_by_id( if form_data.content is not None: vector = await request.app.state.EMBEDDING_FUNCTION(memory.content, user=user) - VECTOR_DB_CLIENT.upsert( + await ASYNC_VECTOR_DB_CLIENT.upsert( collection_name=f'user-memory-{user.id}', items=[ { @@ -318,7 +318,7 @@ async def delete_memory_by_id( result = await Memories.delete_memory_by_id_and_user_id(memory_id, user.id, db=db) if result: - VECTOR_DB_CLIENT.delete(collection_name=f'user-memory-{user.id}', ids=[memory_id]) + await ASYNC_VECTOR_DB_CLIENT.delete(collection_name=f'user-memory-{user.id}', ids=[memory_id]) return True return False diff --git a/backend/open_webui/routers/retrieval.py b/backend/open_webui/routers/retrieval.py index 417441305e..d58ace149f 100644 --- a/backend/open_webui/routers/retrieval.py +++ b/backend/open_webui/routers/retrieval.py @@ -45,6 +45,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT +from open_webui.retrieval.vector.async_client import ASYNC_VECTOR_DB_CLIENT # Document loaders from open_webui.retrieval.loaders.main import Loader @@ -1556,7 +1557,7 @@ async def process_file( try: # /files/{file_id}/data/content/update - VECTOR_DB_CLIENT.delete_collection(collection_name=f'file-{file.id}') + await ASYNC_VECTOR_DB_CLIENT.delete_collection(collection_name=f'file-{file.id}') except Exception: # Audio file upload pipeline pass @@ -1579,7 +1580,9 @@ async def process_file( # Check if the file has already been processed and save the content # Usage: /knowledge/{id}/file/add, /knowledge/{id}/file/update - result = VECTOR_DB_CLIENT.query(collection_name=f'file-{file.id}', filter={'file_id': file.id}) + result = await ASYNC_VECTOR_DB_CLIENT.query( + collection_name=f'file-{file.id}', filter={'file_id': file.id} + ) if result is not None and len(result.ids[0]) > 0: docs = [ @@ -2380,7 +2383,7 @@ async def query_doc_handler( try: if request.app.state.config.ENABLE_RAG_HYBRID_SEARCH and (form_data.hybrid is None or form_data.hybrid): collection_results = {} - collection_results[form_data.collection_name] = VECTOR_DB_CLIENT.get( + 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( @@ -2409,7 +2412,10 @@ async def query_doc_handler( query_embedding = await request.app.state.EMBEDDING_FUNCTION( form_data.query, prefix=RAG_EMBEDDING_QUERY_PREFIX, user=user ) - return query_doc( + # query_doc wraps a blocking VECTOR_DB_CLIENT.search call; + # offload so the request's event loop stays responsive. + return await asyncio.to_thread( + query_doc, collection_name=form_data.collection_name, query_embedding=query_embedding, k=form_data.k if form_data.k else request.app.state.config.TOP_K, @@ -2507,7 +2513,7 @@ async def delete_entries_from_collection( db: AsyncSession = Depends(get_async_session), ): try: - if VECTOR_DB_CLIENT.has_collection(collection_name=form_data.collection_name): + if await ASYNC_VECTOR_DB_CLIENT.has_collection(collection_name=form_data.collection_name): file = await Files.get_file_by_id(form_data.file_id, db=db) if not file: raise HTTPException( @@ -2516,13 +2522,39 @@ async def delete_entries_from_collection( ) hash = file.hash - VECTOR_DB_CLIENT.delete( + # Refuse to issue a `filter={'hash': None}` query — the + # match semantics of a null filter value are + # backend-dependent (some backends ignore the key, some + # match every row whose metadata lacks `hash`) and risk + # deleting unrelated entries. Files without a hash are + # typically unprocessed / failed / legacy records that + # can't be targeted by hash anyway. + if hash is None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=ERROR_MESSAGES.DEFAULT( + 'File has no hash; cannot delete vector entries by hash.' + ), + ) + + # Pre-existing bug: this used `metadata=` which is not a + # parameter on `VectorDBBase.delete` nor on any backend + # implementation, so the call always raised TypeError that + # was silently swallowed by the surrounding `except + # Exception` and the endpoint reported `{'status': False}` + # for every request. Use `filter` to actually do what the + # endpoint name promises. + await ASYNC_VECTOR_DB_CLIENT.delete( collection_name=form_data.collection_name, - metadata={'hash': hash}, + filter={'hash': hash}, ) return {'status': True} else: return {'status': False} + except HTTPException: + # Caller-meaningful errors (404/400 above) must not be + # swallowed and re-shaped as `{'status': False}`. + raise except Exception as e: log.exception(e) return {'status': False} @@ -2530,7 +2562,7 @@ async def delete_entries_from_collection( @router.post('/reset/db') async def reset_vector_db(user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)): - VECTOR_DB_CLIENT.reset() + await ASYNC_VECTOR_DB_CLIENT.reset() await Knowledges.delete_all_knowledge(db=db) diff --git a/backend/open_webui/tools/builtin.py b/backend/open_webui/tools/builtin.py index 61cd5ede4e..8aa27520f3 100644 --- a/backend/open_webui/tools/builtin.py +++ b/backend/open_webui/tools/builtin.py @@ -37,7 +37,7 @@ from open_webui.models.channels import Channels, ChannelMember, Channel from open_webui.models.messages import Messages, Message from open_webui.models.groups import Groups from open_webui.models.memories import Memories -from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT +from open_webui.retrieval.vector.async_client import ASYNC_VECTOR_DB_CLIENT from open_webui.utils.sanitize import sanitize_code log = logging.getLogger(__name__) @@ -653,7 +653,7 @@ async def delete_memory( result = await Memories.delete_memory_by_id_and_user_id(memory_id, user.id) if result: - VECTOR_DB_CLIENT.delete(collection_name=f'user-memory-{user.id}', ids=[memory_id]) + await ASYNC_VECTOR_DB_CLIENT.delete(collection_name=f'user-memory-{user.id}', ids=[memory_id]) return json.dumps( {'status': 'success', 'message': f'Memory {memory_id} deleted'}, ensure_ascii=False, @@ -2202,7 +2202,7 @@ async def query_knowledge_bases( import heapq from open_webui.models.knowledge import Knowledges from open_webui.routers.knowledge import KNOWLEDGE_BASES_COLLECTION - from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT + from open_webui.retrieval.vector.async_client import ASYNC_VECTOR_DB_CLIENT user_id = __user__.get('id') user_group_ids = [group.id for group in await Groups.get_groups_by_member_id(user_id)] @@ -2227,7 +2227,7 @@ async def query_knowledge_bases( accessible_ids = [kb.id for kb in accessible_knowledge_bases.items] - search_results = VECTOR_DB_CLIENT.search( + search_results = await ASYNC_VECTOR_DB_CLIENT.search( collection_name=KNOWLEDGE_BASES_COLLECTION, vectors=[query_embedding], filter={'knowledge_base_id': {'$in': accessible_ids}},