from abc import abstractmethod
from enum import Enum
from typing import Dict, Any, Optional
from chromadb.api.types import (
Embeddings,
Metadatas,
Documents,
URIs,
IDs,
CollectionMetadata,
Where,
WhereDocument,
)
from chromadb.config import Component, System
class Action(str, Enum):
CREATE_DATABASE = "create_database"
CREATE_COLLECTION = "create_collection"
LIST_COLLECTIONS = "list_collections"
UPDATE_COLLECTION = "update_collection"
ADD = "add"
GET = "get"
DELETE = "delete"
UPDATE = "update"
UPSERT = "upsert"
QUERY = "query"
class QuotaEnforcer(Component):
"""
Exposes hooks to enforce quotas.
"""
def __init__(self, system: System) -> None:
super().__init__(system)
@abstractmethod
def set_context(self, context: Dict[str, Any]) -> None:
"""
Sets the context for a given request.
"""
pass
@abstractmethod
def enforce(
self,
action: Action,
tenant: str,
metadatas: Optional[Metadatas] = None,
documents: Optional[Documents] = None,
embeddings: Optional[Embeddings] = None,
uris: Optional[URIs] = None,
ids: Optional[IDs] = None,
name: Optional[str] = None,
new_name: Optional[str] = None,
metadata: Optional[CollectionMetadata] = None,
new_metadata: Optional[CollectionMetadata] = None,
limit: Optional[int] = None,
where: Optional[Where] = None,
where_document: Optional[WhereDocument] = None,
n_results: Optional[int] = None,
query_embeddings: Optional[Embeddings] = None,
) -> None:
"""
Enforces a quota.
"""
pass