Spaces:
Runtime error
Runtime error
File size: 3,390 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
from abc import abstractmethod
from typing import List, Sequence, Optional, Tuple
from uuid import UUID
import numpy.typing as npt
from chromadb.api.types import (
Embeddings,
Documents,
IDs,
Metadatas,
Metadata,
Where,
WhereDocument,
)
from chromadb.config import Component
from overrides import override
class DB(Component):
@abstractmethod
def create_collection(
self,
name: str,
metadata: Optional[Metadata] = None,
get_or_create: bool = False,
) -> Sequence: # type: ignore
pass
@abstractmethod
def get_collection(self, name: str) -> Sequence: # type: ignore
pass
@abstractmethod
def list_collections(self) -> Sequence: # type: ignore
pass
@abstractmethod
def update_collection(
self,
id: UUID,
new_name: Optional[str] = None,
new_metadata: Optional[Metadata] = None,
) -> None:
pass
@abstractmethod
def delete_collection(self, name: str) -> None:
pass
@abstractmethod
def get_collection_uuid_from_name(self, collection_name: str) -> UUID:
pass
@abstractmethod
def add(
self,
collection_uuid: UUID,
embeddings: Embeddings,
metadatas: Optional[Metadatas],
documents: Optional[Documents],
ids: List[str],
) -> List[UUID]:
pass
@abstractmethod
def add_incremental(
self, collection_uuid: UUID, ids: List[UUID], embeddings: Embeddings
) -> None:
pass
@abstractmethod
def get(
self,
where: Where = {},
collection_name: Optional[str] = None,
collection_uuid: Optional[UUID] = None,
ids: Optional[IDs] = None,
sort: Optional[str] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
where_document: WhereDocument = {},
columns: Optional[List[str]] = None,
) -> Sequence: # type: ignore
pass
@abstractmethod
def update(
self,
collection_uuid: UUID,
ids: IDs,
embeddings: Optional[Embeddings] = None,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
) -> bool:
pass
@abstractmethod
def count(self, collection_id: UUID) -> int:
pass
@abstractmethod
def delete(
self,
where: Where = {},
collection_uuid: Optional[UUID] = None,
ids: Optional[IDs] = None,
where_document: WhereDocument = {},
) -> List[str]:
pass
@abstractmethod
@override
def reset(self) -> None:
pass
@abstractmethod
def get_nearest_neighbors(
self,
collection_uuid: UUID,
where: Where = {},
embeddings: Optional[Embeddings] = None,
n_results: int = 10,
where_document: WhereDocument = {},
) -> Tuple[List[List[UUID]], npt.NDArray]:
pass
@abstractmethod
def get_by_ids(
self, uuids: List[UUID], columns: Optional[List[str]] = None
) -> Sequence: # type: ignore
pass
@abstractmethod
def raw_sql(self, raw_sql): # type: ignore
pass
@abstractmethod
def create_index(self, collection_uuid: UUID): # type: ignore
pass
@abstractmethod
def persist(self) -> None:
pass
|