Add fix for bedrock_cache, metadata and max_model_budget (#18872)

litellm/proxy/auth/auth_utils.py

@@ -426,38 +426,65 @@ def get_key_model_rpm_limit(
     user_api_key_dict: UserAPIKeyAuth,
 ) -> Optional[Dict[str, int]]:
     """
-    Get the model rpm limit for a given api key
-    - check key metadata
-    - check key model max budget
-    - check team metadata
+    Get the model rpm limit for a given api key.
+
+    Priority order (returns first found):
+    1. Key metadata (model_rpm_limit)
+    2. Key model_max_budget (rpm_limit per model)
+    3. Team metadata (model_rpm_limit)
     """
+    # 1. Check key metadata first (takes priority)
     if user_api_key_dict.metadata:
-        if "model_rpm_limit" in user_api_key_dict.metadata:
-            return user_api_key_dict.metadata["model_rpm_limit"]
-    elif user_api_key_dict.model_max_budget:
+        result = user_api_key_dict.metadata.get("model_rpm_limit")
+        if result:
+            return result
+
+    # 2. Check model_max_budget
+    if user_api_key_dict.model_max_budget:
         model_rpm_limit: Dict[str, Any] = {}
         for model, budget in user_api_key_dict.model_max_budget.items():
-            if "rpm_limit" in budget and budget["rpm_limit"] is not None:
+            if isinstance(budget, dict) and budget.get("rpm_limit") is not None:
                 model_rpm_limit[model] = budget["rpm_limit"]
-        return model_rpm_limit
-    elif user_api_key_dict.team_metadata:
-        if "model_rpm_limit" in user_api_key_dict.team_metadata:
-            return user_api_key_dict.team_metadata["model_rpm_limit"]
+        if model_rpm_limit:
+            return model_rpm_limit
+
+    # 3. Fallback to team metadata
+    if user_api_key_dict.team_metadata:
+        return user_api_key_dict.team_metadata.get("model_rpm_limit")
+
     return None
 
 
 def get_key_model_tpm_limit(
     user_api_key_dict: UserAPIKeyAuth,
 ) -> Optional[Dict[str, int]]:
+    """
+    Get the model tpm limit for a given api key.
+
+    Priority order (returns first found):
+    1. Key metadata (model_tpm_limit)
+    2. Key model_max_budget (tpm_limit per model)
+    3. Team metadata (model_tpm_limit)
+    """
+    # 1. Check key metadata first (takes priority)
     if user_api_key_dict.metadata:
-        if "model_tpm_limit" in user_api_key_dict.metadata:
-            return user_api_key_dict.metadata["model_tpm_limit"]
-    elif user_api_key_dict.model_max_budget:
-        if "tpm_limit" in user_api_key_dict.model_max_budget:
-            return user_api_key_dict.model_max_budget["tpm_limit"]
-    elif user_api_key_dict.team_metadata:
-        if "model_tpm_limit" in user_api_key_dict.team_metadata:
-            return user_api_key_dict.team_metadata["model_tpm_limit"]
+        result = user_api_key_dict.metadata.get("model_tpm_limit")
+        if result:
+            return result
+
+    # 2. Check model_max_budget (iterate per-model like RPM does)
+    if user_api_key_dict.model_max_budget:
+        model_tpm_limit: Dict[str, Any] = {}
+        for model, budget in user_api_key_dict.model_max_budget.items():
+            if isinstance(budget, dict) and budget.get("tpm_limit") is not None:
+                model_tpm_limit[model] = budget["tpm_limit"]
+        if model_tpm_limit:
+            return model_tpm_limit
+
+    # 3. Fallback to team metadata
+    if user_api_key_dict.team_metadata:
+        return user_api_key_dict.team_metadata.get("model_tpm_limit")
+
     return None
 
 
@@ -469,7 +496,8 @@ def get_model_rate_limit_from_metadata(
     if getattr(user_api_key_dict, metadata_accessor_key):
         return getattr(user_api_key_dict, metadata_accessor_key).get(rate_limit_key)
     return None
-  
+
+
 def get_team_model_rpm_limit(
     user_api_key_dict: UserAPIKeyAuth,
 ) -> Optional[Dict[str, int]]:

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -167,7 +167,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             self.token_increment_script = None
 
         self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60))
-        
+
         # Batch rate limiter (lazy loaded)
         self._batch_rate_limiter: Optional[Any] = None
 
@@ -1013,7 +1013,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             )
             # Fail safe: enforce limits if we can't check
             return True
-    
+
     def get_rate_limiter_for_call_type(self, call_type: str) -> Optional[Any]:
         """Get the rate limiter for the call type."""
         if call_type == "acreate_batch":
@@ -1095,9 +1095,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
 
                 now = self._get_current_time().timestamp()
                 reset_time = now + self.window_size
-                reset_time_formatted = datetime.fromtimestamp(
-                    reset_time
-                ).strftime("%Y-%m-%d %H:%M:%S UTC")
+                reset_time_formatted = datetime.fromtimestamp(reset_time).strftime(
+                    "%Y-%m-%d %H:%M:%S UTC"
+                )
 
                 remaining_display = max(0, status["limit_remaining"])
                 rate_limit_type = status["rate_limit_type"]
@@ -1137,7 +1137,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         # Check if the call type has a specific rate limiter
         # eg. for Batch APIs we need to use the batch rate limiter to read the input file and count the tokens and requests
         #########################################################
-        call_type_specific_rate_limiter = self.get_rate_limiter_for_call_type(call_type=call_type)
+        call_type_specific_rate_limiter = self.get_rate_limiter_for_call_type(
+            call_type=call_type
+        )
         if call_type_specific_rate_limiter:
             return await call_type_specific_rate_limiter.async_pre_call_hook(
                 user_api_key_dict=user_api_key_dict,
@@ -1233,26 +1235,58 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
 
         return pipeline_operations
 
-    def _get_total_tokens_from_usage(self, usage: Any | None, rate_limit_type: Literal["output", "input", "total"]) -> int:
-        # Get total tokens from response
+    def _get_total_tokens_from_usage(
+        self, usage: Any | None, rate_limit_type: Literal["output", "input", "total"]
+    ) -> int:
+        """
+        Get total tokens from response usage for rate limiting.
+
+        For 'input' and 'total' rate limit types, cached tokens are excluded
+        because providers like AWS Bedrock don't count cached tokens toward
+        rate limits. This aligns LiteLLM's TPM calculation with provider behavior.
+        """
         total_tokens = 0
-        # spot fix for /responses api
+        cached_tokens = 0
+
         if usage:
             if isinstance(usage, Usage):
                 if rate_limit_type == "output":
-                    total_tokens = usage.completion_tokens
+                    total_tokens = usage.completion_tokens or 0
                 elif rate_limit_type == "input":
-                    total_tokens = usage.prompt_tokens
+                    total_tokens = usage.prompt_tokens or 0
                 elif rate_limit_type == "total":
-                    total_tokens = usage.total_tokens
+                    total_tokens = usage.total_tokens or 0
+
+                # Get cached tokens to exclude from input/total
+                if rate_limit_type in ("input", "total"):
+                    if (
+                        hasattr(usage, "prompt_tokens_details")
+                        and usage.prompt_tokens_details is not None
+                    ):
+                        cached_tokens = (
+                            getattr(usage.prompt_tokens_details, "cached_tokens", 0)
+                            or 0
+                        )
+
             elif isinstance(usage, dict):
-                # Responses API usage comes as a dict in ResponsesAPIResponse
+                # Responses API usage comes as a dict
                 if rate_limit_type == "output":
-                    total_tokens = usage.get("completion_tokens", 0)
+                    total_tokens = usage.get("completion_tokens", 0) or 0
                 elif rate_limit_type == "input":
-                    total_tokens = usage.get("prompt_tokens", 0)
+                    total_tokens = usage.get("prompt_tokens", 0) or 0
                 elif rate_limit_type == "total":
-                    total_tokens = usage.get("total_tokens", 0)
+                    total_tokens = usage.get("total_tokens", 0) or 0
+
+                # Get cached tokens from dict
+                if rate_limit_type in ("input", "total"):
+                    prompt_details = usage.get("prompt_tokens_details") or {}
+                    if isinstance(prompt_details, dict):
+                        cached_tokens = prompt_details.get("cached_tokens", 0) or 0
+
+        # Subtract cached tokens for input/total (providers don't count them)
+        if cached_tokens > 0:
+            total_tokens = max(0, total_tokens - cached_tokens)
+
         return total_tokens
 
     async def _execute_token_increment_script(
@@ -1336,6 +1370,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
 
     def get_rate_limit_type(self) -> Literal["output", "input", "total"]:
         from litellm.proxy.proxy_server import general_settings
+
         specified_rate_limit_type = general_settings.get(
             "token_rate_limit_type", "total"
         )
@@ -1381,9 +1416,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             user_api_key_organization_id = standard_logging_metadata.get(
                 "user_api_key_org_id"
             )
-            user_api_key_end_user_id = kwargs.get("user") or standard_logging_metadata.get(
-                "user_api_key_end_user_id"
-            )
+            user_api_key_end_user_id = kwargs.get(
+                "user"
+            ) or standard_logging_metadata.get("user_api_key_end_user_id")
             model_group = get_model_group_from_litellm_kwargs(kwargs)
 
             # Get total tokens from response
@@ -1393,7 +1428,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
                 response_obj, BaseLiteLLMOpenAIResponseObject
             ):
                 _usage = getattr(response_obj, "usage", None)
-                total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type)
+                total_tokens = self._get_total_tokens_from_usage(
+                    usage=_usage, rate_limit_type=rate_limit_type
+                )
 
             # Create pipeline operations for TPM increments
             pipeline_operations: List[RedisPipelineIncrementOperation] = []
@@ -1518,9 +1555,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         from litellm.types.caching import RedisPipelineIncrementOperation
 
         try:
-            litellm_parent_otel_span: Union[
-                Span, None
-            ] = _get_parent_otel_span_from_kwargs(kwargs)
+            litellm_parent_otel_span: Union[Span, None] = (
+                _get_parent_otel_span_from_kwargs(kwargs)
+            )
             # Get metadata from standard_logging_object - this correctly handles both
             # 'metadata' and 'litellm_metadata' fields from litellm_params
             standard_logging_object = kwargs.get("standard_logging_object") or {}
@@ -1555,7 +1592,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
                 f"Error in rate limit failure event: {str(e)}"
             )
 
-
     async def async_post_call_success_hook(
         self, data: dict, user_api_key_dict: UserAPIKeyAuth, response
     ):

tests/test_litellm/proxy/auth/test_auth_utils.py

@@ -0,0 +1,131 @@
+"""
+Unit tests for auth_utils functions related to rate limiting.
+"""
+
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.auth.auth_utils import (
+    get_key_model_rpm_limit,
+    get_key_model_tpm_limit,
+)
+
+
+class TestGetKeyModelRpmLimit:
+    """Tests for get_key_model_rpm_limit function."""
+
+    def test_returns_key_metadata_when_present(self):
+        """Key metadata takes priority over team metadata."""
+        user_api_key_dict = UserAPIKeyAuth(
+            api_key="sk-123",
+            metadata={"model_rpm_limit": {"gpt-4": 100}},
+            team_metadata={"model_rpm_limit": {"gpt-4": 50}},
+        )
+        result = get_key_model_rpm_limit(user_api_key_dict)
+        assert result == {"gpt-4": 100}
+
+    def test_falls_back_to_team_metadata_when_key_has_other_metadata(self):
+        """Should fall back to team metadata when key metadata exists but has no model_rpm_limit."""
+        user_api_key_dict = UserAPIKeyAuth(
+            api_key="sk-123",
+            metadata={
+                "some_other_key": "value"
+            },  # Has metadata, but not model_rpm_limit
+            team_metadata={"model_rpm_limit": {"gpt-4": 50}},
+        )
+        result = get_key_model_rpm_limit(user_api_key_dict)
+        assert result == {"gpt-4": 50}
+
+    def test_extracts_from_model_max_budget(self):
+        """Should extract rpm_limit from model_max_budget when metadata is empty."""
+        user_api_key_dict = UserAPIKeyAuth(
+            api_key="sk-123",
+            model_max_budget={
+                "gpt-4": {"rpm_limit": 100, "tpm_limit": 1000},
+                "gpt-3.5-turbo": {"rpm_limit": 200},
+            },
+        )
+        result = get_key_model_rpm_limit(user_api_key_dict)
+        assert result == {"gpt-4": 100, "gpt-3.5-turbo": 200}
+
+    def test_skips_models_without_rpm_limit(self):
+        """Should skip models that don't have rpm_limit in model_max_budget."""
+        user_api_key_dict = UserAPIKeyAuth(
+            api_key="sk-123",
+            model_max_budget={
+                "gpt-4": {"rpm_limit": 100},
+                "gpt-3.5-turbo": {"tpm_limit": 1000},  # No rpm_limit
+            },
+        )
+        result = get_key_model_rpm_limit(user_api_key_dict)
+        assert result == {"gpt-4": 100}
+
+    def test_returns_none_when_no_limits_configured(self):
+        """Should return None when no rate limits are configured."""
+        user_api_key_dict = UserAPIKeyAuth(api_key="sk-123")
+        result = get_key_model_rpm_limit(user_api_key_dict)
+        assert result is None
+
+
+class TestGetKeyModelTpmLimit:
+    """Tests for get_key_model_tpm_limit function."""
+
+    def test_returns_key_metadata_when_present(self):
+        """Key metadata takes priority over team metadata."""
+        user_api_key_dict = UserAPIKeyAuth(
+            api_key="sk-123",
+            metadata={"model_tpm_limit": {"gpt-4": 10000}},
+            team_metadata={"model_tpm_limit": {"gpt-4": 5000}},
+        )
+        result = get_key_model_tpm_limit(user_api_key_dict)
+        assert result == {"gpt-4": 10000}
+
+    def test_falls_back_to_team_metadata_when_key_has_other_metadata(self):
+        """Should fall back to team metadata when key metadata exists but has no model_tpm_limit."""
+        user_api_key_dict = UserAPIKeyAuth(
+            api_key="sk-123",
+            metadata={
+                "some_other_key": "value"
+            },  # Has metadata, but not model_tpm_limit
+            team_metadata={"model_tpm_limit": {"gpt-4": 5000}},
+        )
+        result = get_key_model_tpm_limit(user_api_key_dict)
+        assert result == {"gpt-4": 5000}
+
+    def test_extracts_from_model_max_budget(self):
+        """Should extract tpm_limit from model_max_budget when metadata is empty."""
+        user_api_key_dict = UserAPIKeyAuth(
+            api_key="sk-123",
+            model_max_budget={
+                "gpt-4": {"tpm_limit": 10000, "rpm_limit": 100},
+                "gpt-3.5-turbo": {"tpm_limit": 20000},
+            },
+        )
+        result = get_key_model_tpm_limit(user_api_key_dict)
+        assert result == {"gpt-4": 10000, "gpt-3.5-turbo": 20000}
+
+    def test_skips_models_without_tpm_limit(self):
+        """Should skip models that don't have tpm_limit in model_max_budget."""
+        user_api_key_dict = UserAPIKeyAuth(
+            api_key="sk-123",
+            model_max_budget={
+                "gpt-4": {"tpm_limit": 10000},
+                "gpt-3.5-turbo": {"rpm_limit": 100},  # No tpm_limit
+            },
+        )
+        result = get_key_model_tpm_limit(user_api_key_dict)
+        assert result == {"gpt-4": 10000}
+
+    def test_returns_none_when_no_limits_configured(self):
+        """Should return None when no rate limits are configured."""
+        user_api_key_dict = UserAPIKeyAuth(api_key="sk-123")
+        result = get_key_model_tpm_limit(user_api_key_dict)
+        assert result is None
+
+    def test_model_max_budget_priority_over_team(self):
+        """model_max_budget should take priority over team_metadata."""
+        user_api_key_dict = UserAPIKeyAuth(
+            api_key="sk-123",
+            model_max_budget={"gpt-4": {"tpm_limit": 10000}},
+            team_metadata={"model_tpm_limit": {"gpt-4": 5000}},
+        )
+        result = get_key_model_tpm_limit(user_api_key_dict)
+        assert result == {"gpt-4": 10000}