Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[ENH] HTTP API routes overhaul #2954

Merged
merged 29 commits into from
Oct 18, 2024
Merged

[ENH] HTTP API routes overhaul #2954

Show file tree
Hide file tree
Changes from 27 commits
Commits
Show all changes
29 commits
Select commit Hold shift + click to select a range
94a8593
[ENH] REST API routes overhaul
drewkim Oct 14, 2024
9eb16ea
Fix route bug
drewkim Oct 14, 2024
4219528
More route bugs
drewkim Oct 14, 2024
6f170d0
Update JS client
drewkim Oct 15, 2024
1e1d7a1
Cleanup
drewkim Oct 15, 2024
13b6ffb
Bug fix + tenant/db validation
drewkim Oct 15, 2024
ab93ad4
Add old routes back
drewkim Oct 15, 2024
55182b3
Generate JS client with v1 routes
drewkim Oct 15, 2024
2cddf58
Patch resolve_tenant_and_databases to resolve test failure
drewkim Oct 15, 2024
da34188
Error handling
drewkim Oct 15, 2024
371ea7f
Address comments
drewkim Oct 16, 2024
66429a8
JS client update
drewkim Oct 16, 2024
c0ef648
JS client update + nits
drewkim Oct 16, 2024
3cea6fd
database -> database_name
drewkim Oct 16, 2024
b8741ce
Update JS client
drewkim Oct 16, 2024
10c4453
Update dos and terraform
drewkim Oct 16, 2024
70a1c2e
Update dos and terraform
drewkim Oct 16, 2024
983f11f
House auth utils in auth path, error when resolved tenant/db don't ma…
drewkim Oct 16, 2024
1c8f562
Move v1 routes out into separate file/wrapper class
drewkim Oct 16, 2024
c4cfe18
Update v1 trace names
drewkim Oct 17, 2024
5773156
Add tenant to auth'ed request spans
drewkim Oct 17, 2024
e459823
Move post route actions into V1
drewkim Oct 18, 2024
d728c2d
Move v1 routes back into FastAPI
drewkim Oct 18, 2024
a6d9d20
Tests and nits
drewkim Oct 18, 2024
41c8fb4
cleanup
drewkim Oct 18, 2024
7133f2c
Rename auth utils
drewkim Oct 18, 2024
a15e07d
fix test and update js client
drewkim Oct 18, 2024
3aaf8a8
fix test port issue
drewkim Oct 18, 2024
3b222b9
Update JS client
drewkim Oct 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
134 changes: 134 additions & 0 deletions chromadb/api/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
GetResult,
WhereDocument,
)
from chromadb.auth import UserIdentity
from chromadb.config import Component, Settings
from chromadb.types import Database, Tenant, Collection as CollectionModel
import chromadb.utils.embedding_functions as ef
Expand Down Expand Up @@ -328,6 +329,14 @@ def get_max_batch_size(self) -> int:
"""Return the maximum number of records that can be created or mutated in a single call."""
pass

@abstractmethod
def get_user_identity(self) -> UserIdentity:
"""Resolve the tenant and databases for the client. Returns the default
values if can't be resolved.

"""
pass


class ClientAPI(BaseAPI, ABC):
tenant: str
Expand Down Expand Up @@ -594,3 +603,128 @@ def delete_collection(
database: str = DEFAULT_DATABASE,
) -> None:
pass

@abstractmethod
@override
def _modify(
self,
id: UUID,
new_name: Optional[str] = None,
new_metadata: Optional[CollectionMetadata] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> None:
pass

@abstractmethod
@override
def _count(
self,
collection_id: UUID,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> int:
pass

@abstractmethod
@override
def _peek(
self,
collection_id: UUID,
n: int = 10,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> GetResult:
pass

@abstractmethod
@override
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 = ["metadatas", "documents"], # type: ignore[list-item]
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> GetResult:
pass

@abstractmethod
@override
def _add(
self,
ids: IDs,
collection_id: UUID,
embeddings: Embeddings,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
uris: Optional[URIs] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> bool:
pass

@abstractmethod
@override
def _update(
self,
collection_id: UUID,
ids: IDs,
embeddings: Optional[Embeddings] = None,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
uris: Optional[URIs] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> bool:
pass

@abstractmethod
@override
def _upsert(
self,
collection_id: UUID,
ids: IDs,
embeddings: Embeddings,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
uris: Optional[URIs] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> bool:
pass

@abstractmethod
@override
def _query(
self,
collection_id: UUID,
query_embeddings: Embeddings,
n_results: int = 10,
where: Where = {},
where_document: WhereDocument = {},
include: Include = ["metadatas", "documents", "distances"], # type: ignore[list-item]
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> QueryResult:
pass

@abstractmethod
@override
def _delete(
self,
collection_id: UUID,
ids: Optional[IDs] = None,
where: Optional[Where] = {},
where_document: Optional[WhereDocument] = {},
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> None:
pass
134 changes: 134 additions & 0 deletions chromadb/api/async_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
CollectionConfiguration,
CollectionConfigurationInternal,
)
from chromadb.auth import UserIdentity
from chromadb.api.models.AsyncCollection import AsyncCollection
from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT
from chromadb.api.types import (
Expand Down Expand Up @@ -319,6 +320,14 @@ async def get_max_batch_size(self) -> int:
"""Return the maximum number of records that can be created or mutated in a single call."""
pass

@abstractmethod
async def get_user_identity(self) -> UserIdentity:
"""Resolve the tenant and databases for the client. Returns the default
values if can't be resolved.

"""
pass


class AsyncClientAPI(AsyncBaseAPI, ABC):
tenant: str
Expand Down Expand Up @@ -585,3 +594,128 @@ async def delete_collection(
database: str = DEFAULT_DATABASE,
) -> None:
pass

@abstractmethod
@override
async def _modify(
self,
id: UUID,
new_name: Optional[str] = None,
new_metadata: Optional[CollectionMetadata] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> None:
pass

@abstractmethod
@override
async def _count(
self,
collection_id: UUID,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> int:
pass

@abstractmethod
@override
async def _peek(
self,
collection_id: UUID,
n: int = 10,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> GetResult:
pass

@abstractmethod
@override
async 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 = ["metadatas", "documents"], # type: ignore[list-item]
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> GetResult:
pass

@abstractmethod
@override
async def _add(
self,
ids: IDs,
collection_id: UUID,
embeddings: Embeddings,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
uris: Optional[URIs] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> bool:
pass

@abstractmethod
@override
async def _update(
self,
collection_id: UUID,
ids: IDs,
embeddings: Optional[Embeddings] = None,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
uris: Optional[URIs] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> bool:
pass

@abstractmethod
@override
async def _upsert(
self,
collection_id: UUID,
ids: IDs,
embeddings: Embeddings,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
uris: Optional[URIs] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> bool:
pass

@abstractmethod
@override
async def _query(
self,
collection_id: UUID,
query_embeddings: Embeddings,
n_results: int = 10,
where: Where = {},
where_document: WhereDocument = {},
include: Include = ["metadatas", "documents", "distances"], # type: ignore[list-item]
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> QueryResult:
pass

@abstractmethod
@override
async def _delete(
self,
collection_id: UUID,
ids: Optional[IDs] = None,
where: Optional[Where] = {},
where_document: Optional[WhereDocument] = {},
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> None:
pass
Loading
Loading