Coverage for open_webui/retrieval/vector/async_client.py: 83%

1""" 

2Async facade over the synchronous VECTOR_DB_CLIENT. 

3 

4The vector DB backends bundled with Open WebUI (Chroma, pgvector, Qdrant, 

5Milvus, OpenSearch, Pinecone, Weaviate, …) all expose a uniformly 

6synchronous API. Each method performs blocking network or disk I/O — and 

7some, like `insert`/`upsert`, can run for several seconds. 

8 

9When such a sync method is awaited from an async route handler, it blocks 

10the event loop for its entire duration, freezing every other in-flight 

11HTTP request, websocket message and background task. 

12 

13This module wraps the sync client in an `AsyncVectorDBClient` that 

14transparently dispatches each call to a worker thread via 

15`asyncio.to_thread`. Async callers can `await ASYNC_VECTOR_DB_CLIENT.x(...)` 

16in place of `VECTOR_DB_CLIENT.x(...)` and the loop stays responsive. 

17 

18Client initialization and calls run in the worker thread. Synchronous callers 

19already inside `run_in_threadpool` use `get_vector_db_client()` directly. 

20 

21Thread-safety expectations 

22-------------------------- 

23Every async caller now invokes `VECTOR_DB_CLIENT` from a worker thread 

24rather than the event-loop thread, and many can run concurrently. The 

25sync client (and its underlying backend driver) is therefore expected 

26to be safe for concurrent use across threads, which is the standard 

27contract for the bundled drivers (chroma, pgvector via SQLAlchemy 

28pool, qdrant-client, opensearch-py, …). This is *not* a new exposure 

29introduced by this facade — `save_docs_to_vector_db` already called 

30the sync client from `run_in_threadpool`, so concurrent threaded 

31access has always been a requirement of the codebase. Adding a global 

32serialization lock here would defeat the responsiveness this facade 

33exists to provide; any backend that genuinely cannot tolerate 

34concurrent access should grow its own internal serialization. 

35 

36API surface 

37----------- 

38Method signatures mirror `VectorDBBase` exactly. This is deliberate: 

39permissive `*args/**kwargs` forwarding hides typos at the call site 

40(an earlier revision of this file shipped that, and a `metadata=` 

41typo silently broke an entire endpoint until explicit signatures 

42surfaced it). Callers that need a backend-specific parameter not on 

43`VectorDBBase` should reach for the `.sync` escape hatch and wrap 

44their own `asyncio.to_thread`, e.g. :: 

45 

46 await asyncio.to_thread( 

47 ASYNC_VECTOR_DB_CLIENT.sync.some_backend_specific_op, 

48 collection_name, special_kwarg=value, 

49 ) 

50""" 

51 

52from __future__ import annotations 

53 

54import asyncio 

55from typing import Dict, List, Optional, Union 

56 

57from open_webui.retrieval.vector.factory import get_vector_db_client 

58from open_webui.retrieval.vector.main import ( 

59 GetResult, 

60 SearchResult, 

61 VectorDBBase, 

62 VectorItem, 

63) 

64 

65 

66class AsyncVectorDBClient: 

67 """Awaitable mirror of `VectorDBBase` that off-loads each call to a thread. 

68 

69 Method signatures mirror `VectorDBBase` exactly so static analysis 

70 catches bad kwargs at the call site instead of letting them surface 

71 deep inside the worker thread (where the resulting ``TypeError`` is 

72 typically swallowed by surrounding ``try/except``). 

73 """ 

74 

75 def __init__(self, sync_client: Optional[VectorDBBase] = None) -> None: 

76 self._sync = sync_client 

77 

78 @property 

79 def sync(self) -> VectorDBBase: 

80 """Escape hatch for code that must call the sync client directly 

81 (e.g. already inside a worker thread).""" 

82 return self._sync if self._sync is not None else get_vector_db_client() 

83 

84 @property 

85 def supports_hybrid_search(self) -> bool: 

86 return type(self.sync).hybrid_search is not VectorDBBase.hybrid_search 

87 

88 async def has_collection(self, collection_name: str) -> bool: 

89 return await asyncio.to_thread(lambda: self.sync.has_collection(collection_name)) 

90 

91 async def delete_collection(self, collection_name: str) -> None: 

92 return await asyncio.to_thread(lambda: self.sync.delete_collection(collection_name)) 

93 

94 async def insert(self, collection_name: str, items: List[VectorItem]) -> None: 

95 return await asyncio.to_thread(lambda: self.sync.insert(collection_name, items)) 

96 

97 async def upsert(self, collection_name: str, items: List[VectorItem]) -> None: 

98 return await asyncio.to_thread(lambda: self.sync.upsert(collection_name, items)) 

99 

100 async def search( 

101 self, 

102 collection_name: str, 

103 vectors: List[List[Union[float, int]]], 

104 filter: Optional[Dict] = None, 

105 limit: int = 10, 

106 ) -> Optional[SearchResult]: 

107 return await asyncio.to_thread(lambda: self.sync.search(collection_name, vectors, filter, limit)) 

108 

109 async def hybrid_search( 

110 self, 

111 collection_name: str, 

112 query: str, 

113 vectors: List[List[Union[float, int]]], 

114 filter: Optional[Dict] = None, 

115 limit: int = 10, 

116 hybrid_bm25_weight: float = 0.5, 

117 ) -> Optional[SearchResult]: 

118 return await asyncio.to_thread( 

119 lambda: self.sync.hybrid_search(collection_name, query, vectors, filter, limit, hybrid_bm25_weight) 

120 ) 

121 

122 async def query( 

123 self, 

124 collection_name: str, 

125 filter: Dict, 

126 limit: Optional[int] = None, 

127 ) -> Optional[GetResult]: 

128 return await asyncio.to_thread(lambda: self.sync.query(collection_name, filter, limit)) 

129 

130 async def get(self, collection_name: str) -> Optional[GetResult]: 

131 return await asyncio.to_thread(lambda: self.sync.get(collection_name)) 

132 

133 async def delete( 

134 self, 

135 collection_name: str, 

136 ids: Optional[List[str]] = None, 

137 filter: Optional[Dict] = None, 

138 ) -> None: 

139 return await asyncio.to_thread(lambda: self.sync.delete(collection_name, ids, filter)) 

140 

141 async def reset(self) -> None: 

142 return await asyncio.to_thread(lambda: self.sync.reset()) 

143 

144 

145ASYNC_VECTOR_DB_CLIENT = AsyncVectorDBClient()