import httpx
from typing import Optional, Sequence
from uuid import UUID
from overrides import override
from chromadb.api.models.Collection import CollectionName
from chromadb.auth import UserIdentity
from chromadb.auth.utils import maybe_set_tenant_and_database
from chromadb.api import AsyncAdminAPI, AsyncClientAPI, AsyncServerAPI
from chromadb.api.configuration import CollectionConfiguration
from chromadb.api.models.AsyncCollection import AsyncCollection
from chromadb.api.shared_system_client import SharedSystemClient
from chromadb.api.types import (
CollectionMetadata,
DataLoader,
Documents,
Embeddable,
EmbeddingFunction,
Embeddings,
GetResult,
IDs,
Include,
Loadable,
Metadatas,
QueryResult,
URIs,
)
from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Settings, System
from chromadb.errors import ChromaError
from chromadb.types import Database, Tenant, Where, WhereDocument
import chromadb.utils.embedding_functions as ef
class AsyncClient(SharedSystemClient, AsyncClientAPI):
"""A client for Chroma. This is the main entrypoint for interacting with Chroma.
A client internally stores its tenant and database and proxies calls to a
Server API instance of Chroma. It treats the Server API and corresponding System
as a singleton, so multiple clients connecting to the same resource will share the
same API instance.
Client implementations should be implement their own API-caching strategies.
"""
# An internal admin client for verifying that databases and tenants exist
_admin_client: AsyncAdminAPI
tenant: str = DEFAULT_TENANT
database: str = DEFAULT_DATABASE
_server: AsyncServerAPI
@classmethod
async def create(
cls,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
settings: Settings = Settings(),
) -> "AsyncClient":
# Create an admin client for verifying that databases and tenants exist
self = cls(settings=settings)
SharedSystemClient._populate_data_from_system(self._system)
self.tenant = tenant
self.database = database
# Get the root system component we want to interact with
self._server = self._system.instance(AsyncServerAPI)
user_identity = await self.get_user_identity()
maybe_tenant, maybe_database = maybe_set_tenant_and_database(
user_identity,
overwrite_singleton_tenant_database_access_from_auth=settings.chroma_overwrite_singleton_tenant_database_access_from_auth,
user_provided_tenant=tenant,
user_provided_database=database,
)
if maybe_tenant:
self.tenant = maybe_tenant
if maybe_database:
self.database = maybe_database
self._admin_client = AsyncAdminClient.from_system(self._system)
await self._validate_tenant_database(tenant=self.tenant, database=self.database)
self._submit_client_start_event()
return self
@classmethod
# (we can't override and use from_system() because it's synchronous)
async def from_system_async(
cls,
system: System,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> "AsyncClient":
"""Create a client from an existing system. This is useful for testing and debugging."""
return await AsyncClient.create(tenant, database, system.settings)
@classmethod
@override
def from_system(
cls,
system: System,
) -> "SharedSystemClient":
"""AsyncClient cannot be created synchronously. Use .from_system_async() instead."""
raise NotImplementedError(
"AsyncClient cannot be created synchronously. Use .from_system_async() instead."
)
@override
async def get_user_identity(self) -> UserIdentity:
return await self._server.get_user_identity()
@override
async def set_tenant(self, tenant: str, database: str = DEFAULT_DATABASE) -> None:
await self._validate_tenant_database(tenant=tenant, database=database)
self.tenant = tenant
self.database = database
@override
async def set_database(self, database: str) -> None:
await self._validate_tenant_database(tenant=self.tenant, database=database)
self.database = database
async def _validate_tenant_database(self, tenant: str, database: str) -> None:
try:
await self._admin_client.get_tenant(name=tenant)
except httpx.ConnectError:
raise ValueError(
"Could not connect to a Chroma server. Are you sure it is running?"
)
# Propagate ChromaErrors
except ChromaError as e:
raise e
except Exception:
raise ValueError(
f"Could not connect to tenant {tenant}. Are you sure it exists?"
)
try:
await self._admin_client.get_database(name=database, tenant=tenant)
except httpx.ConnectError:
raise ValueError(
"Could not connect to a Chroma server. Are you sure it is running?"
)
# region BaseAPI Methods
# Note - we could do this in less verbose ways, but they break type checking
@override
async def heartbeat(self) -> int:
return await self._server.heartbeat()
@override
async def list_collections(
self, limit: Optional[int] = None, offset: Optional[int] = None
) -> Sequence[CollectionName]:
models = await self._server.list_collections(
limit, offset, tenant=self.tenant, database=self.database
)
return [CollectionName(model.name) for model in models]
@override
async def count_collections(self) -> int:
return await self._server.count_collections(
tenant=self.tenant, database=self.database
)
@override
async def create_collection(
self,
name: str,
configuration: Optional[CollectionConfiguration] = None,
metadata: Optional[CollectionMetadata] = None,
embedding_function: Optional[
EmbeddingFunction[Embeddable]
] = ef.DefaultEmbeddingFunction(), # type: ignore
data_loader: Optional[DataLoader[Loadable]] = None,
get_or_create: bool = False,
) -> AsyncCollection:
model = await self._server.create_collection(
name=name,
configuration=configuration,
metadata=metadata,
tenant=self.tenant,
database=self.database,
get_or_create=get_or_create,
)
return AsyncCollection(
client=self._server,
model=model,
embedding_function=embedding_function,
data_loader=data_loader,
)
@override
async def get_collection(
self,
name: str,
embedding_function: Optional[
EmbeddingFunction[Embeddable]
] = ef.DefaultEmbeddingFunction(), # type: ignore
data_loader: Optional[DataLoader[Loadable]] = None,
) -> AsyncCollection:
model = await self._server.get_collection(
name=name,
tenant=self.tenant,
database=self.database,
)
return AsyncCollection(
client=self._server,
model=model,
embedding_function=embedding_function,
data_loader=data_loader,
)
@override
async def get_or_create_collection(
self,
name: str,
configuration: Optional[CollectionConfiguration] = None,
metadata: Optional[CollectionMetadata] = None,
embedding_function: Optional[
EmbeddingFunction[Embeddable]
] = ef.DefaultEmbeddingFunction(), # type: ignore
data_loader: Optional[DataLoader[Loadable]] = None,
) -> AsyncCollection:
model = await self._server.get_or_create_collection(
name=name,
configuration=configuration,
metadata=metadata,
tenant=self.tenant,
database=self.database,
)
return AsyncCollection(
client=self._server,
model=model,
embedding_function=embedding_function,
data_loader=data_loader,
)
@override
async def _modify(
self,
id: UUID,
new_name: Optional[str] = None,
new_metadata: Optional[CollectionMetadata] = None,
) -> None:
return await self._server._modify(
id=id,
new_name=new_name,
new_metadata=new_metadata,
tenant=self.tenant,
database=self.database,
)
@override
async def delete_collection(
self,
name: str,
) -> None:
return await self._server.delete_collection(
name=name,
tenant=self.tenant,
database=self.database,
)
#
# ITEM METHODS
#
@override
async def _add(
self,
ids: IDs,
collection_id: UUID,
embeddings: Embeddings,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
uris: Optional[URIs] = None,
) -> bool:
return await self._server._add(
ids=ids,
collection_id=collection_id,
embeddings=embeddings,
metadatas=metadatas,
documents=documents,
uris=uris,
tenant=self.tenant,
database=self.database,
)
@override
async def _update(
self,
collection_id: UUID,
ids: IDs,
embeddings: Optional[Embeddings] = None,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
uris: Optional[URIs] = None,
) -> bool:
return await self._server._update(
collection_id=collection_id,
ids=ids,
embeddings=embeddings,
metadatas=metadatas,
documents=documents,
uris=uris,
tenant=self.tenant,
database=self.database,
)
@override
async def _upsert(
self,
collection_id: UUID,
ids: IDs,
embeddings: Embeddings,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
uris: Optional[URIs] = None,
) -> bool:
return await self._server._upsert(
collection_id=collection_id,
ids=ids,
embeddings=embeddings,
metadatas=metadatas,
documents=documents,
uris=uris,
tenant=self.tenant,
database=self.database,
)
@override
async def _count(self, collection_id: UUID) -> int:
return await self._server._count(
collection_id=collection_id,
)
@override
async def _peek(self, collection_id: UUID, n: int = 10) -> GetResult:
return await self._server._peek(
collection_id=collection_id,
n=n,
)
@override
async def _get(
self,
collection_id: UUID,
ids: Optional[IDs] = None,
where: Optional[Where] = None,
sort: Optional[str] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
page: Optional[int] = None,
page_size: Optional[int] = None,
where_document: Optional[WhereDocument] = None,
include: Include = ["embeddings", "metadatas", "documents"], # type: ignore[list-item]
) -> GetResult:
return await self._server._get(
collection_id=collection_id,
ids=ids,
where=where,
sort=sort,
limit=limit,
offset=offset,
page=page,
page_size=page_size,
where_document=where_document,
include=include,
tenant=self.tenant,
database=self.database,
)
async def _delete(
self,
collection_id: UUID,
ids: Optional[IDs],
where: Optional[Where] = None,
where_document: Optional[WhereDocument] = None,
) -> None:
await self._server._delete(
collection_id=collection_id,
ids=ids,
where=where,
where_document=where_document,
tenant=self.tenant,
database=self.database,
)
@override
async def _query(
self,
collection_id: UUID,
query_embeddings: Embeddings,
n_results: int = 10,
where: Optional[Where] = None,
where_document: Optional[WhereDocument] = None,
include: Include = ["embeddings", "metadatas", "documents", "distances"], # type: ignore[list-item]
) -> QueryResult:
return await self._server._query(
collection_id=collection_id,
query_embeddings=query_embeddings,
n_results=n_results,
where=where,
where_document=where_document,
include=include,
tenant=self.tenant,
database=self.database,
)
@override
async def reset(self) -> bool:
return await self._server.reset()
@override
async def get_version(self) -> str:
return await self._server.get_version()
@override
def get_settings(self) -> Settings:
return self._server.get_settings()
@override
async def get_max_batch_size(self) -> int:
return await self._server.get_max_batch_size()
# endregion
class AsyncAdminClient(SharedSystemClient, AsyncAdminAPI):
_server: AsyncServerAPI
def __init__(self, settings: Settings = Settings()) -> None:
super().__init__(settings)
self._server = self._system.instance(AsyncServerAPI)
@override
async def create_database(self, name: str, tenant: str = DEFAULT_TENANT) -> None:
return await self._server.create_database(name=name, tenant=tenant)
@override
async def get_database(self, name: str, tenant: str = DEFAULT_TENANT) -> Database:
return await self._server.get_database(name=name, tenant=tenant)
@override
async def delete_database(self, name: str, tenant: str = DEFAULT_TENANT) -> None:
return await self._server.delete_database(name=name, tenant=tenant)
@override
async def list_databases(
self,
limit: Optional[int] = None,
offset: Optional[int] = None,
tenant: str = DEFAULT_TENANT,
) -> Sequence[Database]:
return await self._server.list_databases(
limit=limit, offset=offset, tenant=tenant
)
@override
async def create_tenant(self, name: str) -> None:
return await self._server.create_tenant(name=name)
@override
async def get_tenant(self, name: str) -> Tenant:
return await self._server.get_tenant(name=name)
@classmethod
@override
def from_system(
cls,
system: System,
) -> "AsyncAdminClient":
SharedSystemClient._populate_data_from_system(system)
instance = cls(settings=system.settings)
return instance