File size: 2,612 Bytes
c88e992 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
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
|