Spaces:
Runtime error
Runtime error
File size: 10,898 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 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 |
from abc import ABC, abstractmethod
from typing import Sequence, Optional
import pandas as pd
from uuid import UUID
from chromadb.api.models.Collection import Collection
from chromadb.api.types import (
CollectionMetadata,
Documents,
EmbeddingFunction,
Embeddings,
IDs,
Include,
Metadatas,
Where,
QueryResult,
GetResult,
WhereDocument,
)
from chromadb.config import Component
import chromadb.utils.embedding_functions as ef
from overrides import override
class API(Component, ABC):
@abstractmethod
def heartbeat(self) -> int:
"""Returns the current server time in nanoseconds to check if the server is alive
Args:
None
Returns:
int: The current server time in nanoseconds
"""
pass
@abstractmethod
def list_collections(self) -> Sequence[Collection]:
"""Returns all collections in the database
Args:
None
Returns:
dict: A dictionary of collections
"""
pass
@abstractmethod
def create_collection(
self,
name: str,
metadata: Optional[CollectionMetadata] = None,
embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
get_or_create: bool = False,
) -> Collection:
"""Creates a new collection in the database
Args:
name The name of the collection to create. The name must be unique.
metadata: A dictionary of metadata to associate with the collection. Defaults to None.
embedding_function: A function that takes documents and returns an embedding. Defaults to None.
get_or_create: If True, will return the collection if it already exists,
and update the metadata (if applicable). Defaults to False.
Returns:
dict: the created collection
"""
pass
@abstractmethod
def delete_collection(
self,
name: str,
) -> None:
"""Deletes a collection from the database
Args:
name: The name of the collection to delete
"""
@abstractmethod
def get_or_create_collection(
self,
name: str,
metadata: Optional[CollectionMetadata] = None,
embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
) -> Collection:
"""Calls create_collection with get_or_create=True.
If the collection exists, but with different metadata, the metadata will be replaced.
Args:
name: The name of the collection to create. The name must be unique.
metadata: A dictionary of metadata to associate with the collection. Defaults to None.
embedding_function: A function that takes documents and returns an embedding. Should be the same as the one used to create the collection. Defaults to None.
Returns:
the created collection
"""
pass
@abstractmethod
def get_collection(
self,
name: str,
embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
) -> Collection:
"""Gets a collection from the database by either name or uuid
Args:
name: The name of the collection to get. Defaults to None.
embedding_function: A function that takes documents and returns an embedding. Should be the same as the one used to create the collection. Defaults to None.
Returns:
dict: the requested collection
"""
pass
def _modify(
self,
id: UUID,
new_name: Optional[str] = None,
new_metadata: Optional[CollectionMetadata] = None,
) -> None:
"""Modify a collection in the database - can update the name and/or metadata
Args:
current_name: The name of the collection to modify
new_name: The new name of the collection. Defaults to None.
new_metadata: The new metadata to associate with the collection. Defaults to None.
"""
pass
@abstractmethod
def _add(
self,
ids: IDs,
collection_id: UUID,
embeddings: Embeddings,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
increment_index: bool = True,
) -> bool:
"""Add embeddings to the data store. This is the most general way to add embeddings to the database.
⚠️ It is recommended to use the more specific methods below when possible.
Args:
collection_id: The collection to add the embeddings to
embedding: The sequence of embeddings to add
metadata: The metadata to associate with the embeddings. Defaults to None.
documents: The documents to associate with the embeddings. Defaults to None.
ids: The ids to associate with the embeddings. Defaults to None.
"""
pass
@abstractmethod
def _update(
self,
collection_id: UUID,
ids: IDs,
embeddings: Optional[Embeddings] = None,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
) -> bool:
"""Add embeddings to the data store. This is the most general way to add embeddings to the database.
⚠️ It is recommended to use the more specific methods below when possible.
Args:
collection_id: The collection to add the embeddings to
embedding: The sequence of embeddings to add
"""
pass
@abstractmethod
def _upsert(
self,
collection_id: UUID,
ids: IDs,
embeddings: Embeddings,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
increment_index: bool = True,
) -> bool:
"""Add or update entries in the embedding store.
If an entry with the same id already exists, it will be updated, otherwise it will be added.
Args:
collection_id: The collection to add the embeddings to
ids: The ids to associate with the embeddings. Defaults to None.
embeddings: The sequence of embeddings to add
metadatas: The metadata to associate with the embeddings. Defaults to None.
documents: The documents to associate with the embeddings. Defaults to None.
increment_index: If True, will incrementally add to the ANN index of the collection. Defaults to True.
"""
pass
@abstractmethod
def _count(self, collection_id: UUID) -> int:
"""Returns the number of embeddings in the database
Args:
collection_id: The collection to count the embeddings in.
Returns:
int: The number of embeddings in the collection
"""
pass
@abstractmethod
def _peek(self, collection_id: UUID, n: int = 10) -> GetResult:
pass
@abstractmethod
def _get(
self,
collection_id: UUID,
ids: Optional[IDs] = None,
where: Optional[Where] = {},
sort: Optional[str] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
page: Optional[int] = None,
page_size: Optional[int] = None,
where_document: Optional[WhereDocument] = {},
include: Include = ["embeddings", "metadatas", "documents"],
) -> GetResult:
"""Gets embeddings from the database. Supports filtering, sorting, and pagination.
⚠️ This method should not be used directly.
Args:
where: A dictionary of key-value pairs to filter the embeddings by. Defaults to {}.
sort: The column to sort the embeddings by. Defaults to None.
limit: The maximum number of embeddings to return. Defaults to None.
offset: The number of embeddings to skip before returning. Defaults to None.
page: The page number to return. Defaults to None.
page_size: The number of embeddings to return per page. Defaults to None.
Returns:
pd.DataFrame: A pandas dataframe containing the embeddings and metadata
"""
pass
@abstractmethod
def _delete(
self,
collection_id: UUID,
ids: Optional[IDs],
where: Optional[Where] = {},
where_document: Optional[WhereDocument] = {},
) -> IDs:
"""Deletes embeddings from the database
⚠️ This method should not be used directly.
Args:
where: A dictionary of key-value pairs to filter the embeddings by. Defaults to {}.
Returns:
List: The list of internal UUIDs of the deleted embeddings
"""
pass
@abstractmethod
def _query(
self,
collection_id: UUID,
query_embeddings: Embeddings,
n_results: int = 10,
where: Where = {},
where_document: WhereDocument = {},
include: Include = ["embeddings", "metadatas", "documents", "distances"],
) -> QueryResult:
"""Gets the nearest neighbors of a single embedding
⚠️ This method should not be used directly.
Args:
embedding: The embedding to find the nearest neighbors of
n_results: The number of nearest neighbors to return. Defaults to 10.
where: A dictionary of key-value pairs to filter the embeddings by. Defaults to {}.
"""
pass
@override
@abstractmethod
def reset(self) -> None:
"""Resets the database
⚠️ This is destructive and will delete all data in the database.
Args:
None
Returns:
None
"""
pass
@abstractmethod
def raw_sql(self, sql: str) -> pd.DataFrame:
"""Runs a raw SQL query against the database
⚠️ This method should not be used directly.
Args:
sql: The SQL query to run
Returns:
pd.DataFrame: A pandas dataframe containing the results of the query
"""
pass
@abstractmethod
def create_index(self, collection_name: str) -> bool:
"""Creates an index for the given collection
⚠️ This method should not be used directly.
Args:
collection_name: The collection to create the index for. Uses the client's collection if None. Defaults to None.
Returns:
bool: True if the index was created successfully
"""
pass
@abstractmethod
def persist(self) -> bool:
"""Persist the database to disk"""
pass
@abstractmethod
def get_version(self) -> str:
"""Get the version of Chroma.
Returns:
str: The version of Chroma
"""
pass
|