codekingpro/portable-devtools

1from abc import abstractmethod
2from typing import Awaitable, Callable, TypeVar, Any
3from chromadb.config import Component, System
4
5T = TypeVar("T", bound=Callable[..., Any])
6A = TypeVar("A", bound=Awaitable[Any])
7
8
9class RateLimitEnforcer(Component):
10    """
11    Rate limit enforcer.
12
13    Implemented as a wrapper around server functions to block requests if rate limits are exceeded.
14    """
15
16    def __init__(self, system: System) -> None:
17        super().__init__(system)
18
19    @abstractmethod
20    def rate_limit(self, func: T) -> T:
21        pass
22
23
24class AsyncRateLimitEnforcer(Component):
25    """
26    Rate limit enforcer.
27
28    Implemented as a wrapper around async functions to block requests if rate limits are exceeded.
29    """
30
31    def __init__(self, system: System) -> None:
32        super().__init__(system)
33
34    @abstractmethod
35    def rate_limit(self, func: A) -> A:
36        pass
37