types.py

from pydantic import BaseModel
from typing import Any, Dict, List, Optional
#from chromadb.api.types import (
#    CollectionMetadata,
#    Include,
#)
from typing import Union, Literal

#TODO:Update version-Step1::This file is a copy from chromadb/server/fastapi/types.py
# Copy CollectionMetadata and Include from chromadb.api.types and comment that import
# Additionally import Union and Literal from typing for "Include to work"
CollectionMetadata = Dict[Any, Any]

Include = List[
    Union[
        Literal["documents"],
        Literal["embeddings"],
        Literal["metadatas"],
        Literal["distances"],
    ]
]


class AddEmbedding(BaseModel):  # type: ignore
    # Pydantic doesn't handle Union types cleanly like Embeddings which has
    # Union[int, float] so we use Any here to ensure data is parsed
    # to its original type.
    embeddings: Optional[List[Any]] = None
    metadatas: Optional[List[Dict[Any, Any]]] = None
    documents: Optional[List[str]] = None
    ids: List[str]
    increment_index: bool = True


class UpdateEmbedding(BaseModel):  # type: ignore
    embeddings: Optional[List[Any]] = None
    metadatas: Optional[List[Dict[Any, Any]]] = None
    documents: Optional[List[str]] = None
    ids: List[str]
    increment_index: bool = True


class QueryEmbedding(BaseModel):  # type: ignore
    # TODO: Pydantic doesn't bode well with recursive types so we use generic Dicts
    # for Where and WhereDocument. This is not ideal, but it works for now since
    # there is a lot of downstream validation.
    where: Optional[Dict[Any, Any]] = {}
    where_document: Optional[Dict[Any, Any]] = {}
    query_embeddings: List[Any]
    n_results: int = 10
    include: Include = ["metadatas", "documents", "distances"]


class GetEmbedding(BaseModel):  # type: ignore
    ids: Optional[List[str]] = None
    where: Optional[Dict[Any, Any]] = None
    where_document: Optional[Dict[Any, Any]] = None
    sort: Optional[str] = None
    limit: Optional[int] = None
    offset: Optional[int] = None
    include: Include = ["metadatas", "documents"]


class RawSql(BaseModel):  # type: ignore
    raw_sql: str


class DeleteEmbedding(BaseModel):  # type: ignore
    ids: Optional[List[str]] = None
    where: Optional[Dict[Any, Any]] = None
    where_document: Optional[Dict[Any, Any]] = None


class CreateCollection(BaseModel):  # type: ignore
    name: str
    metadata: Optional[CollectionMetadata] = None
    get_or_create: bool = False


class UpdateCollection(BaseModel):  # type: ignore
    new_name: Optional[str] = None
    new_metadata: Optional[CollectionMetadata] = None