codekingpro/portable-devtools

1from typing import Optional, Sequence, TypeVar
2from abc import abstractmethod
3from chromadb.types import (
4    Collection,
5    MetadataEmbeddingRecord,
6    Operation,
7    RequestVersionContext,
8    VectorEmbeddingRecord,
9    Where,
10    WhereDocument,
11    VectorQuery,
12    VectorQueryResult,
13    Segment,
14    SeqId,
15    Metadata,
16)
17from chromadb.config import Component, System
18from uuid import UUID
19from enum import Enum
20
21
22class SegmentType(Enum):
23    SQLITE = "urn:chroma:segment/metadata/sqlite"
24    HNSW_LOCAL_MEMORY = "urn:chroma:segment/vector/hnsw-local-memory"
25    HNSW_LOCAL_PERSISTED = "urn:chroma:segment/vector/hnsw-local-persisted"
26    HNSW_DISTRIBUTED = "urn:chroma:segment/vector/hnsw-distributed"
27    BLOCKFILE_RECORD = "urn:chroma:segment/record/blockfile"
28    BLOCKFILE_METADATA = "urn:chroma:segment/metadata/blockfile"
29
30
31class SegmentImplementation(Component):
32    @abstractmethod
33    def __init__(self, sytstem: System, segment: Segment):
34        pass
35
36    @abstractmethod
37    def count(self, request_version_context: RequestVersionContext) -> int:
38        """Get the number of embeddings in this segment"""
39        pass
40
41    @abstractmethod
42    def max_seqid(self) -> SeqId:
43        """Get the maximum SeqID currently indexed by this segment"""
44        pass
45
46    @staticmethod
47    def propagate_collection_metadata(metadata: Metadata) -> Optional[Metadata]:
48        """Given an arbitrary metadata map (e.g, from a collection), validate it and
49        return metadata (if any) that is applicable and should be applied to the
50        segment. Validation errors will be reported to the user."""
51        return None
52
53    @abstractmethod
54    def delete(self) -> None:
55        """Delete the segment and all its data"""
56        ...
57
58
59S = TypeVar("S", bound=SegmentImplementation)
60
61
62class MetadataReader(SegmentImplementation):
63    """Embedding Metadata segment interface"""
64
65    @abstractmethod
66    def get_metadata(
67        self,
68        request_version_context: RequestVersionContext,
69        where: Optional[Where] = None,
70        where_document: Optional[WhereDocument] = None,
71        ids: Optional[Sequence[str]] = None,
72        limit: Optional[int] = None,
73        offset: Optional[int] = None,
74        include_metadata: bool = True,
75    ) -> Sequence[MetadataEmbeddingRecord]:
76        """Query for embedding metadata."""
77        pass
78
79
80class VectorReader(SegmentImplementation):
81    """Embedding Vector segment interface"""
82
83    @abstractmethod
84    def get_vectors(
85        self,
86        request_version_context: RequestVersionContext,
87        ids: Optional[Sequence[str]] = None,
88    ) -> Sequence[VectorEmbeddingRecord]:
89        """Get embeddings from the segment. If no IDs are provided, all embeddings are
90        returned."""
91        pass
92
93    @abstractmethod
94    def query_vectors(
95        self, query: VectorQuery
96    ) -> Sequence[Sequence[VectorQueryResult]]:
97        """Given a vector query, return the top-k nearest neighbors for vector in the
98        query."""
99        pass
100
101
102class SegmentManager(Component):
103    """Interface for a pluggable strategy for creating, retrieving and instantiating
104    segments as required"""
105
106    @abstractmethod
107    def prepare_segments_for_new_collection(
108        self, collection: Collection
109    ) -> Sequence[Segment]:
110        """Return the segments required for a new collection. Returns only segment data,
111        does not persist to the SysDB"""
112        pass
113
114    @abstractmethod
115    def delete_segments(self, collection_id: UUID) -> Sequence[UUID]:
116        """Delete any local state for all the segments associated with a collection, and
117        returns a sequence of their IDs. Does not update the SysDB."""
118        pass
119
120    @abstractmethod
121    def hint_use_collection(self, collection_id: UUID, hint_type: Operation) -> None:
122        """Signal to the segment manager that a collection is about to be used, so that
123        it can preload segments as needed. This is only a hint, and implementations are
124        free to ignore it."""
125        pass
126