codekingpro/portable-devtools
115k
1from abc import ABC, abstractmethod
2from typing import Sequence, Optional, List
3from uuid import UUID
4
5from overrides import override
6from chromadb.api.collection_configuration import (
7 CreateCollectionConfiguration,
8 UpdateCollectionConfiguration,
9)
10from chromadb.auth import UserIdentity
11from chromadb.api.models.AsyncCollection import AsyncCollection
12from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT
13from chromadb.api.types import (
14 CollectionMetadata,
15 DeleteResult,
16 Documents,
17 Embeddable,
18 EmbeddingFunction,
19 DataLoader,
20 Embeddings,
21 IDs,
22 Include,
23 IndexingStatus,
24 Loadable,
25 Metadatas,
26 ReadLevel,
27 Schema,
28 URIs,
29 Where,
30 QueryResult,
31 GetResult,
32 WhereDocument,
33 IncludeMetadataDocuments,
34 IncludeMetadataDocumentsDistances,
35 SearchResult,
36 DefaultEmbeddingFunction,
37)
38from chromadb.execution.expression.plan import Search
39from chromadb.config import Component, Settings
40from chromadb.types import Database, Tenant, Collection as CollectionModel
41
42
43class AsyncBaseAPI(ABC):
44 @abstractmethod
45 async def heartbeat(self) -> int:
46 """Get the current time in nanoseconds since epoch.
47 Used to check if the server is alive.
48
49 Returns:
50 int: The current time in nanoseconds since epoch
51
52 """
53 pass
54
55 #
56 # COLLECTION METHODS
57 #
58
59 @abstractmethod
60 async def count_collections(self) -> int:
61 """Count the number of collections.
62
63 Returns:
64 int: The number of collections.
65
66 Examples:
67 ```python
68 await client.count_collections()
69 # 1
70 ```
71 """
72 pass
73
74 @abstractmethod
75 async def _modify(
76 self,
77 id: UUID,
78 new_name: Optional[str] = None,
79 new_metadata: Optional[CollectionMetadata] = None,
80 new_configuration: Optional[UpdateCollectionConfiguration] = None,
81 ) -> None:
82 """[Internal] Modify a collection by UUID. Can update the name and/or metadata.
83
84 Args:
85 id: The internal UUID of the collection to modify.
86 new_name: The new name of the collection.
87 If None, the existing name will remain. Defaults to None.
88 new_metadata: The new metadata to associate with the collection.
89 Defaults to None.
90 new_configuration: The new configuration to associate with the collection.
91 Defaults to None.
92 """
93 pass
94
95 @abstractmethod
96 async def delete_collection(
97 self,
98 name: str,
99 ) -> None:
100 """Delete a collection with the given name.
101 Args:
102 name: The name of the collection to delete.
103
104 Raises:
105 ValueError: If the collection does not exist.
106
107 Examples:
108 ```python
109 await client.delete_collection("my_collection")
110 ```
111 """
112 pass
113
114 #
115 # ITEM METHODS
116 #
117
118 @abstractmethod
119 async def _add(
120 self,
121 ids: IDs,
122 collection_id: UUID,
123 embeddings: Embeddings,
124 metadatas: Optional[Metadatas] = None,
125 documents: Optional[Documents] = None,
126 uris: Optional[URIs] = None,
127 ) -> bool:
128 """[Internal] Add embeddings to a collection specified by UUID.
129 If (some) ids already exist, only the new embeddings will be added.
130
131 Args:
132 ids: The ids to associate with the embeddings.
133 collection_id: The UUID of the collection to add the embeddings to.
134 embedding: The sequence of embeddings to add.
135 metadata: The metadata to associate with the embeddings. Defaults to None.
136 documents: The documents to associate with the embeddings. Defaults to None.
137 uris: URIs of data sources for each embedding. Defaults to None.
138
139 Returns:
140 True if the embeddings were added successfully.
141 """
142 pass
143
144 @abstractmethod
145 async def _update(
146 self,
147 collection_id: UUID,
148 ids: IDs,
149 embeddings: Optional[Embeddings] = None,
150 metadatas: Optional[Metadatas] = None,
151 documents: Optional[Documents] = None,
152 uris: Optional[URIs] = None,
153 ) -> bool:
154 """[Internal] Update entries in a collection specified by UUID.
155
156 Args:
157 collection_id: The UUID of the collection to update the embeddings in.
158 ids: The IDs of the entries to update.
159 embeddings: The sequence of embeddings to update. Defaults to None.
160 metadatas: The metadata to associate with the embeddings. Defaults to None.
161 documents: The documents to associate with the embeddings. Defaults to None.
162 uris: URIs of data sources for each embedding. Defaults to None.
163 Returns:
164 True if the embeddings were updated successfully.
165 """
166 pass
167
168 @abstractmethod
169 async def _upsert(
170 self,
171 collection_id: UUID,
172 ids: IDs,
173 embeddings: Embeddings,
174 metadatas: Optional[Metadatas] = None,
175 documents: Optional[Documents] = None,
176 uris: Optional[URIs] = None,
177 ) -> bool:
178 """[Internal] Add or update entries in the a collection specified by UUID.
179 If an entry with the same id already exists, it will be updated,
180 otherwise it will be added.
181
182 Args:
183 collection_id: The collection to add the embeddings to
184 ids: The ids to associate with the embeddings. Defaults to None.
185 embeddings: The sequence of embeddings to add
186 metadatas: The metadata to associate with the embeddings. Defaults to None.
187 documents: The documents to associate with the embeddings. Defaults to None.
188 uris: URIs of data sources for each embedding. Defaults to None.
189 """
190 pass
191
192 @abstractmethod
193 async def _count(self, collection_id: UUID) -> int:
194 """[Internal] Returns the number of entries in a collection specified by UUID.
195
196 Args:
197 collection_id: The UUID of the collection to count the embeddings in.
198
199 Returns:
200 int: The number of embeddings in the collection
201
202 """
203 pass
204
205 @abstractmethod
206 async def _peek(self, collection_id: UUID, n: int = 10) -> GetResult:
207 """[Internal] Returns the first n entries in a collection specified by UUID.
208
209 Args:
210 collection_id: The UUID of the collection to peek into.
211 n: The number of entries to peek. Defaults to 10.
212
213 Returns:
214 GetResult: The first n entries in the collection.
215
216 """
217
218 pass
219
220 @abstractmethod
221 async def _get(
222 self,
223 collection_id: UUID,
224 ids: Optional[IDs] = None,
225 where: Optional[Where] = None,
226 limit: Optional[int] = None,
227 offset: Optional[int] = None,
228 where_document: Optional[WhereDocument] = None,
229 include: Include = IncludeMetadataDocuments,
230 ) -> GetResult:
231 """[Internal] Returns entries from a collection specified by UUID.
232
233 Args:
234 ids: The IDs of the entries to get. Defaults to None.
235 where: Conditional filtering on metadata. Defaults to None.
236 limit: The maximum number of entries to return. Defaults to None.
237 offset: The number of entries to skip before returning. Defaults to None.
238 where_document: Conditional filtering on documents. Defaults to None.
239 include: The fields to include in the response.
240 Defaults to ["embeddings", "metadatas", "documents"].
241 Returns:
242 GetResult: The entries in the collection that match the query.
243
244 """
245 pass
246
247 @abstractmethod
248 async def _delete(
249 self,
250 collection_id: UUID,
251 ids: Optional[IDs],
252 where: Optional[Where] = None,
253 where_document: Optional[WhereDocument] = None,
254 limit: Optional[int] = None,
255 ) -> DeleteResult:
256 """[Internal] Deletes entries from a collection specified by UUID.
257
258 Args:
259 collection_id: The UUID of the collection to delete the entries from.
260 ids: The IDs of the entries to delete. Defaults to None.
261 where: Conditional filtering on metadata. Defaults to None.
262 where_document: Conditional filtering on documents. Defaults to None.
263 limit: Maximum number of records to delete. Can only be used with
264 where or where_document. Defaults to None (no limit).
265
266 Returns:
267 DeleteResult: A dict containing the number of records deleted.
268 """
269 pass
270
271 @abstractmethod
272 async def _query(
273 self,
274 collection_id: UUID,
275 query_embeddings: Embeddings,
276 ids: Optional[IDs] = None,
277 n_results: int = 10,
278 where: Optional[Where] = None,
279 where_document: Optional[WhereDocument] = None,
280 include: Include = IncludeMetadataDocumentsDistances,
281 ) -> QueryResult:
282 """[Internal] Performs a nearest neighbors query on a collection specified by UUID.
283
284 Args:
285 collection_id: The UUID of the collection to query.
286 query_embeddings: The embeddings to use as the query.
287 n_results: The number of results to return. Defaults to 10.
288 where: Conditional filtering on metadata. Defaults to None.
289 where_document: Conditional filtering on documents. Defaults to None.
290 include: The fields to include in the response.
291 Defaults to ["embeddings", "metadatas", "documents", "distances"].
292
293 Returns:
294 QueryResult: The results of the query.
295 """
296 pass
297
298 @abstractmethod
299 async def reset(self) -> bool:
300 """Resets the database. This will delete all collections and entries.
301
302 Returns:
303 bool: True if the database was reset successfully.
304 """
305 pass
306
307 @abstractmethod
308 async def get_version(self) -> str:
309 """Get the version of Chroma.
310
311 Returns:
312 str: The version of Chroma
313
314 """
315 pass
316
317 @abstractmethod
318 def get_settings(self) -> Settings:
319 """Get the settings used to initialize.
320
321 Returns:
322 Settings: The settings used to initialize.
323
324 """
325 pass
326
327 @abstractmethod
328 async def get_max_batch_size(self) -> int:
329 """Return the maximum number of records that can be created or mutated in a single call."""
330 pass
331
332 @abstractmethod
333 async def get_user_identity(self) -> UserIdentity:
334 """Resolve the tenant and databases for the client. Returns the default
335 values if can't be resolved.
336
337 """
338 pass
339
340
341class AsyncClientAPI(AsyncBaseAPI, ABC):
342 tenant: str
343 database: str
344
345 @abstractmethod
346 async def list_collections(
347 self,
348 limit: Optional[int] = None,
349 offset: Optional[int] = None,
350 ) -> Sequence[AsyncCollection]:
351 """List all collections.
352 Args:
353 limit: The maximum number of entries to return. Defaults to None.
354 offset: The number of entries to skip before returning. Defaults to None.
355
356 Returns:
357 Sequence[AsyncCollection]: A list of collections.
358
359 Examples:
360 ```python
361 await client.list_collections()
362 # [collection(name="my_collection", metadata={})]
363 ```
364 """
365 pass
366
367 @abstractmethod
368 async def create_collection(
369 self,
370 name: str,
371 schema: Optional[Schema] = None,
372 configuration: Optional[CreateCollectionConfiguration] = None,
373 metadata: Optional[CollectionMetadata] = None,
374 embedding_function: Optional[
375 EmbeddingFunction[Embeddable]
376 ] = DefaultEmbeddingFunction(), # type: ignore
377 data_loader: Optional[DataLoader[Loadable]] = None,
378 get_or_create: bool = False,
379 ) -> AsyncCollection:
380 """Create a new collection with the given name and metadata.
381 Args:
382 name: The name of the collection to create.
383 metadata: Optional metadata to associate with the collection.
384 embedding_function: Optional function to use to embed documents.
385 Uses the default embedding function if not provided.
386 get_or_create: If True, return the existing collection if it exists.
387 data_loader: Optional function to use to load records (documents, images, etc.)
388
389 Returns:
390 Collection: The newly created collection.
391
392 Raises:
393 ValueError: If the collection already exists and get_or_create is False.
394 ValueError: If the collection name is invalid.
395
396 Examples:
397 ```python
398 await client.create_collection("my_collection")
399 # collection(name="my_collection", metadata={})
400
401 await client.create_collection("my_collection", metadata={"foo": "bar"})
402 # collection(name="my_collection", metadata={"foo": "bar"})
403 ```
404 """
405 pass
406
407 @abstractmethod
408 async def get_collection(
409 self,
410 name: str,
411 embedding_function: Optional[
412 EmbeddingFunction[Embeddable]
413 ] = DefaultEmbeddingFunction(), # type: ignore
414 data_loader: Optional[DataLoader[Loadable]] = None,
415 ) -> AsyncCollection:
416 """Get a collection with the given name.
417 Args:
418 name: The name of the collection to get
419 embedding_function: Optional function to use to embed documents.
420 Uses the default embedding function if not provided.
421 data_loader: Optional function to use to load records (documents, images, etc.)
422
423 Returns:
424 Collection: The collection
425
426 Raises:
427 ValueError: If the collection does not exist
428
429 Examples:
430 ```python
431 await client.get_collection("my_collection")
432 # collection(name="my_collection", metadata={})
433 ```
434 """
435 pass
436
437 @abstractmethod
438 async def get_collection_by_id(
439 self,
440 id: UUID,
441 embedding_function: Optional[
442 EmbeddingFunction[Embeddable]
443 ] = DefaultEmbeddingFunction(), # type: ignore
444 data_loader: Optional[DataLoader[Loadable]] = None,
445 ) -> AsyncCollection:
446 """Get a collection by its ID.
447
448 Args:
449 id: The UUID of the collection to get.
450 embedding_function: Optional function to use to embed documents.
451 Uses the default embedding function if not provided.
452 data_loader: Optional function to use to load records (documents, images, etc.)
453
454 Returns:
455 Collection: The collection
456
457 Raises:
458 NotFoundError: If no collection with the given ID exists.
459
460 Examples:
461 ```python
462 await client.get_collection_by_id(uuid.UUID("..."))
463 # collection(name="my_collection", metadata={})
464 ```
465 """
466 pass
467
468 @abstractmethod
469 async def get_or_create_collection(
470 self,
471 name: str,
472 schema: Optional[Schema] = None,
473 configuration: Optional[CreateCollectionConfiguration] = None,
474 metadata: Optional[CollectionMetadata] = None,
475 embedding_function: Optional[
476 EmbeddingFunction[Embeddable]
477 ] = DefaultEmbeddingFunction(), # type: ignore
478 data_loader: Optional[DataLoader[Loadable]] = None,
479 ) -> AsyncCollection:
480 """Get or create a collection with the given name and metadata.
481 Args:
482 name: The name of the collection to get or create
483 metadata: Optional metadata to associate with the collection. If
484 the collection already exists, the metadata provided is ignored.
485 If the collection does not exist, the new collection will be created
486 with the provided metadata.
487 embedding_function: Optional function to use to embed documents
488 data_loader: Optional function to use to load records (documents, images, etc.)
489
490 Returns:
491 The collection
492
493 Examples:
494 ```python
495 await client.get_or_create_collection("my_collection")
496 # collection(name="my_collection", metadata={})
497 ```
498 """
499 pass
500
501 @abstractmethod
502 async def set_tenant(self, tenant: str, database: str = DEFAULT_DATABASE) -> None:
503 """Set the tenant and database for the client. Raises an error if the tenant or
504 database does not exist.
505
506 Args:
507 tenant: The tenant to set.
508 database: The database to set.
509
510 """
511 pass
512
513 @abstractmethod
514 async def set_database(self, database: str) -> None:
515 """Set the database for the client. Raises an error if the database does not exist.
516
517 Args:
518 database: The database to set.
519
520 """
521 pass
522
523 @staticmethod
524 @abstractmethod
525 def clear_system_cache() -> None:
526 """Clear the system cache so that new systems can be created for an existing path.
527 This should only be used for testing purposes."""
528 pass
529
530
531class AsyncAdminAPI(ABC):
532 @abstractmethod
533 async def create_database(self, name: str, tenant: str = DEFAULT_TENANT) -> None:
534 """Create a new database. Raises an error if the database already exists.
535
536 Args:
537 database: The name of the database to create.
538
539 """
540 pass
541
542 @abstractmethod
543 async def get_database(self, name: str, tenant: str = DEFAULT_TENANT) -> Database:
544 """Get a database. Raises an error if the database does not exist.
545
546 Args:
547 database: The name of the database to get.
548 tenant: The tenant of the database to get.
549
550 """
551 pass
552
553 @abstractmethod
554 async def delete_database(self, name: str, tenant: str = DEFAULT_TENANT) -> None:
555 """Delete a database. Raises an error if the database does not exist.
556
557 Args:
558 database: The name of the database to delete.
559 tenant: The tenant of the database to delete.
560
561 """
562 pass
563
564 @abstractmethod
565 async def list_databases(
566 self,
567 limit: Optional[int] = None,
568 offset: Optional[int] = None,
569 tenant: str = DEFAULT_TENANT,
570 ) -> Sequence[Database]:
571 """List all databases for a tenant. Raises an error if the tenant does not exist.
572
573 Args:
574 tenant: The tenant to list databases for.
575
576 """
577 pass
578
579 @abstractmethod
580 async def create_tenant(self, name: str) -> None:
581 """Create a new tenant. Raises an error if the tenant already exists.
582
583 Args:
584 tenant: The name of the tenant to create.
585
586 """
587 pass
588
589 @abstractmethod
590 async def get_tenant(self, name: str) -> Tenant:
591 """Get a tenant. Raises an error if the tenant does not exist.
592
593 Args:
594 tenant: The name of the tenant to get.
595
596 """
597 pass
598
599
600class AsyncServerAPI(AsyncBaseAPI, AsyncAdminAPI, Component):
601 """An API instance that extends the relevant Base API methods by passing
602 in a tenant and database. This is the root component of the Chroma System"""
603
604 @abstractmethod
605 async def list_collections(
606 self,
607 limit: Optional[int] = None,
608 offset: Optional[int] = None,
609 tenant: str = DEFAULT_TENANT,
610 database: str = DEFAULT_DATABASE,
611 ) -> Sequence[CollectionModel]:
612 pass
613
614 @abstractmethod
615 @override
616 async def count_collections(
617 self, tenant: str = DEFAULT_TENANT, database: str = DEFAULT_DATABASE
618 ) -> int:
619 pass
620
621 @abstractmethod
622 async def create_collection(
623 self,
624 name: str,
625 schema: Optional[Schema] = None,
626 configuration: Optional[CreateCollectionConfiguration] = None,
627 metadata: Optional[CollectionMetadata] = None,
628 get_or_create: bool = False,
629 tenant: str = DEFAULT_TENANT,
630 database: str = DEFAULT_DATABASE,
631 ) -> CollectionModel:
632 pass
633
634 @abstractmethod
635 async def get_collection(
636 self,
637 name: str,
638 tenant: str = DEFAULT_TENANT,
639 database: str = DEFAULT_DATABASE,
640 ) -> CollectionModel:
641 pass
642
643 @abstractmethod
644 async def get_collection_by_id(
645 self,
646 collection_id: UUID,
647 tenant: str = DEFAULT_TENANT,
648 database: str = DEFAULT_DATABASE,
649 ) -> CollectionModel:
650 """Get a collection by its ID.
651
652 Args:
653 collection_id: The UUID of the collection to retrieve.
654 tenant: The tenant to search within.
655 database: The database to search within.
656
657 Returns:
658 CollectionModel: The collection with the given ID.
659
660 Raises:
661 NotFoundError: If no collection with the given ID exists.
662 """
663 pass
664
665 @abstractmethod
666 async def get_or_create_collection(
667 self,
668 name: str,
669 schema: Optional[Schema] = None,
670 configuration: Optional[CreateCollectionConfiguration] = None,
671 metadata: Optional[CollectionMetadata] = None,
672 tenant: str = DEFAULT_TENANT,
673 database: str = DEFAULT_DATABASE,
674 ) -> CollectionModel:
675 pass
676
677 @abstractmethod
678 @override
679 async def delete_collection(
680 self,
681 name: str,
682 tenant: str = DEFAULT_TENANT,
683 database: str = DEFAULT_DATABASE,
684 ) -> None:
685 pass
686
687 @abstractmethod
688 @override
689 async def _modify(
690 self,
691 id: UUID,
692 new_name: Optional[str] = None,
693 new_metadata: Optional[CollectionMetadata] = None,
694 new_configuration: Optional[UpdateCollectionConfiguration] = None,
695 tenant: str = DEFAULT_TENANT,
696 database: str = DEFAULT_DATABASE,
697 ) -> None:
698 pass
699
700 @abstractmethod
701 async def _fork(
702 self,
703 collection_id: UUID,
704 new_name: str,
705 tenant: str = DEFAULT_TENANT,
706 database: str = DEFAULT_DATABASE,
707 ) -> CollectionModel:
708 pass
709
710 @abstractmethod
711 async def _fork_count(
712 self,
713 collection_id: UUID,
714 tenant: str = DEFAULT_TENANT,
715 database: str = DEFAULT_DATABASE,
716 ) -> int:
717 pass
718
719 @abstractmethod
720 async def _get_indexing_status(
721 self,
722 collection_id: UUID,
723 tenant: str = DEFAULT_TENANT,
724 database: str = DEFAULT_DATABASE,
725 ) -> "IndexingStatus":
726 pass
727
728 @abstractmethod
729 async def _search(
730 self,
731 collection_id: UUID,
732 searches: List[Search],
733 tenant: str = DEFAULT_TENANT,
734 database: str = DEFAULT_DATABASE,
735 read_level: ReadLevel = ReadLevel.INDEX_AND_WAL,
736 ) -> SearchResult:
737 pass
738
739 @abstractmethod
740 @override
741 async def _count(
742 self,
743 collection_id: UUID,
744 tenant: str = DEFAULT_TENANT,
745 database: str = DEFAULT_DATABASE,
746 read_level: ReadLevel = ReadLevel.INDEX_AND_WAL,
747 ) -> int:
748 pass
749
750 @abstractmethod
751 @override
752 async def _peek(
753 self,
754 collection_id: UUID,
755 n: int = 10,
756 tenant: str = DEFAULT_TENANT,
757 database: str = DEFAULT_DATABASE,
758 ) -> GetResult:
759 pass
760
761 @abstractmethod
762 @override
763 async def _get(
764 self,
765 collection_id: UUID,
766 ids: Optional[IDs] = None,
767 where: Optional[Where] = None,
768 limit: Optional[int] = None,
769 offset: Optional[int] = None,
770 where_document: Optional[WhereDocument] = None,
771 include: Include = IncludeMetadataDocuments,
772 tenant: str = DEFAULT_TENANT,
773 database: str = DEFAULT_DATABASE,
774 ) -> GetResult:
775 pass
776
777 @abstractmethod
778 @override
779 async def _add(
780 self,
781 ids: IDs,
782 collection_id: UUID,
783 embeddings: Embeddings,
784 metadatas: Optional[Metadatas] = None,
785 documents: Optional[Documents] = None,
786 uris: Optional[URIs] = None,
787 tenant: str = DEFAULT_TENANT,
788 database: str = DEFAULT_DATABASE,
789 ) -> bool:
790 pass
791
792 @abstractmethod
793 @override
794 async def _update(
795 self,
796 collection_id: UUID,
797 ids: IDs,
798 embeddings: Optional[Embeddings] = None,
799 metadatas: Optional[Metadatas] = None,
800 documents: Optional[Documents] = None,
801 uris: Optional[URIs] = None,
802 tenant: str = DEFAULT_TENANT,
803 database: str = DEFAULT_DATABASE,
804 ) -> bool:
805 pass
806
807 @abstractmethod
808 @override
809 async def _upsert(
810 self,
811 collection_id: UUID,
812 ids: IDs,
813 embeddings: Embeddings,
814 metadatas: Optional[Metadatas] = None,
815 documents: Optional[Documents] = None,
816 uris: Optional[URIs] = None,
817 tenant: str = DEFAULT_TENANT,
818 database: str = DEFAULT_DATABASE,
819 ) -> bool:
820 pass
821
822 @abstractmethod
823 @override
824 async def _query(
825 self,
826 collection_id: UUID,
827 query_embeddings: Embeddings,
828 ids: Optional[IDs] = None,
829 n_results: int = 10,
830 where: Optional[Where] = None,
831 where_document: Optional[WhereDocument] = None,
832 include: Include = IncludeMetadataDocumentsDistances,
833 tenant: str = DEFAULT_TENANT,
834 database: str = DEFAULT_DATABASE,
835 ) -> QueryResult:
836 pass
837
838 @abstractmethod
839 @override
840 async def _delete(
841 self,
842 collection_id: UUID,
843 ids: Optional[IDs] = None,
844 where: Optional[Where] = None,
845 where_document: Optional[WhereDocument] = None,
846 limit: Optional[int] = None,
847 tenant: str = DEFAULT_TENANT,
848 database: str = DEFAULT_DATABASE,
849 ) -> DeleteResult:
850 pass
851