196 lines
5.6 KiB
Python
196 lines
5.6 KiB
Python
from typing import List, Optional, Sequence
|
|
from uuid import UUID
|
|
from chromadb import CollectionMetadata, Embeddings, IDs
|
|
from chromadb.api.types import (
|
|
CollectionMetadata,
|
|
Documents,
|
|
Embeddings,
|
|
IDs,
|
|
Metadatas,
|
|
URIs,
|
|
Include,
|
|
)
|
|
from chromadb.types import Tenant, Collection as CollectionModel
|
|
from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT
|
|
from enum import Enum
|
|
|
|
class DatabaseFromBindings:
|
|
id: UUID
|
|
name: str
|
|
tenant: str
|
|
|
|
# Result Types
|
|
|
|
class GetResponse:
|
|
ids: IDs
|
|
embeddings: Embeddings
|
|
documents: Documents
|
|
uris: URIs
|
|
metadatas: Metadatas
|
|
include: Include
|
|
|
|
class QueryResponse:
|
|
ids: List[IDs]
|
|
embeddings: Optional[List[Embeddings]]
|
|
documents: Optional[List[Documents]]
|
|
uris: Optional[List[URIs]]
|
|
metadatas: Optional[List[Metadatas]]
|
|
distances: Optional[List[List[float]]]
|
|
include: Include
|
|
|
|
class GetTenantResponse:
|
|
name: str
|
|
|
|
# SqliteDBConfig types
|
|
class MigrationMode(Enum):
|
|
Apply = 0
|
|
Validate = 1
|
|
|
|
class MigrationHash(Enum):
|
|
SHA256 = 0
|
|
MD5 = 1
|
|
|
|
class SqliteDBConfig:
|
|
url: str
|
|
hash_type: MigrationHash
|
|
migration_mode: MigrationMode
|
|
|
|
def __init__(
|
|
self, url: str, hash_type: MigrationHash, migration_mode: MigrationMode
|
|
) -> None: ...
|
|
|
|
class Bindings:
|
|
def __init__(
|
|
self,
|
|
allow_reset: bool,
|
|
sqlite_db_config: SqliteDBConfig,
|
|
persist_path: str,
|
|
hnsw_cache_size: int,
|
|
) -> None: ...
|
|
def heartbeat(self) -> int: ...
|
|
def create_database(self, name: str, tenant: str = DEFAULT_TENANT) -> None: ...
|
|
def get_database(
|
|
self, name: str, tenant: str = DEFAULT_TENANT
|
|
) -> DatabaseFromBindings: ...
|
|
def delete_database(self, name: str, tenant: str = DEFAULT_TENANT) -> None: ...
|
|
def list_databases(
|
|
self,
|
|
limit: Optional[int] = None,
|
|
offset: Optional[int] = None,
|
|
tenant: str = DEFAULT_TENANT,
|
|
) -> Sequence[DatabaseFromBindings]: ...
|
|
def create_tenant(self, name: str) -> None: ...
|
|
def get_tenant(self, name: str) -> GetTenantResponse: ...
|
|
def count_collections(
|
|
self, tenant: str = DEFAULT_TENANT, database: str = DEFAULT_DATABASE
|
|
) -> int: ...
|
|
def list_collections(
|
|
self,
|
|
limit: Optional[int] = None,
|
|
offset: Optional[int] = None,
|
|
tenant: str = DEFAULT_TENANT,
|
|
database: str = DEFAULT_DATABASE,
|
|
) -> Sequence[CollectionModel]: ...
|
|
def create_collection(
|
|
self,
|
|
name: str,
|
|
configuration_json_str: Optional[str] = None,
|
|
schema_str: Optional[str] = None,
|
|
metadata: Optional[CollectionMetadata] = None,
|
|
get_or_create: bool = False,
|
|
tenant: str = DEFAULT_TENANT,
|
|
database: str = DEFAULT_DATABASE,
|
|
) -> CollectionModel: ...
|
|
def get_collection(
|
|
self,
|
|
name: str,
|
|
tenant: str = DEFAULT_TENANT,
|
|
database: str = DEFAULT_DATABASE,
|
|
) -> CollectionModel: ...
|
|
def update_collection(
|
|
self,
|
|
id: str,
|
|
new_name: Optional[str] = None,
|
|
new_metadata: Optional[CollectionMetadata] = None,
|
|
new_configuration_json_str: Optional[str] = None,
|
|
) -> None: ...
|
|
def delete_collection(
|
|
self,
|
|
name: str,
|
|
tenant: str = DEFAULT_TENANT,
|
|
database: str = DEFAULT_DATABASE,
|
|
) -> None: ...
|
|
def add(
|
|
self,
|
|
ids: IDs,
|
|
collection_id: str,
|
|
embeddings: Embeddings,
|
|
metadatas: Optional[Metadatas] = None,
|
|
documents: Optional[Documents] = None,
|
|
uris: Optional[URIs] = None,
|
|
tenant: str = DEFAULT_TENANT,
|
|
database: str = DEFAULT_DATABASE,
|
|
) -> bool: ...
|
|
def update(
|
|
self,
|
|
collection_id: str,
|
|
ids: IDs,
|
|
embeddings: Optional[Embeddings] = None,
|
|
metadatas: Optional[Metadatas] = None,
|
|
documents: Optional[Documents] = None,
|
|
uris: Optional[URIs] = None,
|
|
tenant: str = DEFAULT_TENANT,
|
|
database: str = DEFAULT_DATABASE,
|
|
) -> bool: ...
|
|
def upsert(
|
|
self,
|
|
collection_id: str,
|
|
ids: IDs,
|
|
embeddings: Embeddings,
|
|
metadatas: Optional[Metadatas] = None,
|
|
documents: Optional[Documents] = None,
|
|
uris: Optional[URIs] = None,
|
|
tenant: str = DEFAULT_TENANT,
|
|
database: str = DEFAULT_DATABASE,
|
|
) -> bool: ...
|
|
def delete(
|
|
self,
|
|
collection_id: str,
|
|
ids: Optional[IDs] = None,
|
|
where: Optional[str] = None,
|
|
where_document: Optional[str] = None,
|
|
tenant: str = DEFAULT_TENANT,
|
|
database: str = DEFAULT_DATABASE,
|
|
) -> None: ...
|
|
def count(
|
|
self,
|
|
collection_id: str,
|
|
tenant: str = DEFAULT_TENANT,
|
|
database: str = DEFAULT_DATABASE,
|
|
) -> int: ...
|
|
def get(
|
|
self,
|
|
collection_id: str,
|
|
ids: Optional[IDs] = None,
|
|
where: Optional[str] = None,
|
|
limit: Optional[int] = None,
|
|
offset: Optional[int] = None,
|
|
where_document: Optional[str] = None,
|
|
include: Include = ["metadatas", "documents"], # type: ignore[list-item]
|
|
tenant: str = DEFAULT_TENANT,
|
|
database: str = DEFAULT_DATABASE,
|
|
) -> GetResponse: ...
|
|
def query(
|
|
self,
|
|
collection_id: str,
|
|
query_embeddings: Embeddings,
|
|
n_results: int = 10,
|
|
where: Optional[str] = None,
|
|
where_document: Optional[str] = None,
|
|
include: Include = ["metadatas", "documents", "distances"], # type: ignore[list-item]
|
|
tenant: str = DEFAULT_TENANT,
|
|
database: str = DEFAULT_DATABASE,
|
|
) -> QueryResponse: ...
|
|
def reset(self) -> bool: ...
|
|
def get_version(self) -> str: ...
|