chore: sync upstream 2026-07-06

litellm/litellm_core_utils/get_litellm_params.py

@@ -36,6 +36,8 @@
         "aws_bedrock_project_id",
         "tpm",
         "rpm",
+        "itpm",
+        "otpm",
         "use_xai_oauth",
     }
 )

litellm/proxy/common_request_processing.py

@@ -52,6 +52,7 @@
 from litellm.proxy.route_llm_request import route_request
 from litellm.proxy.utils import ProxyLogging
 from litellm.router import Router
+from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.router import RouterRateLimitError
 from litellm.types.utils import ServerToolUse
@@ -600,7 +601,7 @@ def _override_openai_response_model(
     if not requested_model:
         return
 
-    hidden_params = getattr(response_obj, "_hidden_params", {}) or {}
+    hidden_params = get_hidden_params_dict(response_obj)
     if isinstance(hidden_params, dict):
         # Check if a fallback occurred - if so, preserve the actual model used
         fallback_headers = hidden_params.get("additional_headers", {}) or {}
@@ -900,7 +901,7 @@ async def build_litellm_proxy_success_headers_from_llm_response(
         (e.g. Google native :generateContent) instead of base_process_llm_request.
         """
         if isinstance(response, dict):
-            hidden_params = response.get("_hidden_params") or {}
+            hidden_params = get_hidden_params_dict(response)
         else:
             hidden_params = getattr(response, "_hidden_params", None) or {}
         if not isinstance(hidden_params, dict):
@@ -1433,7 +1434,7 @@ async def base_process_llm_request(
 
         _exception_raised = False
         try:
-            hidden_params = getattr(response, "_hidden_params", {}) or {}
+            hidden_params = get_hidden_params_dict(response)
             model_id = self._get_model_id_from_response(hidden_params, self.data)
 
             cache_key, api_base, response_cost = (
@@ -1708,7 +1709,7 @@ async def _on_deferred_stream_complete(assembled_response, cache_hit):
                 log_context=f"litellm_call_id={logging_obj.litellm_call_id}",
             )
 
-        hidden_params = getattr(response, "_hidden_params", {}) or {}  # get any updated response headers
+        hidden_params = get_hidden_params_dict(response)  # get any updated response headers
         additional_headers = hidden_params.get("additional_headers", {}) or {}
 
         recover_response_cost = not response_cost and hidden_params.get("response_cost") is None
@@ -1736,6 +1737,9 @@ async def _on_deferred_stream_complete(assembled_response, cache_hit):
             )
         )
 
+        if isinstance(response, dict):
+            response.pop("_hidden_params", None)
+
         # Call response headers hook for non-streaming success
         callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
             data=self.data,

litellm/router.py

@@ -90,7 +90,12 @@
     _HiddenParamsHost,
     add_fallback_headers_to_response,
     add_retry_headers_to_response,
+    apply_quality_router_decision_headers,
+    apply_remaining_usage_headers,
+    ensure_response_additional_headers,
     get_hidden_params_dict,
+    prepare_response_for_header_attachment,
+    response_in_flight_token_count,
 )
 from litellm.router_utils.batch_utils import (
     _get_router_metadata_variable_name,
@@ -133,6 +138,12 @@
 from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
     ModelRateLimitingCheck,
 )
+from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
+    build_io_token_rate_limit_headers,
+    deployment_has_io_token_limits,
+    refund_stale_reservation_before_retry,
+    set_io_token_rate_limit_request_kwargs,
+)
 from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import (
     PromptCachingDeploymentCheck,
 )
@@ -1649,6 +1660,7 @@ def _completion(
                 )
                 thread.start()
 
+            kwargs.setdefault("messages", messages)
             self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
             kwargs.pop("silent_model", None)  # Ensure it's not in kwargs either
             model_name = litellm_params["model"]
@@ -2672,6 +2684,7 @@ async def _acompletion(
                     )
                 )
 
+            kwargs.setdefault("messages", messages)
             self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
             kwargs.pop("silent_model", None)  # Ensure it's not in kwargs either
 
@@ -2959,6 +2972,13 @@ def _update_kwargs_with_deployment(
             }
         )
 
+        # A retry/fallback reuses this same kwargs dict for the next deployment.
+        # Refund and clear any reservation the previous deployment attempt left
+        # here before it's wiped below, instead of relying on that attempt's
+        # (possibly still-pending) failure event to do it.
+        refund_stale_reservation_before_retry(self.cache, kwargs)
+        set_io_token_rate_limit_request_kwargs(kwargs)
+
         ## DEPLOYMENT-LEVEL TAGS
         deployment_tags = deployment.get("litellm_params", {}).get("tags")
         if deployment_tags:
@@ -6766,14 +6786,21 @@ async def deployment_callback_on_success(
                     deployment_id=id,
                 )
 
-                ## if all are none, return - no need to track current tpm/rpm usage for models with no tpm/rpm set
+                deployment_dict = deployment_info if isinstance(deployment_info, dict) else deployment_info.model_dump()
+                has_io_token_limits = deployment_has_io_token_limits(deployment_dict)
+
+                ## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are
+                ## set. IO deployments still record TPM/RPM usage here so TPM-aware
+                ## routing strategies see their real load in mixed model groups; their
+                ## itpm/otpm enforcement runs separately in ModelRateLimitingCheck.
                 if (
                     tpm is None
                     and rpm is None
                     and tpm_litellm_params is None
                     and rpm_litellm_params is None
                     and tpm_model_info is None
                     and rpm_model_info is None
+                    and not has_io_token_limits
                 ):
                     return
 
@@ -8610,6 +8637,8 @@ def _set_model_group_info(self, model_group: str, user_facing_model_group_name:
 
         total_tpm: Optional[int] = None
         total_rpm: Optional[int] = None
+        total_itpm: Optional[int] = None
+        total_otpm: Optional[int] = None
         configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None
         model_list = self.get_model_list(model_name=model_group)
         if model_list is None:
@@ -8650,6 +8679,18 @@ def _set_model_group_info(self, model_group: str, user_facing_model_group_name:
             if _deployment_rpm is None:
                 _deployment_rpm = model_info_dict.get("rpm", None)  # type: ignore
 
+            _deployment_itpm: Optional[int] = model.get("itpm")
+            if _deployment_itpm is None:
+                _deployment_itpm = model_litellm_params.get("itpm", None)
+            if _deployment_itpm is None:
+                _deployment_itpm = model_info_dict.get("itpm", None)
+
+            _deployment_otpm: Optional[int] = model.get("otpm")
+            if _deployment_otpm is None:
+                _deployment_otpm = model_litellm_params.get("otpm", None)
+            if _deployment_otpm is None:
+                _deployment_otpm = model_info_dict.get("otpm", None)
+
             # get model info
             try:
                 model_id = model_info_dict.get("id", None)
@@ -8790,6 +8831,16 @@ def _set_model_group_info(self, model_group: str, user_facing_model_group_name:
                 if total_rpm is None:
                     total_rpm = 0
                 total_rpm += _deployment_rpm  # type: ignore
+
+            if _deployment_itpm is not None:
+                if total_itpm is None:
+                    total_itpm = 0
+                total_itpm += _deployment_itpm
+
+            if _deployment_otpm is not None:
+                if total_otpm is None:
+                    total_otpm = 0
+                total_otpm += _deployment_otpm
         if model_group_info is not None:
             ## UPDATE WITH TOTAL TPM/RPM FOR MODEL GROUP
             if total_tpm is not None:
@@ -8798,6 +8849,12 @@ def _set_model_group_info(self, model_group: str, user_facing_model_group_name:
             if total_rpm is not None:
                 model_group_info.rpm = total_rpm
 
+            if total_itpm is not None:
+                model_group_info.itpm = total_itpm
+
+            if total_otpm is not None:
+                model_group_info.otpm = total_otpm
+
             ## UPDATE WITH CONFIGURABLE CLIENTSIDE AUTH PARAMS FOR MODEL GROUP
             if configurable_clientside_auth_params is not None:
                 model_group_info.configurable_clientside_auth_params = configurable_clientside_auth_params
@@ -8900,6 +8957,58 @@ async def get_model_group_usage(self, model_group: str) -> Tuple[Optional[int],
                     rpm_usage += t
         return tpm_usage, rpm_usage
 
+    async def get_model_group_io_token_usage(self, model_group: str) -> tuple[Optional[int], Optional[int]]:
+        """
+        Returns current ITPM/OTPM usage for a model group (sum across deployments).
+        """
+        dt = get_utc_datetime()
+        current_minute = dt.strftime("%H-%M")
+        itpm_keys: list[str] = []
+        otpm_keys: list[str] = []
+
+        model_list = self.get_model_list(model_name=model_group)
+        if model_list is None:
+            return None, None
+
+        for model in model_list:
+            model_id: Optional[str] = model.get("model_info", {}).get("id")
+            litellm_model: Optional[str] = model["litellm_params"].get("model")
+            if model_id is None or litellm_model is None:
+                continue
+            itpm_keys.append(
+                RouterCacheEnum.ITPM.value.format(
+                    id=model_id,
+                    model=litellm_model,
+                    current_minute=current_minute,
+                )
+            )
+            otpm_keys.append(
+                RouterCacheEnum.OTPM.value.format(
+                    id=model_id,
+                    model=litellm_model,
+                    current_minute=current_minute,
+                )
+            )
+
+        combined_values = await self.cache.async_batch_get_cache(keys=itpm_keys + otpm_keys)
+        if combined_values is None:
+            return None, None
+
+        itpm_values = combined_values[: len(itpm_keys)]
+        otpm_values = combined_values[len(itpm_keys) :]
+
+        total_itpm: Optional[int] = None
+        for value in itpm_values:
+            if isinstance(value, int):
+                total_itpm = (total_itpm or 0) + value
+
+        total_otpm: Optional[int] = None
+        for value in otpm_values:
+            if isinstance(value, int):
+                total_otpm = (total_otpm or 0) + value
+
+        return total_itpm, total_otpm
+
     @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
     def _cached_get_model_group_info(self, model_group: str) -> Optional[ModelGroupInfo]:
         """
@@ -8909,25 +9018,33 @@ def _cached_get_model_group_info(self, model_group: str) -> Optional[ModelGroupI
         """
         return self.get_model_group_info(model_group)
 
-    async def get_remaining_model_group_usage(self, model_group: str) -> Dict[str, int]:
+    async def get_remaining_model_group_usage(self, model_group: str) -> dict[str, int]:
         model_group_info = self._cached_get_model_group_info(model_group)
 
-        if model_group_info is not None and model_group_info.tpm is not None:
-            tpm_limit = model_group_info.tpm
-        else:
-            tpm_limit = None
+        returned_dict: dict[str, int] = {}
 
-        if model_group_info is not None and model_group_info.rpm is not None:
-            rpm_limit = model_group_info.rpm
-        else:
-            rpm_limit = None
+        # ITPM/OTPM groups emit input/output token headers, but they may also set
+        # tpm/rpm, so build both sets rather than returning early - clients and
+        # prometheus gauges that read the standard headers still get data.
+        if model_group_info is not None and (model_group_info.itpm is not None or model_group_info.otpm is not None):
+            current_itpm, current_otpm = await self.get_model_group_io_token_usage(model_group)
+            returned_dict.update(
+                build_io_token_rate_limit_headers(
+                    itpm_limit=model_group_info.itpm,
+                    otpm_limit=model_group_info.otpm,
+                    current_itpm=current_itpm,
+                    current_otpm=current_otpm,
+                )
+            )
+
+        tpm_limit = model_group_info.tpm if model_group_info is not None else None
+        rpm_limit = model_group_info.rpm if model_group_info is not None else None
 
         if tpm_limit is None and rpm_limit is None:
-            return {}
+            return returned_dict
 
         current_tpm, current_rpm = await self.get_model_group_usage(model_group)
 
-        returned_dict = {}
         if tpm_limit is not None:
             returned_dict["x-ratelimit-remaining-tokens"] = tpm_limit - (current_tpm or 0)
             returned_dict["x-ratelimit-limit-tokens"] = tpm_limit
@@ -8950,69 +9067,25 @@ async def set_response_headers(
         # - if healthy_deployments > 1, return model group rate limit headers
         # - else return the model's rate limit headers
         """
-        if response is not None and hasattr(response, "_hidden_params"):
-            hidden_params = getattr(response, "_hidden_params", {}) or {}
-            if hasattr(hidden_params, "model_dump"):
-                hidden_params = hidden_params.model_dump()
-            if not isinstance(hidden_params, dict):
-                return response
-            response._hidden_params = hidden_params
-
-            additional_headers = hidden_params.get("additional_headers")
-            if not isinstance(additional_headers, dict):
-                additional_headers = {}
-                hidden_params["additional_headers"] = additional_headers
-            additional_headers["x-litellm-model-group"] = model_group
-
-            # Lift QualityRouter routing decision into response headers for
-            # transparency. The decision is stashed in request_kwargs.metadata
-            # by QualityRouter.async_pre_routing_hook.
-            metadata = (request_kwargs.get("metadata") or {}) if isinstance(request_kwargs, dict) else {}
-            decision = metadata.get("quality_router_decision") if isinstance(metadata, dict) else None
-            if isinstance(decision, dict):
-                # Only emit headers for fields that have a meaningful value.
-                # `complexity_tier` and `matched_keyword` are mutually exclusive
-                # (the keyword path short-circuits classification), so each
-                # request emits one or the other but not both.
-                if decision.get("routed_model") is not None:
-                    additional_headers["x-litellm-quality-router-model"] = str(decision["routed_model"])
-                if decision.get("quality_tier") is not None:
-                    additional_headers["x-litellm-quality-router-tier"] = str(decision["quality_tier"])
-                if decision.get("routed_via") is not None:
-                    additional_headers["x-litellm-quality-router-via"] = str(decision["routed_via"])
-                if decision.get("matched_keyword") is not None:
-                    additional_headers["x-litellm-quality-router-keyword"] = str(decision["matched_keyword"])
-                if decision.get("complexity_tier") is not None:
-                    additional_headers["x-litellm-quality-router-complexity"] = str(decision["complexity_tier"])
-
-            if (
-                "x-ratelimit-remaining-tokens" not in additional_headers
-                and "x-ratelimit-remaining-requests" not in additional_headers
-                and model_group is not None
-            ):
-                remaining_usage = await self.get_remaining_model_group_usage(model_group)
-
-                # get_remaining_model_group_usage reads the router's TPM/RPM
-                # counter, which is incremented post-response by
-                # deployment_callback_on_success. So the values returned here
-                # are pre-decrement for the current request, while vendor
-                # headers (OpenAI/Anthropic/Azure) are post-decrement. Replay
-                # the in-flight increment so router-derived headers match
-                # vendor-derived semantics — for both the HTTP response sent
-                # to the client and the prometheus gauges that read these
-                # headers downstream (LIT-2719).
-                in_flight_tokens = 0
-                usage = getattr(response, "usage", None)
-                if usage is not None:
-                    in_flight_tokens = getattr(usage, "total_tokens", 0) or 0
-                in_flight_delta = {
-                    "x-ratelimit-remaining-tokens": in_flight_tokens,
-                    "x-ratelimit-remaining-requests": 1,
-                }
+        response = prepare_response_for_header_attachment(response)
+        if response is None:
+            return response
 
-                for header, value in remaining_usage.items():
-                    if value is not None:
-                        additional_headers[header] = value - in_flight_delta.get(header, 0)
+        additional_headers = ensure_response_additional_headers(response)
+        additional_headers["x-litellm-model-group"] = model_group
+        apply_quality_router_decision_headers(additional_headers, request_kwargs)
+
+        if model_group is not None:
+            remaining_usage = await self.get_remaining_model_group_usage(model_group)
+            # get_remaining_model_group_usage reads the router's TPM/RPM counter,
+            # which is incremented post-response by deployment_callback_on_success.
+            # Replay the in-flight increment for TPM/RPM only (LIT-2719); ITPM/OTPM
+            # counters are incremented at reservation time and must not be adjusted.
+            apply_remaining_usage_headers(
+                additional_headers,
+                remaining_usage,
+                response_in_flight_token_count(response),
+            )
         return response
 
     def _build_model_name_index(self, model_list: list) -> None: