Spaces:
Runtime error
Runtime error
File size: 3,059 Bytes
4a51346 |
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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
from typing import Optional, Sequence, Set, TypeVar, Type
from abc import ABC, abstractmethod
from chromadb.types import (
Collection,
MetadataEmbeddingRecord,
VectorEmbeddingRecord,
Where,
WhereDocument,
VectorQuery,
VectorQueryResult,
Segment,
SeqId,
)
from chromadb.config import Component, System
from overrides import EnforceOverrides
from uuid import UUID
class SegmentImplementation(ABC, EnforceOverrides):
@abstractmethod
def __init__(self, sytstem: System, segment: Segment):
pass
@abstractmethod
def count(self) -> int:
"""Get the number of embeddings in this segment"""
pass
@abstractmethod
def max_seqid(self) -> SeqId:
"""Get the maximum SeqID currently indexed by this segment"""
pass
class MetadataReader(SegmentImplementation):
"""Embedding Metadata segment interface"""
@abstractmethod
def get_metadata(
self,
where: Optional[Where] = None,
where_document: Optional[WhereDocument] = None,
ids: Optional[Sequence[str]] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
) -> Sequence[MetadataEmbeddingRecord]:
"""Query for embedding metadata."""
pass
class VectorReader(SegmentImplementation):
"""Embedding Vector segment interface"""
@abstractmethod
def get_vectors(
self, ids: Optional[Sequence[str]] = None
) -> Sequence[VectorEmbeddingRecord]:
"""Get embeddings from the segment. If no IDs are provided, all embeddings are
returned."""
pass
@abstractmethod
def query_vectors(
self, query: VectorQuery
) -> Sequence[Sequence[VectorQueryResult]]:
"""Given a vector query, return the top-k nearest neighbors for vector in the
query."""
pass
class SegmentManager(Component):
"""Interface for a pluggable strategy for creating, retrieving and instantiating
segments as required"""
@abstractmethod
def create_segments(self, collection: Collection) -> Set[Segment]:
"""Create the segments required for a new collection."""
pass
@abstractmethod
def delete_segments(self, collection_id: UUID) -> None:
"""Delete all the segments associated with a collection"""
pass
T = TypeVar("T", bound="SegmentImplementation")
# Future Note: To support time travel, add optional parameters to this method to
# retrieve Segment instances that are bounded to events from a specific range of
# time
@abstractmethod
def get_segment(self, collection_id: UUID, type: Type[T]) -> SegmentImplementation:
"""Return the segment that should be used for servicing queries to a collection.
Implementations should cache appropriately; clients are intended to call this
method repeatedly rather than storing the result (thereby giving this
implementation full control over which segment impls are in or out of memory at
a given time.)"""
pass
|