Merge 7fcf2c54ba into 0c98afa780

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 ":<model>:tpm:<minute>" 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 ":<model>:tpm:<minute>" while the real
+                # deployment's key stays at zero, so limits never trigger.
                 return
 
             total_tokens: Final = standard_logging_object.get("total_tokens", 0)

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 ":<model>:tpm:<minute>" 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."""