From 6d86808eaffc1b5dbd9f59d0e4a4928fb6a0aae0 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sun, 1 Feb 2026 04:06:35 +0530 Subject: [PATCH] =?UTF-8?q?feat:=20enforce=20model-level=20TPM/RPM=20limit?= =?UTF-8?q?s=20(enforce=5Fmodel=5Frate=5Flimits)=20=E2=80=A6=20(#19230)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: enforce model-level TPM/RPM limits (enforce_model_rate_limits) flag * fix lint errors --- docs/my-website/docs/proxy/load_balancing.md | 61 +++ litellm/router.py | 5 + .../pre_call_checks/model_rate_limit_check.py | 373 ++++++++++++++++++ litellm/types/router.py | 29 +- .../test_enforce_model_rate_limits.py | 315 +++++++++++++++ 5 files changed, 768 insertions(+), 15 deletions(-) create mode 100644 litellm/router_utils/pre_call_checks/model_rate_limit_check.py create mode 100644 tests/test_litellm/test_router/test_enforce_model_rate_limits.py diff --git a/docs/my-website/docs/proxy/load_balancing.md b/docs/my-website/docs/proxy/load_balancing.md index 42f6ef1aa5..186307d649 100644 --- a/docs/my-website/docs/proxy/load_balancing.md +++ b/docs/my-website/docs/proxy/load_balancing.md @@ -69,6 +69,67 @@ router_settings: redis_port: 1992 ``` +## Enforce Model Rate Limits + +Strictly enforce RPM/TPM limits set on deployments. When limits are exceeded, requests are blocked **before** reaching the LLM provider with a `429 Too Many Requests` error. + +:::info +By default, `rpm` and `tpm` values are only used for **routing decisions** (picking deployments with capacity). With `enforce_model_rate_limits`, they become **hard limits**. +::: + +### Quick Start + +```yaml +model_list: + - model_name: gpt-4 + litellm_params: + model: openai/gpt-4 + api_key: os.environ/OPENAI_API_KEY + rpm: 60 # 60 requests per minute + tpm: 90000 # 90k tokens per minute + +router_settings: + optional_pre_call_checks: + - enforce_model_rate_limits # 👈 Enables strict enforcement +``` + +### How It Works + +| Limit Type | Enforcement | Accuracy | +|------------|-------------|----------| +| **RPM** | Hard limit - blocked at exact threshold | 100% accurate | +| **TPM** | Best-effort - may slightly exceed | Blocked when already over limit | + +**Why TPM is best-effort:** Token count is unknown until the LLM responds. TPM is checked before each request (blocks if already over), and tracked after (adds actual tokens used). + +### Error Response + +```json +{ + "error": { + "message": "Model rate limit exceeded. RPM limit=60, current usage=60", + "type": "rate_limit_error", + "code": 429 + } +} +``` + +Response includes `retry-after: 60` header. + +### Multi-Instance Deployment + +For multiple LiteLLM proxy instances, add Redis to share rate limit state: + +```yaml +router_settings: + optional_pre_call_checks: + - enforce_model_rate_limits + redis_host: redis.example.com + redis_port: 6379 + redis_password: your-password +``` + + :::info Detailed information about [routing strategies can be found here](../routing) ::: diff --git a/litellm/router.py b/litellm/router.py index e117e8c09a..65445e29c4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -118,6 +118,9 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import from litellm.router_utils.pre_call_checks.responses_api_deployment_check import ( ResponsesApiDeploymentCheck, ) +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( + ModelRateLimitingCheck, +) from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, increment_deployment_successes_for_current_minute, @@ -1195,6 +1198,8 @@ class Router: ) elif pre_call_check == "responses_api_deployment_check": _callback = ResponsesApiDeploymentCheck() + elif pre_call_check == "enforce_model_rate_limits": + _callback = ModelRateLimitingCheck(dual_cache=self.cache) if _callback is not None: if self.optional_callbacks is None: self.optional_callbacks = [] diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py new file mode 100644 index 0000000000..e5be61690b --- /dev/null +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -0,0 +1,373 @@ +""" +Enforce TPM/RPM rate limits set on model deployments. + +This pre-call check ensures that model-level TPM/RPM limits are enforced +across all requests, regardless of routing strategy. + +When enabled via `enforce_model_rate_limits: true` in litellm_settings, +requests that exceed the configured TPM/RPM limits will receive a 429 error. +""" + +from typing import TYPE_CHECKING, Any, Dict, Optional, Union + +import httpx + +import litellm +from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.router import RouterErrors +from litellm.types.utils import StandardLoggingPayload +from litellm.utils import get_utc_datetime + +if TYPE_CHECKING: + from opentelemetry.trace import Span as _Span + + Span = Union[_Span, Any] +else: + Span = Any + + +class RoutingArgs: + ttl: int = 60 # 1min (RPM/TPM expire key) + + +class ModelRateLimitingCheck(CustomLogger): + """ + Pre-call check that enforces TPM/RPM limits on model deployments. + + This check runs before each request and raises a RateLimitError + if the deployment has exceeded its configured TPM or RPM limits. + + Unlike the usage-based-routing strategy which uses limits for routing decisions, + this check actively enforces those limits across ALL routing strategies. + """ + + def __init__(self, dual_cache: DualCache): + self.dual_cache = dual_cache + + def _get_deployment_limits( + self, deployment: Dict + ) -> tuple[Optional[int], Optional[int]]: + """ + Extract TPM and RPM limits from a deployment configuration. + + Checks in order: + 1. Top-level 'tpm'/'rpm' fields + 2. litellm_params.tpm/rpm + 3. model_info.tpm/rpm + + Returns: + Tuple of (tpm_limit, rpm_limit) + """ + # Check top-level + tpm = deployment.get("tpm") + rpm = deployment.get("rpm") + + # Check litellm_params + if tpm is None: + tpm = deployment.get("litellm_params", {}).get("tpm") + if rpm is None: + rpm = deployment.get("litellm_params", {}).get("rpm") + + # Check model_info + if tpm is None: + tpm = deployment.get("model_info", {}).get("tpm") + if rpm is None: + rpm = deployment.get("model_info", {}).get("rpm") + + return tpm, rpm + + def _get_cache_keys(self, deployment: Dict, current_minute: str) -> tuple[str, str]: + """Get the cache keys for TPM and RPM tracking.""" + model_id = deployment.get("model_info", {}).get("id") + deployment_name = deployment.get("litellm_params", {}).get("model") + + tpm_key = f"{model_id}:{deployment_name}:tpm:{current_minute}" + rpm_key = f"{model_id}:{deployment_name}:rpm:{current_minute}" + + return tpm_key, rpm_key + + def pre_call_check(self, deployment: Dict) -> Optional[Dict]: + """ + Synchronous pre-call check for model rate limits. + + Raises RateLimitError if deployment exceeds TPM/RPM limits. + """ + try: + tpm_limit, rpm_limit = self._get_deployment_limits(deployment) + + # If no limits are set, allow the request + if tpm_limit is None and rpm_limit is None: + return deployment + + dt = get_utc_datetime() + current_minute = dt.strftime("%H-%M") + tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) + + model_id = deployment.get("model_info", {}).get("id") + model_name = deployment.get("litellm_params", {}).get("model") + model_group = deployment.get("model_name", "") + + # Check TPM limit + if tpm_limit is not None: + # First check local cache + current_tpm = self.dual_cache.get_cache(key=tpm_key, local_only=True) + if current_tpm is not None and current_tpm >= tpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} tpm limit={tpm_limit}. current usage={current_tpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + ) + + # Check RPM limit + if rpm_limit is not None: + # First check local cache + current_rpm = self.dual_cache.get_cache(key=rpm_key, local_only=True) + if current_rpm >= rpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + ) + + # Check redis cache and increment + current_rpm = self.dual_cache.increment_cache( + key=rpm_key, value=1, ttl=RoutingArgs.ttl + ) + if current_rpm is not None and current_rpm > rpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + ) + + return deployment + + except litellm.RateLimitError: + raise + except Exception as e: + verbose_router_logger.debug( + f"Error in ModelRateLimitingCheck.pre_call_check: {str(e)}" + ) + # Don't fail the request if rate limit check fails + return deployment + + async def async_pre_call_check( + self, deployment: Dict, parent_otel_span: Optional[Span] = None + ) -> Optional[Dict]: + """ + Async pre-call check for model rate limits. + + Raises RateLimitError if deployment exceeds TPM/RPM limits. + """ + try: + tpm_limit, rpm_limit = self._get_deployment_limits(deployment) + + # If no limits are set, allow the request + if tpm_limit is None and rpm_limit is None: + return deployment + + dt = get_utc_datetime() + current_minute = dt.strftime("%H-%M") + tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) + + model_id = deployment.get("model_info", {}).get("id") + model_name = deployment.get("litellm_params", {}).get("model") + model_group = deployment.get("model_name", "") + + # Check TPM limit + if tpm_limit is not None: + # First check local cache + current_tpm = await self.dual_cache.async_get_cache( + key=tpm_key, local_only=True + ) + if current_tpm is not None and current_tpm >= tpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} tpm limit={tpm_limit}. current usage={current_tpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + num_retries=0, # Don't retry - return 429 immediately + ) + + # Check RPM limit + if rpm_limit is not None: + # First check local cache + current_rpm = await self.dual_cache.async_get_cache( + key=rpm_key, local_only=True + ) + if current_rpm is not None and current_rpm >= rpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + num_retries=0, # Don't retry - return 429 immediately + ) + + # Check redis cache and increment + current_rpm = await self.dual_cache.async_increment_cache( + key=rpm_key, + value=1, + ttl=RoutingArgs.ttl, + parent_otel_span=parent_otel_span, + ) + if current_rpm is not None and current_rpm > rpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + num_retries=0, # Don't retry - return 429 immediately + ) + + return deployment + + except litellm.RateLimitError: + raise + except Exception as e: + verbose_router_logger.debug( + f"Error in ModelRateLimitingCheck.async_pre_call_check: {str(e)}" + ) + # Don't fail the request if rate limit check fails + return deployment + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + Track TPM usage after successful request. + + This updates the TPM counter with the actual tokens used. + Always tracks tokens - the pre-call check handles enforcement. + """ + try: + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object" + ) + if standard_logging_object is None: + return + + model_id = standard_logging_object.get("model_id") + if model_id is None: + return + + total_tokens = standard_logging_object.get("total_tokens", 0) + model = standard_logging_object.get("hidden_params", {}).get( + "litellm_model_name" + ) + + verbose_router_logger.debug( + f"[TPM TRACKING] model_id={model_id}, total_tokens={total_tokens}, model={model}" + ) + + if not model or not total_tokens: + return + + dt = get_utc_datetime() + current_minute = dt.strftime("%H-%M") + tpm_key = f"{model_id}:{model}:tpm:{current_minute}" + + verbose_router_logger.debug( + f"[TPM TRACKING] Incrementing {tpm_key} by {total_tokens}" + ) + + await self.dual_cache.async_increment_cache( + key=tpm_key, + value=total_tokens, + ttl=RoutingArgs.ttl, + ) + + except Exception as e: + verbose_router_logger.debug( + f"Error in ModelRateLimitingCheck.async_log_success_event: {str(e)}" + ) + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + Sync version of tracking TPM usage after successful request. + Always tracks tokens - the pre-call check handles enforcement. + """ + try: + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object" + ) + if standard_logging_object is None: + return + + model_id = standard_logging_object.get("model_id") + if model_id is None: + return + + total_tokens = standard_logging_object.get("total_tokens", 0) + model = standard_logging_object.get("hidden_params", {}).get( + "litellm_model_name" + ) + + if not model or not total_tokens: + return + + dt = get_utc_datetime() + current_minute = dt.strftime("%H-%M") + tpm_key = f"{model_id}:{model}:tpm:{current_minute}" + + self.dual_cache.increment_cache( + key=tpm_key, + value=total_tokens, + ttl=RoutingArgs.ttl, + ) + + except Exception as e: + verbose_router_logger.debug( + f"Error in ModelRateLimitingCheck.log_success_event: {str(e)}" + ) diff --git a/litellm/types/router.py b/litellm/types/router.py index f31c6df300..f78789c977 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -95,18 +95,16 @@ class ModelInfo(BaseModel): id: Optional[ str ] # Allow id to be optional on input, but it will always be present as a str in the model instance - db_model: bool = ( - False # used for proxy - to separate models which are stored in the db vs. config. - ) + db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. updated_at: Optional[datetime.datetime] = None updated_by: Optional[str] = None created_at: Optional[datetime.datetime] = None created_by: Optional[str] = None - base_model: Optional[str] = ( - None # specify if the base model is azure/gpt-3.5-turbo etc for accurate cost tracking - ) + base_model: Optional[ + str + ] = None # specify if the base model is azure/gpt-3.5-turbo etc for accurate cost tracking tier: Optional[Literal["free", "paid"]] = None """ @@ -172,12 +170,12 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): custom_llm_provider: Optional[str] = None tpm: Optional[int] = None rpm: Optional[int] = None - timeout: Optional[Union[float, str, httpx.Timeout]] = ( - None # if str, pass in as os.environ/ - ) - stream_timeout: Optional[Union[float, str]] = ( - None # timeout when making stream=True calls, if str, pass in as os.environ/ - ) + timeout: Optional[ + Union[float, str, httpx.Timeout] + ] = None # if str, pass in as os.environ/ + stream_timeout: Optional[ + Union[float, str] + ] = None # timeout when making stream=True calls, if str, pass in as os.environ/ max_retries: Optional[int] = None organization: Optional[str] = None # for openai orgs configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None @@ -276,9 +274,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): if max_retries is not None and isinstance(max_retries, str): max_retries = int(max_retries) # cast to int # We need to keep max_retries in args since it's a parameter of GenericLiteLLMParams - args["max_retries"] = ( - max_retries # Put max_retries back in args after popping it - ) + args[ + "max_retries" + ] = max_retries # Put max_retries back in args after popping it super().__init__(**args, **params) def __contains__(self, key): @@ -805,6 +803,7 @@ OptionalPreCallChecks = List[ "router_budget_limiting", "responses_api_deployment_check", "forward_client_headers_by_model_group", + "enforce_model_rate_limits", ] ] diff --git a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py b/tests/test_litellm/test_router/test_enforce_model_rate_limits.py new file mode 100644 index 0000000000..3bca3df4e1 --- /dev/null +++ b/tests/test_litellm/test_router/test_enforce_model_rate_limits.py @@ -0,0 +1,315 @@ +""" +Tests for enforce_model_rate_limits feature. + +This feature allows users to enforce TPM/RPM limits set on model deployments +regardless of the routing strategy being used. +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm +from litellm import Router +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( + ModelRateLimitingCheck, +) + + +class TestModelRateLimitingCheck: + """Test the ModelRateLimitingCheck class directly.""" + + def test_get_deployment_limits_from_top_level(self): + """Test extracting limits from top-level deployment config.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "tpm": 1000, + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 1000 + assert rpm == 10 + + def test_get_deployment_limits_from_litellm_params(self): + """Test extracting limits from litellm_params.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4", "tpm": 2000, "rpm": 20}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 2000 + assert rpm == 20 + + def test_get_deployment_limits_from_model_info(self): + """Test extracting limits from model_info.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id", "tpm": 3000, "rpm": 30}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 3000 + assert rpm == 30 + + def test_get_deployment_limits_none_when_not_set(self): + """Test that None is returned when limits are not set.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm is None + assert rpm is None + + def test_pre_call_check_allows_request_when_no_limits(self): + """Test that requests are allowed when no limits are set.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + result = check.pre_call_check(deployment) + assert result == deployment + + def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self): + """Test that RateLimitError is raised when RPM limit is exceeded.""" + mock_cache = MagicMock() + mock_cache.get_cache.return_value = 10 # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "RPM limit=10" in str(exc_info.value) + assert "current usage=10" in str(exc_info.value) + + def test_pre_call_check_allows_request_under_limit(self): + """Test that requests are allowed when under the limit.""" + mock_cache = MagicMock() + mock_cache.get_cache.return_value = 5 + mock_cache.increment_cache.return_value = 6 + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + result = check.pre_call_check(deployment) + assert result == deployment + + def test_pre_call_check_raises_rate_limit_error_when_over_tpm(self): + """Test that RateLimitError is raised when TPM limit is exceeded.""" + mock_cache = MagicMock() + mock_cache.get_cache.return_value = 1000 # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "TPM limit=1000" in str(exc_info.value) + assert "current usage=1000" in str(exc_info.value) + + def test_log_success_event_increments_cache(self): + """Test that log_success_event correctly increments the cache.""" + mock_cache = MagicMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } + } + + check.log_success_event(kwargs, None, None, None) + + # Verify increment_cache was called + mock_cache.increment_cache.assert_called_once() + _, kwarg_params = mock_cache.increment_cache.call_args + assert "test-id:gpt-4:tpm:" in kwarg_params["key"] + assert kwarg_params["value"] == 50 + + +class TestModelRateLimitingCheckAsync: + """Test async methods of ModelRateLimitingCheck.""" + + @pytest.mark.asyncio + async def test_async_pre_call_check_allows_request_when_no_limits(self): + """Test that requests are allowed when no limits are set (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + result = await check.async_pre_call_check(deployment) + assert result == deployment + + @pytest.mark.asyncio + async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self): + """Test that RateLimitError is raised when RPM limit is exceeded (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=10) # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "RPM limit=10" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_allows_request_under_limit(self): + """Test that requests are allowed when under the limit (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=5) + mock_cache.async_increment_cache = AsyncMock(return_value=6) + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + result = await check.async_pre_call_check(deployment) + assert result == deployment + + @pytest.mark.asyncio + async def test_async_pre_call_check_raises_rate_limit_error_when_over_tpm(self): + """Test that RateLimitError is raised when TPM limit is exceeded (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=1000) # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "TPM limit=1000" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_log_success_event_increments_cache(self): + """Test that async_log_success_event correctly increments the cache.""" + mock_cache = MagicMock() + mock_cache.async_increment_cache = AsyncMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } + } + + await check.async_log_success_event(kwargs, None, None, None) + + # Verify async_increment_cache was called + mock_cache.async_increment_cache.assert_called_once() + _, kwarg_params = mock_cache.async_increment_cache.call_args + assert "test-id:gpt-4:tpm:" in kwarg_params["key"] + assert kwarg_params["value"] == 50 + + +class TestRouterWithEnforceModelRateLimits: + """Test Router integration with enforce_model_rate_limits.""" + + def test_router_initializes_with_enforce_model_rate_limits(self): + """Test that Router properly initializes the ModelRateLimitingCheck.""" + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "test"}, + "rpm": 10, + } + ] + + router = Router( + model_list=model_list, + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + # Check that the callback was added + assert router.optional_callbacks is not None + assert len(router.optional_callbacks) == 1 + assert isinstance(router.optional_callbacks[0], ModelRateLimitingCheck) + + def test_router_optional_callbacks_contains_model_rate_limiting(self): + """Test that ModelRateLimitingCheck is in the callbacks list.""" + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "test"}, + "rpm": 10, + } + ] + + Router( + model_list=model_list, + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + # Find the ModelRateLimitingCheck in litellm.callbacks + found = False + for callback in litellm.callbacks: + if isinstance(callback, ModelRateLimitingCheck): + found = True + break + + assert found, "ModelRateLimitingCheck should be in litellm.callbacks"