from abc import abstractmethod
from typing import Awaitable, Callable, TypeVar, Any
from chromadb.config import Component, System
T = TypeVar("T", bound=Callable[..., Any])
A = TypeVar("A", bound=Awaitable[Any])
class RateLimitEnforcer(Component):
"""
Rate limit enforcer.
Implemented as a wrapper around server functions to block requests if rate limits are exceeded.
"""
def __init__(self, system: System) -> None:
super().__init__(system)
@abstractmethod
def rate_limit(self, func: T) -> T:
pass
class AsyncRateLimitEnforcer(Component):
"""
Rate limit enforcer.
Implemented as a wrapper around async functions to block requests if rate limits are exceeded.
"""
def __init__(self, system: System) -> None:
super().__init__(system)
@abstractmethod
def rate_limit(self, func: A) -> A:
pass