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 index af3d7ddfac7..b7e0a58523c 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -11,6 +11,8 @@ is logged the first time such a deployment is seen. """ import contextlib +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx @@ -46,6 +48,9 @@ class RoutingArgs: ttl: int = 60 # 1min (RPM/TPM expire key) +_NO_LITELLM_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({}) + + class ModelRateLimitingCheck(CustomLogger): """ Pre-call check that enforces TPM/RPM or ITPM/OTPM limits on model deployments. @@ -304,6 +309,24 @@ class ModelRateLimitingCheck(CustomLogger): # Don't fail the request if rate limit check fails return deployment + @staticmethod + def _model_id_from_kwargs(kwargs: Mapping[str, Any]) -> str | None: + """Recover the deployment id when the logging payload is missing it. + + On the streaming path the payload is built from a ``litellm_params`` + snapshot taken at ``Logging`` init, before the router stamped + ``model_info`` in, so ``model_id`` comes back as "". The router does + stamp ``kwargs["model_info"]``, so prefer that before giving up. + """ + litellm_params: Final = kwargs.get("litellm_params") or _NO_LITELLM_PARAMS + for source in (kwargs.get("model_info"), litellm_params.get("model_info")): + if isinstance(source, Mapping): + candidate = source.get("id") + if candidate: + return str(candidate) + fallback: Final = litellm_params.get("model_id") + return str(fallback) if fallback else None + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, @@ -333,8 +356,13 @@ class ModelRateLimitingCheck(CustomLogger): if standard_logging_object is None: return - model_id: Final = standard_logging_object.get("model_id") - if model_id is None: + model_id: Final = standard_logging_object.get("model_id") or self._model_id_from_kwargs(kwargs) + if not model_id: + # An empty string reaches here as readily as None: the standard + # logging payload defaults model_id to "" when metadata carries + # no model_info. Keying the counter on it buckets every such + # deployment under "::tpm:" while the real + # deployment's key stays at zero, so limits never trigger. return total_tokens: Final = standard_logging_object.get("total_tokens", 0) @@ -397,8 +425,13 @@ class ModelRateLimitingCheck(CustomLogger): if standard_logging_object is None: return - model_id: Final = standard_logging_object.get("model_id") - if model_id is None: + model_id: Final = standard_logging_object.get("model_id") or self._model_id_from_kwargs(kwargs) + if not model_id: + # An empty string reaches here as readily as None: the standard + # logging payload defaults model_id to "" when metadata carries + # no model_info. Keying the counter on it buckets every such + # deployment under "::tpm:" while the real + # deployment's key stays at zero, so limits never trigger. return total_tokens: Final = standard_logging_object.get("total_tokens", 0) 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 index 1def253ac93..04ad96ae15d 100644 --- a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py +++ b/tests/test_litellm/test_router/test_enforce_model_rate_limits.py @@ -268,6 +268,78 @@ class TestModelRateLimitingCheckAsync: assert "test-id:gpt-4:tpm:" in kwarg_params["key"] assert kwarg_params["value"] == 50 + @pytest.mark.asyncio + async def test_async_log_success_event_recovers_model_id_from_kwargs(self): + """An empty model_id in the payload falls back to the router's model_info. + + On the streaming path the standard logging payload is built from a + litellm_params snapshot taken before the router stamped model_info in, + so model_id arrives as "". The deployment id is still on kwargs. + """ + mock_cache = MagicMock() + mock_cache.async_increment_cache = AsyncMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + }, + "model_info": {"id": "deployment-1"}, + } + + await check.async_log_success_event(kwargs, None, None, None) + + mock_cache.async_increment_cache.assert_called_once() + _, kwarg_params = mock_cache.async_increment_cache.call_args + assert "deployment-1:gpt-4:tpm:" in kwarg_params["key"] + assert kwarg_params["value"] == 50 + + @pytest.mark.asyncio + async def test_async_log_success_event_falls_back_to_litellm_params(self): + """model_info nested under litellm_params is also accepted.""" + mock_cache = MagicMock() + mock_cache.async_increment_cache = AsyncMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "", + "total_tokens": 10, + "hidden_params": {"litellm_model_name": "gpt-4"}, + }, + "litellm_params": {"model_info": {"id": "deployment-2"}}, + } + + await check.async_log_success_event(kwargs, None, None, None) + + _, kwarg_params = mock_cache.async_increment_cache.call_args + assert "deployment-2:gpt-4:tpm:" in kwarg_params["key"] + + @pytest.mark.asyncio + async def test_async_log_success_event_skips_when_no_model_id_anywhere(self): + """With no deployment id at all, nothing is counted. + + Incrementing on an empty model_id would bucket unrelated deployments + under a shared "::tpm:" key. + """ + mock_cache = MagicMock() + mock_cache.async_increment_cache = AsyncMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } + } + + await check.async_log_success_event(kwargs, None, None, None) + + mock_cache.async_increment_cache.assert_not_called() + class TestRouterWithEnforceModelRateLimits: """Test Router integration with enforce_model_rate_limits."""