from dataclasses import dataclass
from typing import Optional
from chromadb.api.types import Embeddings, IDs, Include, IncludeEnum
from chromadb.types import Collection, RequestVersionContext, Segment, Where, WhereDocument
@dataclass
class Scan:
collection: Collection
knn: Segment
metadata: Segment
record: Segment
@property
def version(self) -> RequestVersionContext:
return RequestVersionContext(
collection_version=self.collection.version,
log_position=self.collection.log_position,
)
@dataclass
class Filter:
user_ids: Optional[IDs] = None
where: Optional[Where] = None
where_document: Optional[WhereDocument] = None
@dataclass
class KNN:
embeddings: Embeddings
fetch: int
@dataclass
class Limit:
skip: int = 0
fetch: Optional[int] = None
@dataclass
class Projection:
document: bool = False
embedding: bool = False
metadata: bool = False
rank: bool = False
uri: bool = False
@property
def included(self) -> Include:
includes = list()
if self.document:
includes.append(IncludeEnum.documents)
if self.embedding:
includes.append(IncludeEnum.embeddings)
if self.metadata:
includes.append(IncludeEnum.metadatas)
if self.rank:
includes.append(IncludeEnum.distances)
if self.uri:
includes.append(IncludeEnum.uris)
return includes