codekingpro/portable-devtools
115k
1from typing import TYPE_CHECKING, Optional, Union, List, cast
2
3from chromadb.api.types import (
4 URI,
5 CollectionMetadata,
6 Embedding,
7 PyEmbedding,
8 Include,
9 IndexingStatus,
10 Metadata,
11 Document,
12 Image,
13 Where,
14 IDs,
15 GetResult,
16 QueryResult,
17 ID,
18 OneOrMany,
19 ReadLevel,
20 WhereDocument,
21 SearchResult,
22 DeleteResult,
23 maybe_cast_one_to_many,
24)
25
26from chromadb.api.models.CollectionCommon import CollectionCommon
27from chromadb.api.collection_configuration import UpdateCollectionConfiguration
28from chromadb.execution.expression.plan import Search
29
30if TYPE_CHECKING:
31 from chromadb.api import AsyncServerAPI # noqa: F401
32
33
34class AsyncCollection(CollectionCommon["AsyncServerAPI"]):
35 async def add(
36 self,
37 ids: OneOrMany[ID],
38 embeddings: Optional[
39 Union[
40 OneOrMany[Embedding],
41 OneOrMany[PyEmbedding],
42 ]
43 ] = None,
44 metadatas: Optional[OneOrMany[Metadata]] = None,
45 documents: Optional[OneOrMany[Document]] = None,
46 images: Optional[OneOrMany[Image]] = None,
47 uris: Optional[OneOrMany[URI]] = None,
48 ) -> None:
49 """Add embeddings to the data store.
50 Args:
51 ids: The ids of the embeddings you wish to add
52 embeddings: The embeddings to add. If None, embeddings will be computed based on the documents or images using the embedding_function set for the Collection. Optional.
53 metadatas: The metadata to associate with the embeddings. When querying, you can filter on this metadata. Optional.
54 documents: The documents to associate with the embeddings. Optional.
55 images: The images to associate with the embeddings. Optional.
56 uris: The uris of the images to associate with the embeddings. Optional.
57
58 Returns:
59 None
60
61 Raises:
62 ValueError: If you don't provide either embeddings or documents
63 ValueError: If the length of ids, embeddings, metadatas, or documents don't match
64 ValueError: If you don't provide an embedding function and don't provide embeddings
65 ValueError: If you provide both embeddings and documents
66 ValueError: If you provide an id that already exists
67
68 """
69 add_request = self._validate_and_prepare_add_request(
70 ids=ids,
71 embeddings=embeddings,
72 metadatas=metadatas,
73 documents=documents,
74 images=images,
75 uris=uris,
76 )
77
78 await self._client._add(
79 collection_id=self.id,
80 ids=add_request["ids"],
81 embeddings=add_request["embeddings"],
82 metadatas=add_request["metadatas"],
83 documents=add_request["documents"],
84 uris=add_request["uris"],
85 tenant=self.tenant,
86 database=self.database,
87 )
88
89 async def count(self, read_level: ReadLevel = ReadLevel.INDEX_AND_WAL) -> int:
90 """Return the number of records in the collection.
91
92 Args:
93 read_level: Controls whether to read from the write-ahead log (WAL):
94 - ReadLevel.INDEX_AND_WAL: Read from both the compacted index and WAL (default).
95 All committed writes will be visible.
96 - ReadLevel.INDEX_ONLY: Read only from the compacted index, skipping the WAL.
97 Faster, but recent writes that haven't been compacted may not be visible.
98 - ReadLevel.INDEX_AND_BOUNDED_WAL: Read from the index and up to a
99 server-configured number of WAL entries for bounded query latency.
100
101 Returns:
102 int: The total number of embeddings added to the database
103 """
104 return await self._client._count(
105 collection_id=self.id,
106 tenant=self.tenant,
107 database=self.database,
108 read_level=read_level,
109 )
110
111 async def get_indexing_status(self) -> IndexingStatus:
112 """Get the indexing status of this collection.
113
114 Returns:
115 IndexingStatus: An object containing:
116 - num_indexed_ops: Number of user operations that have been indexed
117 - num_unindexed_ops: Number of user operations pending indexing
118 - total_ops: Total number of user operations in collection
119 - op_indexing_progress: Proportion of user operations that have been indexed as a float between 0 and 1
120 """
121 return await self._client._get_indexing_status(
122 collection_id=self.id,
123 tenant=self.tenant,
124 database=self.database,
125 )
126
127 async def get(
128 self,
129 ids: Optional[OneOrMany[ID]] = None,
130 where: Optional[Where] = None,
131 limit: Optional[int] = None,
132 offset: Optional[int] = None,
133 where_document: Optional[WhereDocument] = None,
134 include: Include = ["metadatas", "documents"],
135 ) -> GetResult:
136 """Get embeddings and their associate data from the data store. If no ids or where filter is provided returns
137 all embeddings up to limit starting at offset.
138
139 Args:
140 ids: The ids of the embeddings to get. Optional.
141 where: A Where type dict used to filter results by. E.g. `{"$and": [{"color" : "red"}, {"price": {"$gte": 4.20}}]}`. Optional.
142 limit: The number of documents to return. Optional.
143 offset: The offset to start returning results from. Useful for paging results with limit. Optional.
144 where_document: A WhereDocument type dict used to filter by the documents. E.g. `{"$contains": "hello"}`. Optional.
145 include: A list of what to include in the results. Can contain `"embeddings"`, `"metadatas"`, `"documents"`. Ids are always included. Defaults to `["metadatas", "documents"]`. Optional.
146
147 Returns:
148 GetResult: A GetResult object containing the results.
149
150 """
151 get_request = self._validate_and_prepare_get_request(
152 ids=ids,
153 where=where,
154 where_document=where_document,
155 include=include,
156 )
157
158 get_results = await self._client._get(
159 collection_id=self.id,
160 ids=get_request["ids"],
161 where=get_request["where"],
162 where_document=get_request["where_document"],
163 include=get_request["include"],
164 limit=limit,
165 offset=offset,
166 tenant=self.tenant,
167 database=self.database,
168 )
169
170 return self._transform_get_response(
171 response=get_results, include=get_request["include"]
172 )
173
174 async def peek(self, limit: int = 10) -> GetResult:
175 """Get the first few results in the database up to limit
176
177 Args:
178 limit: The number of results to return.
179
180 Returns:
181 GetResult: A GetResult object containing the results.
182 """
183 return self._transform_peek_response(
184 await self._client._peek(
185 collection_id=self.id,
186 n=limit,
187 tenant=self.tenant,
188 database=self.database,
189 )
190 )
191
192 async def query(
193 self,
194 query_embeddings: Optional[
195 Union[
196 OneOrMany[Embedding],
197 OneOrMany[PyEmbedding],
198 ]
199 ] = None,
200 query_texts: Optional[OneOrMany[Document]] = None,
201 query_images: Optional[OneOrMany[Image]] = None,
202 query_uris: Optional[OneOrMany[URI]] = None,
203 ids: Optional[OneOrMany[ID]] = None,
204 n_results: int = 10,
205 where: Optional[Where] = None,
206 where_document: Optional[WhereDocument] = None,
207 include: Include = [
208 "metadatas",
209 "documents",
210 "distances",
211 ],
212 ) -> QueryResult:
213 """Get the n_results nearest neighbor embeddings for provided query_embeddings or query_texts.
214
215 Args:
216 query_embeddings: The embeddings to get the closes neighbors of. Optional.
217 query_texts: The document texts to get the closes neighbors of. Optional.
218 query_images: The images to get the closes neighbors of. Optional.
219 ids: A subset of ids to search within. Optional.
220 n_results: The number of neighbors to return for each query_embedding or query_texts. Optional.
221 where: A Where type dict used to filter results by. E.g. `{"$and": [{"color" : "red"}, {"price": {"$gte": 4.20}}]}`. Optional.
222 where_document: A WhereDocument type dict used to filter by the documents. E.g. `{"$contains": "hello"}`. Optional.
223 include: A list of what to include in the results. Can contain `"embeddings"`, `"metadatas"`, `"documents"`, `"distances"`. Ids are always included. Defaults to `["metadatas", "documents", "distances"]`. Optional.
224
225 Returns:
226 QueryResult: A QueryResult object containing the results.
227
228 Raises:
229 ValueError: If you don't provide either query_embeddings, query_texts, or query_images
230 ValueError: If you provide both query_embeddings and query_texts
231 ValueError: If you provide both query_embeddings and query_images
232 ValueError: If you provide both query_texts and query_images
233
234 """
235
236 query_request = self._validate_and_prepare_query_request(
237 query_embeddings=query_embeddings,
238 query_texts=query_texts,
239 query_images=query_images,
240 query_uris=query_uris,
241 ids=ids,
242 n_results=n_results,
243 where=where,
244 where_document=where_document,
245 include=include,
246 )
247
248 query_results = await self._client._query(
249 collection_id=self.id,
250 ids=query_request["ids"],
251 query_embeddings=query_request["embeddings"],
252 n_results=query_request["n_results"],
253 where=query_request["where"],
254 where_document=query_request["where_document"],
255 include=query_request["include"],
256 tenant=self.tenant,
257 database=self.database,
258 )
259
260 return self._transform_query_response(
261 response=query_results, include=query_request["include"]
262 )
263
264 async def modify(
265 self,
266 name: Optional[str] = None,
267 metadata: Optional[CollectionMetadata] = None,
268 configuration: Optional[UpdateCollectionConfiguration] = None,
269 ) -> None:
270 """Modify the collection name or metadata
271
272 Args:
273 name: The updated name for the collection. Optional.
274 metadata: The updated metadata for the collection. Optional.
275
276 Returns:
277 None
278 """
279
280 self._validate_modify_request(metadata)
281
282 # Note there is a race condition here where the metadata can be updated
283 # but another thread sees the cached local metadata.
284 # TODO: fixme
285 await self._client._modify(
286 id=self.id,
287 new_name=name,
288 new_metadata=metadata,
289 new_configuration=configuration,
290 tenant=self.tenant,
291 database=self.database,
292 )
293
294 self._update_model_after_modify_success(name, metadata, configuration)
295
296 async def fork(
297 self,
298 new_name: str,
299 ) -> "AsyncCollection":
300 """Fork the current collection under a new name. The returning collection should contain identical data to the current collection.
301 This only works for Hosted Chroma for now.
302
303 Args:
304 new_name: The name of the new collection.
305
306 Returns:
307 Collection: A new collection with the specified name and containing identical data to the current collection.
308 """
309 model = await self._client._fork(
310 collection_id=self.id,
311 new_name=new_name,
312 tenant=self.tenant,
313 database=self.database,
314 )
315 return AsyncCollection(
316 client=self._client,
317 model=model,
318 embedding_function=self._embedding_function,
319 data_loader=self._data_loader,
320 )
321
322 async def fork_count(self) -> int:
323 """Get the number of forks that exist for this collection.
324 This only works for Hosted Chroma for now.
325
326 Returns:
327 int: The number of forks for this collection.
328 """
329 return await self._client._fork_count(
330 collection_id=self.id,
331 tenant=self.tenant,
332 database=self.database,
333 )
334
335 async def search(
336 self,
337 searches: OneOrMany[Search],
338 read_level: ReadLevel = ReadLevel.INDEX_AND_WAL,
339 ) -> SearchResult:
340 """Perform hybrid search on the collection.
341 This is an experimental API that only works for Hosted Chroma for now.
342
343 Args:
344 searches: A single Search object or a list of Search objects, each containing:
345 - where: Where expression for filtering
346 - rank: Ranking expression for hybrid search (defaults to Val(0.0))
347 - limit: Limit configuration for pagination (defaults to no limit)
348 - select: Select configuration for keys to return (defaults to empty)
349 read_level: Controls whether to read from the write-ahead log (WAL):
350 - ReadLevel.INDEX_AND_WAL: Read from both the compacted index and WAL (default).
351 All committed writes will be visible.
352 - ReadLevel.INDEX_ONLY: Read only from the compacted index, skipping the WAL.
353 Faster, but recent writes that haven't been compacted may not be visible.
354 - ReadLevel.INDEX_AND_BOUNDED_WAL: Read from the index and up to a
355 server-configured number of WAL entries for bounded query latency.
356
357 Returns:
358 SearchResult: Column-major format response with:
359 - ids: List of result IDs for each search payload
360 - documents: Optional documents for each payload
361 - embeddings: Optional embeddings for each payload
362 - metadatas: Optional metadata for each payload
363 - scores: Optional scores for each payload
364 - select: List of selected keys for each payload
365
366 Raises:
367 NotImplementedError: For local/segment API implementations
368
369 Examples:
370 # Using builder pattern with Key constants
371 from chromadb.execution.expression import (
372 Search, Key, K, Knn, Val
373 )
374
375 # Note: K is an alias for Key, so K.DOCUMENT == Key.DOCUMENT
376 search = (Search()
377 .where((K("category") == "science") & (K("score") > 0.5))
378 .rank(Knn(query=[0.1, 0.2, 0.3]) * 0.8 + Val(0.5) * 0.2)
379 .limit(10, offset=0)
380 .select(K.DOCUMENT, K.SCORE, "title"))
381
382 # Direct construction
383 from chromadb.execution.expression import (
384 Search, Eq, And, Gt, Knn, Limit, Select, Key
385 )
386
387 search = Search(
388 where=And([Eq("category", "science"), Gt("score", 0.5)]),
389 rank=Knn(query=[0.1, 0.2, 0.3]),
390 limit=Limit(offset=0, limit=10),
391 select=Select(keys={Key.DOCUMENT, Key.SCORE, "title"})
392 )
393
394 # Single search
395 result = await collection.search(search)
396
397 # Multiple searches at once
398 searches = [
399 Search().where(K("type") == "article").rank(Knn(query=[0.1, 0.2])),
400 Search().where(K("type") == "paper").rank(Knn(query=[0.3, 0.4]))
401 ]
402 results = await collection.search(searches)
403
404 # Skip WAL for faster queries (may miss recent uncommitted writes)
405 from chromadb.api.types import ReadLevel
406 result = await collection.search(search, read_level=ReadLevel.INDEX_ONLY)
407 """
408 # Convert single search to list for consistent handling
409 searches_list = maybe_cast_one_to_many(searches)
410 if searches_list is None:
411 searches_list = []
412
413 # Embed any string queries in Knn objects
414 embedded_searches = [
415 self._embed_search_string_queries(search) for search in searches_list
416 ]
417
418 return await self._client._search(
419 collection_id=self.id,
420 searches=cast(List[Search], embedded_searches),
421 tenant=self.tenant,
422 database=self.database,
423 read_level=read_level,
424 )
425
426 async def update(
427 self,
428 ids: OneOrMany[ID],
429 embeddings: Optional[
430 Union[
431 OneOrMany[Embedding],
432 OneOrMany[PyEmbedding],
433 ]
434 ] = None,
435 metadatas: Optional[OneOrMany[Metadata]] = None,
436 documents: Optional[OneOrMany[Document]] = None,
437 images: Optional[OneOrMany[Image]] = None,
438 uris: Optional[OneOrMany[URI]] = None,
439 ) -> None:
440 """Update the embeddings, metadatas or documents for provided ids.
441
442 Args:
443 ids: The ids of the embeddings to update
444 embeddings: The embeddings to update. If None, embeddings will be computed based on the documents or images using the embedding_function set for the Collection. Optional.
445 metadatas: The metadata to associate with the embeddings. When querying, you can filter on this metadata. Optional.
446 documents: The documents to associate with the embeddings. Optional.
447 images: The images to associate with the embeddings. Optional.
448 Returns:
449 None
450 """
451 update_request = self._validate_and_prepare_update_request(
452 ids=ids,
453 embeddings=embeddings,
454 metadatas=metadatas,
455 documents=documents,
456 images=images,
457 uris=uris,
458 )
459
460 await self._client._update(
461 collection_id=self.id,
462 ids=update_request["ids"],
463 embeddings=update_request["embeddings"],
464 metadatas=update_request["metadatas"],
465 documents=update_request["documents"],
466 uris=update_request["uris"],
467 tenant=self.tenant,
468 database=self.database,
469 )
470
471 async def upsert(
472 self,
473 ids: OneOrMany[ID],
474 embeddings: Optional[
475 Union[
476 OneOrMany[Embedding],
477 OneOrMany[PyEmbedding],
478 ]
479 ] = None,
480 metadatas: Optional[OneOrMany[Metadata]] = None,
481 documents: Optional[OneOrMany[Document]] = None,
482 images: Optional[OneOrMany[Image]] = None,
483 uris: Optional[OneOrMany[URI]] = None,
484 ) -> None:
485 """Update the embeddings, metadatas or documents for provided ids, or create them if they don't exist.
486
487 Args:
488 ids: The ids of the embeddings to update
489 embeddings: The embeddings to add. If None, embeddings will be computed based on the documents using the embedding_function set for the Collection. Optional.
490 metadatas: The metadata to associate with the embeddings. When querying, you can filter on this metadata. Optional.
491 documents: The documents to associate with the embeddings. Optional.
492
493 Returns:
494 None
495 """
496 upsert_request = self._validate_and_prepare_upsert_request(
497 ids=ids,
498 embeddings=embeddings,
499 metadatas=metadatas,
500 documents=documents,
501 images=images,
502 uris=uris,
503 )
504
505 await self._client._upsert(
506 collection_id=self.id,
507 ids=upsert_request["ids"],
508 embeddings=upsert_request["embeddings"],
509 metadatas=upsert_request["metadatas"],
510 documents=upsert_request["documents"],
511 uris=upsert_request["uris"],
512 tenant=self.tenant,
513 database=self.database,
514 )
515
516 async def delete(
517 self,
518 ids: Optional[IDs] = None,
519 where: Optional[Where] = None,
520 where_document: Optional[WhereDocument] = None,
521 limit: Optional[int] = None,
522 ) -> DeleteResult:
523 """Delete the embeddings based on ids and/or a where filter
524
525 Args:
526 ids: The ids of the embeddings to delete
527 where: A Where type dict used to filter the delection by. E.g. `{"$and": [{"color" : "red"}, {"price": {"$gte": 4.20}}]}`. Optional.
528 where_document: A WhereDocument type dict used to filter the deletion by the document content. E.g. `{"$contains": "hello"}`. Optional.
529 limit: Maximum number of records to delete. Can only be used with where or where_document filters.
530
531 Returns:
532 DeleteResult: A dict containing the number of records deleted.
533
534 Raises:
535 ValueError: If you don't provide either ids, where, or where_document
536 ValueError: If limit is specified without a where or where_document clause.
537 """
538 delete_request = self._validate_and_prepare_delete_request(
539 ids, where, where_document, limit=limit
540 )
541
542 return await self._client._delete(
543 collection_id=self.id,
544 ids=delete_request["ids"],
545 where=delete_request["where"],
546 where_document=delete_request["where_document"],
547 limit=delete_request["limit"],
548 tenant=self.tenant,
549 database=self.database,
550 )
551