codekingpro/portable-devtools

1from abc import abstractmethod
2from enum import Enum
3from typing import Dict, Any, Optional
4from uuid import UUID
5
6from chromadb.api.types import (
7    Embeddings,
8    Metadatas,
9    Documents,
10    URIs,
11    IDs,
12    CollectionMetadata,
13    Where,
14    WhereDocument,
15)
16from chromadb.config import Component, System
17
18
19class Action(str, Enum):
20    CREATE_DATABASE = "create_database"
21    CREATE_COLLECTION = "create_collection"
22    LIST_COLLECTIONS = "list_collections"
23    UPDATE_COLLECTION = "update_collection"
24    ADD = "add"
25    GET = "get"
26    DELETE = "delete"
27    UPDATE = "update"
28    UPSERT = "upsert"
29    QUERY = "query"
30
31
32class QuotaEnforcer(Component):
33    """
34    Exposes hooks to enforce quotas.
35    """
36
37    def __init__(self, system: System) -> None:
38        super().__init__(system)
39
40    @abstractmethod
41    def set_context(self, context: Dict[str, Any]) -> None:
42        """
43        Sets the context for a given request.
44        """
45        pass
46
47    @abstractmethod
48    def enforce(
49        self,
50        action: Action,
51        tenant: str,
52        metadatas: Optional[Metadatas] = None,
53        documents: Optional[Documents] = None,
54        embeddings: Optional[Embeddings] = None,
55        uris: Optional[URIs] = None,
56        ids: Optional[IDs] = None,
57        name: Optional[str] = None,
58        new_name: Optional[str] = None,
59        metadata: Optional[CollectionMetadata] = None,
60        new_metadata: Optional[CollectionMetadata] = None,
61        limit: Optional[int] = None,
62        where: Optional[Where] = None,
63        where_document: Optional[WhereDocument] = None,
64        n_results: Optional[int] = None,
65        query_embeddings: Optional[Embeddings] = None,
66        collection_id: Optional[UUID] = None,
67    ) -> None:
68        """
69        Enforces a quota.
70        """
71        pass
72