codekingpro/portable-devtools

1from abc import abstractmethod
2from typing import Callable, Optional, Sequence
3from chromadb.types import (
4    OperationRecord,
5    LogRecord,
6    SeqId,
7    Vector,
8    ScalarEncoding,
9)
10from chromadb.config import Component
11from uuid import UUID
12import numpy as np
13
14
15def encode_vector(vector: Vector, encoding: ScalarEncoding) -> bytes:
16    """Encode a vector into a byte array."""
17
18    if encoding == ScalarEncoding.FLOAT32:
19        return np.array(vector, dtype=np.float32).tobytes()
20    elif encoding == ScalarEncoding.INT32:
21        return np.array(vector, dtype=np.int32).tobytes()
22    else:
23        raise ValueError(f"Unsupported encoding: {encoding.value}")
24
25
26def decode_vector(vector: bytes, encoding: ScalarEncoding) -> Vector:
27    """Decode a byte array into a vector"""
28
29    if encoding == ScalarEncoding.FLOAT32:
30        return np.frombuffer(vector, dtype=np.float32)
31    elif encoding == ScalarEncoding.INT32:
32        return np.frombuffer(vector, dtype=np.float32)
33    else:
34        raise ValueError(f"Unsupported encoding: {encoding.value}")
35
36
37class Producer(Component):
38    """Interface for writing embeddings to an ingest stream"""
39
40    @abstractmethod
41    def delete_log(self, collection_id: UUID) -> None:
42        pass
43
44    @abstractmethod
45    def purge_log(self, collection_id: UUID) -> None:
46        """Truncates the log for the given collection, removing all seen records."""
47        pass
48
49    @abstractmethod
50    def submit_embedding(
51        self, collection_id: UUID, embedding: OperationRecord
52    ) -> SeqId:
53        """Add an embedding record to the given collections log. Returns the SeqID of the record."""
54        pass
55
56    @abstractmethod
57    def submit_embeddings(
58        self, collection_id: UUID, embeddings: Sequence[OperationRecord]
59    ) -> Sequence[SeqId]:
60        """Add a batch of embedding records to the given collections log. Returns the SeqIDs of
61        the records. The returned SeqIDs will be in the same order as the given
62        SubmitEmbeddingRecords. However, it is not guaranteed that the SeqIDs will be
63        processed in the same order as the given SubmitEmbeddingRecords. If the number
64        of records exceeds the maximum batch size, an exception will be thrown."""
65        pass
66
67    @property
68    @abstractmethod
69    def max_batch_size(self) -> int:
70        """Return the maximum number of records that can be submitted in a single call
71        to submit_embeddings."""
72        pass
73
74
75ConsumerCallbackFn = Callable[[Sequence[LogRecord]], None]
76
77
78class Consumer(Component):
79    """Interface for reading embeddings off an ingest stream"""
80
81    @abstractmethod
82    def subscribe(
83        self,
84        collection_id: UUID,
85        consume_fn: ConsumerCallbackFn,
86        start: Optional[SeqId] = None,
87        end: Optional[SeqId] = None,
88        id: Optional[UUID] = None,
89    ) -> UUID:
90        """Register a function that will be called to receive embeddings for a given
91        collections log stream. The given function may be called any number of times, with any number of
92        records, and may be called concurrently.
93
94        Only records between start (exclusive) and end (inclusive) SeqIDs will be
95        returned. If start is None, the first record returned will be the next record
96        generated, not including those generated before creating the subscription. If
97        end is None, the consumer will consume indefinitely, otherwise it will
98        automatically be unsubscribed when the end SeqID is reached.
99
100        If the function throws an exception, the function may be called again with the
101        same or different records.
102
103        Takes an optional UUID as a unique subscription ID. If no ID is provided, a new
104        ID will be generated and returned."""
105        pass
106
107    @abstractmethod
108    def unsubscribe(self, subscription_id: UUID) -> None:
109        """Unregister a subscription. The consume function will no longer be invoked,
110        and resources associated with the subscription will be released."""
111        pass
112
113    @abstractmethod
114    def min_seqid(self) -> SeqId:
115        """Return the minimum possible SeqID in this implementation."""
116        pass
117
118    @abstractmethod
119    def max_seqid(self) -> SeqId:
120        """Return the maximum possible SeqID in this implementation."""
121        pass
122