diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index fbed9594a0b..d505cbaf1e0 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -36,6 +36,8 @@ "aws_bedrock_project_id", "tpm", "rpm", + "itpm", + "otpm", "use_xai_oauth", } ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ec79080f7e2..8f73ce58a9c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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, diff --git a/litellm/router.py b/litellm/router.py index 22e9887a497..0306a36ae01 100644 --- a/litellm/router.py +++ b/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,7 +6786,13 @@ 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 @@ -6774,6 +6800,7 @@ async def deployment_callback_on_success( 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: diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index 0b927714ca9..e2204e18a3c 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -1,5 +1,5 @@ import json -from typing import Protocol, TypedDict, cast +from typing import Any, Protocol, TypedDict, cast from pydantic import BaseModel @@ -15,8 +15,98 @@ class _HiddenParamsHost(Protocol): _hidden_params: dict[str, object] -def get_hidden_params_dict(response: object) -> dict[str, object]: - hidden_params: object = cast(object, getattr(response, "_hidden_params", None)) +class HiddenParamsAsyncIteratorWrapper: + """ + Wraps a bare async generator/iterator (e.g. a provider's raw SSE + streaming response) that cannot itself hold a ``_hidden_params`` + attribute, so router-derived headers (ITPM/OTPM, model-group, retry, + fallback) can attach to a streaming response the same way they attach + to object-based responses (e.g. ``CustomStreamWrapper``). + """ + + def __init__(self, inner: object) -> None: + self._inner = inner + self._hidden_params: dict[str, object] = {} + + def __aiter__(self) -> "HiddenParamsAsyncIteratorWrapper": + return self + + async def __anext__(self) -> object: + return await cast(Any, self._inner).__anext__() + + async def aclose(self) -> None: + aclose = getattr(self._inner, "aclose", None) + if callable(aclose): + await aclose() + + +def prepare_response_for_header_attachment(response: object) -> object | None: + if response is None: + return None + if isinstance(response, dict) or hasattr(response, "_hidden_params"): + return response + if hasattr(response, "__anext__"): + return HiddenParamsAsyncIteratorWrapper(response) + return response + + +def ensure_response_additional_headers(response: object) -> dict[str, object]: + hidden_params = get_hidden_params_dict(response, create=isinstance(response, dict)) + _write_hidden_params(response, hidden_params) + additional_headers = hidden_params.get("additional_headers") + if not isinstance(additional_headers, dict): + additional_headers = {} + hidden_params["additional_headers"] = additional_headers + return additional_headers + + +def apply_quality_router_decision_headers( + additional_headers: dict[str, object], + request_kwargs: object, +) -> None: + 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 not isinstance(decision, dict): + return + quality_header_fields = ( + ("routed_model", "x-litellm-quality-router-model"), + ("quality_tier", "x-litellm-quality-router-tier"), + ("routed_via", "x-litellm-quality-router-via"), + ("matched_keyword", "x-litellm-quality-router-keyword"), + ("complexity_tier", "x-litellm-quality-router-complexity"), + ) + for field, header in quality_header_fields: + if decision.get(field) is not None: + additional_headers[header] = str(decision[field]) + + +def response_in_flight_token_count(response: object) -> int: + usage = response.get("usage") if isinstance(response, dict) else getattr(response, "usage", None) + if usage is None: + return 0 + if isinstance(usage, dict): + total = int(usage.get("total_tokens") or 0) + if total: + return total + return int(usage.get("input_tokens") or 0) + int(usage.get("output_tokens") or 0) + return int(getattr(usage, "total_tokens", 0) or 0) + + +def apply_remaining_usage_headers( + additional_headers: dict[str, object], + remaining_usage: dict[str, int], + in_flight_tokens: int, +) -> None: + in_flight_delta = { + "x-ratelimit-remaining-tokens": in_flight_tokens, + "x-ratelimit-remaining-requests": 1, + } + for header, value in remaining_usage.items(): + if value is not None and header not in additional_headers: + additional_headers[header] = value - in_flight_delta.get(header, 0) + + +def _normalize_hidden_params(hidden_params: object) -> dict[str, object]: if isinstance(hidden_params, BaseModel): return cast("dict[str, object]", hidden_params.model_dump()) if isinstance(hidden_params, dict): @@ -24,6 +114,29 @@ def get_hidden_params_dict(response: object) -> dict[str, object]: return {} +def get_hidden_params_dict( + response: object, + *, + create: bool = False, +) -> dict[str, object]: + if isinstance(response, dict): + hidden_params = _normalize_hidden_params(response.get("_hidden_params")) + if not hidden_params and create: + hidden_params = {} + response["_hidden_params"] = hidden_params + return hidden_params + + hidden_params = _normalize_hidden_params(cast(object, getattr(response, "_hidden_params", None))) + return hidden_params + + +def _write_hidden_params(response: object, hidden_params: dict[str, object]) -> None: + if isinstance(response, dict): + response["_hidden_params"] = hidden_params + elif hasattr(response, "_hidden_params"): + cast(_HiddenParamsHost, response)._hidden_params = hidden_params + + def _ensure_additional_headers_dict( hidden_params: dict[str, object], ) -> dict[str, object]: @@ -73,15 +186,19 @@ def _add_headers_to_response(response: object, headers: dict[str, object]) -> ob if response is None: return response - if not isinstance(response, BaseModel) and not hasattr(response, "_hidden_params"): + if ( + not isinstance(response, BaseModel) + and not isinstance(response, dict) + and not hasattr(response, "_hidden_params") + ): return response - hidden_params = get_hidden_params_dict(response) + hidden_params = get_hidden_params_dict(response, create=isinstance(response, dict)) additional_headers = _ensure_additional_headers_dict(hidden_params) additional_headers.update(headers) hidden_params["additional_headers"] = additional_headers - cast(_HiddenParamsHost, response)._hidden_params = hidden_params + _write_hidden_params(response, hidden_params) return response @@ -127,12 +244,12 @@ def add_fallback_headers_to_response( if fallback_errors is None or response is None: return response - hidden_params = get_hidden_params_dict(response) + hidden_params = get_hidden_params_dict(response, create=isinstance(response, dict)) additional_headers = _ensure_additional_headers_dict(hidden_params) merged_errors = get_fallback_errors_from_headers(additional_headers) + [ cast("dict[str, object]", error) for error in fallback_errors ] additional_headers["x-litellm-fallback-errors"] = json.dumps(merged_errors) hidden_params["additional_headers"] = additional_headers - cast(_HiddenParamsHost, response)._hidden_params = hidden_params + _write_hidden_params(response, hidden_params) return response diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py new file mode 100644 index 00000000000..62c99a1c9f6 --- /dev/null +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -0,0 +1,711 @@ +""" +Separate ITPM/OTPM (input/output tokens per minute) deployment rate limits. + +- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM +- Post-call: reconcile ITPM to actual input tokens and OTPM to actual output tokens +- Cached prompt-read tokens are excluded from ITPM post-call accounting + +Used by ModelRateLimitingCheck when a deployment sets itpm/otpm. +""" + +from __future__ import annotations + +import contextlib +import contextvars +from typing import TYPE_CHECKING, Any, Optional + +import httpx + +import litellm +from litellm import token_counter +from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache +from litellm.types.router import RouterCacheEnum, RouterErrors +from litellm.utils import get_utc_datetime + +if TYPE_CHECKING: + from opentelemetry.trace import Span as _Span + + Span = _Span | Any +else: + Span = Any + +RoutingArgsTTL = 60 + +_io_token_rate_limit_request_kwargs: contextvars.ContextVar[Optional[dict[str, Any]]] = contextvars.ContextVar( + "io_token_rate_limit_request_kwargs", + default=None, +) + +ITPM_RESERVED_KEY = "_litellm_itpm_reserved" +OTPM_RESERVED_KEY = "_litellm_otpm_reserved" +ITPM_CACHE_KEY = "_litellm_itpm_cache_key" +OTPM_CACHE_KEY = "_litellm_otpm_cache_key" + + +def set_io_token_rate_limit_request_kwargs(kwargs: Optional[dict[str, Any]]) -> None: + # The reservation sentinels are server-only, but `metadata` is caller + # controlled on proxy requests. Strip any client-supplied copies here (this + # runs before the router stashes its own reservation) so a forged + # reservation can't drive the post-call reconcile/refund against an + # arbitrary counter and bypass the configured limits. + _clear_reservation_from_kwargs(kwargs) + _io_token_rate_limit_request_kwargs.set(kwargs) + + +def get_io_token_rate_limit_request_kwargs() -> Optional[dict[str, Any]]: + return _io_token_rate_limit_request_kwargs.get() + + +def seconds_until_minute_reset() -> int: + dt = get_utc_datetime() + return max(1, 60 - dt.second) + + +def get_deployment_io_token_limits( + deployment: dict, +) -> tuple[Optional[int], Optional[int]]: + itpm = deployment.get("itpm") + otpm = deployment.get("otpm") + litellm_params = deployment.get("litellm_params") or {} + model_info = deployment.get("model_info") or {} + if itpm is None: + itpm = litellm_params.get("itpm") + if otpm is None: + otpm = litellm_params.get("otpm") + if itpm is None: + itpm = model_info.get("itpm") + if otpm is None: + otpm = model_info.get("otpm") + return itpm, otpm + + +def deployment_has_io_token_limits(deployment: dict) -> bool: + itpm, otpm = get_deployment_io_token_limits(deployment) + return itpm is not None or otpm is not None + + +def _get_cache_keys(deployment: dict, current_minute: str) -> Optional[tuple[str, str]]: + model_id = deployment.get("model_info", {}).get("id") + deployment_name = deployment.get("litellm_params", {}).get("model") + # Without both a deployment id and model name the key would collapse to a + # shared "None:None" bucket across misconfigured deployments, so bail out. + if model_id is None or deployment_name is None: + return None + itpm_key = RouterCacheEnum.ITPM.value.format(id=model_id, model=deployment_name, current_minute=current_minute) + otpm_key = RouterCacheEnum.OTPM.value.format(id=model_id, model=deployment_name, current_minute=current_minute) + return itpm_key, otpm_key + + +def _estimate_input_tokens(request_kwargs: Optional[dict[str, Any]], model: str = "") -> int: + if not request_kwargs: + return 0 + messages = request_kwargs.get("messages") + prompt = request_kwargs.get("prompt") + input_text = request_kwargs.get("input") + # token_counter can raise from any of its tokenizer backends; this is a + # best-effort estimate for the ITPM reservation and must never fail the + # underlying request. Passing the deployment model name uses a model-specific + # tokenizer when available, reducing the reservation over/under-estimate window + # between pre-call and post-call reconcile. + with contextlib.suppress(Exception): + return max(0, int(token_counter(model=model, messages=messages, text=prompt or input_text))) + return 0 + + +def _model_max_output_tokens(model_name: str) -> Optional[int]: + # litellm.get_model_info raises a bare Exception for an unrecognized model; + # this lookup is a fallback default and must never fail the request. + with contextlib.suppress(Exception): + info = litellm.get_model_info(model=model_name) + model_max = info.get("max_output_tokens") or info.get("max_tokens") + if model_max is not None: + return max(0, int(model_max)) + return None + + +def _resolve_max_tokens(request_kwargs: Optional[dict[str, Any]], deployment: dict) -> int: + if request_kwargs: + # An explicit max_tokens=0 must be honored, not treated as absent and + # replaced by the model default. + explicit = request_kwargs.get("max_tokens") + if explicit is None: + explicit = request_kwargs.get("max_completion_tokens") + if explicit is None: + explicit = request_kwargs.get("max_output_tokens") + if explicit is not None: + return max(0, int(explicit)) + + model_name = (deployment.get("litellm_params") or {}).get("model") + if model_name: + model_max = _model_max_output_tokens(model_name) + if model_max is not None: + return model_max + return 4096 + + +def _get_usage_tokens(usage: Any) -> tuple[int, int, int]: + if usage is None: + return 0, 0, 0 + if hasattr(usage, "prompt_tokens") or hasattr(usage, "input_tokens"): + prompt = int(getattr(usage, "prompt_tokens", None) or getattr(usage, "input_tokens", 0) or 0) + completion = int(getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", 0) or 0) + cached = 0 + details = getattr(usage, "prompt_tokens_details", None) + if details is not None: + cached = int(getattr(details, "cached_tokens", 0) or 0) + if not cached: + cached = int(getattr(usage, "cache_read_input_tokens", 0) or 0) + return prompt, completion, cached + if isinstance(usage, dict): + prompt = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0) + completion = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0) + details = usage.get("prompt_tokens_details") or {} + cached = int(details.get("cached_tokens", 0) or 0) if isinstance(details, dict) else 0 + if not cached: + cached = int(usage.get("cache_read_input_tokens") or 0) + return prompt, completion, cached + return 0, 0, 0 + + +def _extract_response_usage(response_obj: Any) -> Any: + if isinstance(response_obj, dict): + return response_obj.get("usage") + return getattr(response_obj, "usage", None) + + +def _usage_is_present(usage: Any) -> bool: + """ + True only if usage carries an actual input/output breakdown. + + ``total_tokens`` alone is deliberately excluded: ``_get_usage_tokens`` has + no way to split a bare total into input vs. output, so treating it as + "present" would resolve to (0, 0) and refund the full reservation as if + zero tokens were used. + """ + if usage is None: + return False + fields = ("prompt_tokens", "completion_tokens", "input_tokens", "output_tokens") + if isinstance(usage, dict): + return any(key in usage for key in fields) + return any(hasattr(usage, key) for key in fields) + + +def _resolve_reconcile_usage_tokens( + kwargs: Any, + response_obj: Any, +) -> tuple[int, int, bool]: + """ + Resolve billable input and output tokens for post-call reconcile. + + Prefer the response usage object; fall back to standard_logging_object token + fields. When usage cannot be resolved, return ``usage_resolved=False`` so + callers keep the pre-call reservation instead of refunding it as zero usage. + """ + usage = _extract_response_usage(response_obj) + if _usage_is_present(usage): + prompt_tokens, completion_tokens, cached_tokens = _get_usage_tokens(usage) + return max(0, prompt_tokens - cached_tokens), completion_tokens, True + + if isinstance(kwargs, dict): + standard_logging_object = kwargs.get("standard_logging_object") + if isinstance(standard_logging_object, dict): + prompt_tokens = int(standard_logging_object.get("prompt_tokens") or 0) + completion_tokens = int(standard_logging_object.get("completion_tokens") or 0) + cached_tokens = 0 + metadata = standard_logging_object.get("metadata") + if isinstance(metadata, dict): + cached_tokens = int(metadata.get("cache_read_input_tokens") or 0) + # Same rationale as _usage_is_present: a bare total_tokens with no + # prompt/completion breakdown can't be split, so it isn't treated + # as resolved usage - the reservation is kept instead of refunded. + if prompt_tokens or completion_tokens: + return max(0, prompt_tokens - cached_tokens), completion_tokens, True + + return 0, 0, False + + +def _stash_reservation_in_metadata( + request_kwargs: Optional[dict[str, Any]], + *, + itpm_reserved: int, + otpm_reserved: int, + itpm_cache_key: Optional[str], + otpm_cache_key: Optional[str], +) -> None: + if not request_kwargs: + return + reservation = { + ITPM_RESERVED_KEY: itpm_reserved, + OTPM_RESERVED_KEY: otpm_reserved, + ITPM_CACHE_KEY: itpm_cache_key, + OTPM_CACHE_KEY: otpm_cache_key, + } + for channel in ("metadata", "litellm_metadata"): + existing = request_kwargs.get(channel) + if isinstance(existing, dict): + existing.update(reservation) + elif channel == "metadata": + request_kwargs[channel] = dict(reservation) + + +def _extract_reservation(reservation: dict[str, Any]) -> tuple[int, int, Optional[str], Optional[str]]: + itpm_cache_key = reservation.get(ITPM_CACHE_KEY) + otpm_cache_key = reservation.get(OTPM_CACHE_KEY) + return ( + int(reservation.get(ITPM_RESERVED_KEY, 0) or 0), + int(reservation.get(OTPM_RESERVED_KEY, 0) or 0), + itpm_cache_key if isinstance(itpm_cache_key, str) else None, + otpm_cache_key if isinstance(otpm_cache_key, str) else None, + ) + + +def _reservation_channels(kwargs: Any) -> tuple[Any, ...]: + """ + Places a reservation may live, in priority order: the top-level metadata + channels win over litellm_params.metadata (so a top-level stash is never + shadowed), which win over the standard_logging_object copy. + """ + if not isinstance(kwargs, dict): + return () + channels = [kwargs.get("metadata"), kwargs.get("litellm_metadata")] + litellm_params = kwargs.get("litellm_params") + if isinstance(litellm_params, dict): + channels.append(litellm_params.get("metadata")) + standard_logging_object = kwargs.get("standard_logging_object") + if isinstance(standard_logging_object, dict): + channels.append(standard_logging_object.get("metadata")) + return tuple(channels) + + +def _read_reservation_from_kwargs(kwargs: Any) -> tuple[int, int, Optional[str], Optional[str]]: + for channel_dict in _reservation_channels(kwargs): + if isinstance(channel_dict, dict) and ITPM_RESERVED_KEY in channel_dict: + return _extract_reservation(channel_dict) + return 0, 0, None, None + + +def _clear_reservation_from_kwargs(kwargs: Any) -> None: + """ + Remove the stashed reservation so a retry on a different (e.g. non-IO) + deployment does not re-process the already-reconciled/refunded reservation. + """ + for channel_dict in _reservation_channels(kwargs): + if isinstance(channel_dict, dict): + for key in (ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY): + channel_dict.pop(key, None) + + +def _reservation_value(value: int, limit: Optional[int]) -> int: + if limit is None: + return 0 + if value > 0: + return value + # Estimation failed (empty messages, unsupported model, tokenizer error). + # Reserve a minimal 1-token slot rather than the full limit: the latter + # would let one request whose estimate failed fill the entire bucket, + # serializing every concurrent request to the deployment until it + # completes and reconciles against actual usage. + return 1 + + +def _rate_limit_error(limit_label: str, limit: int, current: float) -> litellm.RateLimitError: + return litellm.RateLimitError( + message=f"Model rate limit exceeded. {limit_label} limit={limit}, current usage={current}", + llm_provider="", + model="", + response=httpx.Response( + status_code=429, + content=( + f"{RouterErrors.user_defined_ratelimit_error.value} " + f"{limit_label} limit={limit}. current usage={current}." + ), + headers={"retry-after": str(RoutingArgsTTL)}, + request=httpx.Request( + method="io_token_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + num_retries=0, + ) + + +def _sync_increment_with_rollback( + dual_cache: DualCache, + key: str, + value: int, + limit: Optional[int], + *, + limit_label: str, +) -> None: + if value <= 0 or limit is None: + return + current = dual_cache.increment_cache( + key=key, + value=value, + ttl=RoutingArgsTTL, + ) + if current is not None and current > limit: + dual_cache.increment_cache( + key=key, + value=-value, + ttl=RoutingArgsTTL, + ) + raise _rate_limit_error(limit_label, limit, current) + + +async def _increment_with_rollback( + dual_cache: DualCache, + key: str, + value: int, + limit: Optional[int], + *, + parent_otel_span: Optional[Span] = None, + limit_label: str, +) -> None: + if value <= 0 or limit is None: + return + current = await dual_cache.async_increment_cache( + key=key, + value=value, + ttl=RoutingArgsTTL, + parent_otel_span=parent_otel_span, + ) + if current is not None and current > limit: + await dual_cache.async_increment_cache( + key=key, + value=-value, + ttl=RoutingArgsTTL, + parent_otel_span=parent_otel_span, + ) + raise _rate_limit_error(limit_label, limit, current) + + +def io_token_pre_call_check( + dual_cache: DualCache, + deployment: dict, +) -> Optional[dict]: + itpm_limit, otpm_limit = get_deployment_io_token_limits(deployment) + if itpm_limit is None and otpm_limit is None: + return deployment + + request_kwargs = get_io_token_rate_limit_request_kwargs() + _model = (deployment.get("litellm_params") or {}).get("model") or "" + estimated_input = _estimate_input_tokens(request_kwargs, model=_model) + max_tokens = _resolve_max_tokens(request_kwargs, deployment) + + dt = get_utc_datetime() + current_minute = dt.strftime("%H-%M") + cache_keys = _get_cache_keys(deployment, current_minute) + if cache_keys is None: + return deployment + itpm_key, otpm_key = cache_keys + + itpm_reserved = 0 + otpm_reserved = 0 + + if itpm_limit is not None: + itpm_reserved = _reservation_value(estimated_input, itpm_limit) + _sync_increment_with_rollback( + dual_cache, + itpm_key, + itpm_reserved, + itpm_limit, + limit_label="ITPM", + ) + + if otpm_limit is not None: + otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit) + try: + _sync_increment_with_rollback( + dual_cache, + otpm_key, + otpm_reserved, + otpm_limit, + limit_label="OTPM", + ) + except Exception: + if itpm_reserved > 0: + dual_cache.increment_cache( + key=itpm_key, + value=-itpm_reserved, + ttl=RoutingArgsTTL, + ) + raise + + _stash_reservation_in_metadata( + request_kwargs, + itpm_reserved=itpm_reserved, + otpm_reserved=otpm_reserved, + itpm_cache_key=itpm_key if itpm_limit is not None else None, + otpm_cache_key=otpm_key if otpm_limit is not None else None, + ) + return deployment + + +async def async_io_token_pre_call_check( + dual_cache: DualCache, + deployment: dict, + parent_otel_span: Optional[Span] = None, +) -> Optional[dict]: + itpm_limit, otpm_limit = get_deployment_io_token_limits(deployment) + if itpm_limit is None and otpm_limit is None: + return deployment + + request_kwargs = get_io_token_rate_limit_request_kwargs() + _model = (deployment.get("litellm_params") or {}).get("model") or "" + estimated_input = _estimate_input_tokens(request_kwargs, model=_model) + max_tokens = _resolve_max_tokens(request_kwargs, deployment) + + dt = get_utc_datetime() + current_minute = dt.strftime("%H-%M") + cache_keys = _get_cache_keys(deployment, current_minute) + if cache_keys is None: + return deployment + itpm_key, otpm_key = cache_keys + + itpm_reserved = 0 + otpm_reserved = 0 + + if itpm_limit is not None: + itpm_reserved = _reservation_value(estimated_input, itpm_limit) + await _increment_with_rollback( + dual_cache, + itpm_key, + itpm_reserved, + itpm_limit, + parent_otel_span=parent_otel_span, + limit_label="ITPM", + ) + + if otpm_limit is not None: + otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit) + try: + await _increment_with_rollback( + dual_cache, + otpm_key, + otpm_reserved, + otpm_limit, + parent_otel_span=parent_otel_span, + limit_label="OTPM", + ) + except Exception: + # Any failure reserving OTPM (a 429 or a transient cache error) must + # release the ITPM reservation already made, or it stays inflated + # until the TTL expires. + if itpm_reserved > 0: + await dual_cache.async_increment_cache( + key=itpm_key, + value=-itpm_reserved, + ttl=RoutingArgsTTL, + parent_otel_span=parent_otel_span, + ) + raise + + _stash_reservation_in_metadata( + request_kwargs, + itpm_reserved=itpm_reserved, + otpm_reserved=otpm_reserved, + itpm_cache_key=itpm_key if itpm_limit is not None else None, + otpm_cache_key=otpm_key if otpm_limit is not None else None, + ) + return deployment + + +def io_token_reconcile_success( + dual_cache: DualCache, + kwargs: Any, + response_obj: Any, +) -> None: + itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs) + if itpm_key is None and otpm_key is None: + return + + billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(kwargs, response_obj) + + try: + if usage_resolved: + if itpm_key is not None: + itpm_delta = billable_input - itpm_reserved + if itpm_delta != 0: + dual_cache.increment_cache( + key=itpm_key, + value=itpm_delta, + ttl=RoutingArgsTTL, + ) + + if otpm_key is not None: + otpm_delta = completion_tokens - otpm_reserved + if otpm_delta != 0: + dual_cache.increment_cache( + key=otpm_key, + value=otpm_delta, + ttl=RoutingArgsTTL, + ) + else: + verbose_router_logger.debug( + "[IO TOKEN LIMIT] usage missing; keeping reservation " + f"(itpm_reserved={itpm_reserved}, otpm_reserved={otpm_reserved})" + ) + finally: + _clear_reservation_from_kwargs(kwargs) + + verbose_router_logger.debug( + f"[IO TOKEN LIMIT] reconciled " + f"(usage_resolved={usage_resolved}, itpm_reserved={itpm_reserved}, " + f"billable_input={billable_input}, otpm_reserved={otpm_reserved}, output={completion_tokens})" + ) + + +async def async_io_token_reconcile_success( + dual_cache: DualCache, + kwargs: Any, + response_obj: Any, + *, + parent_otel_span: Optional[Span] = None, +) -> None: + itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs) + if itpm_key is None and otpm_key is None: + return + + billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(kwargs, response_obj) + + # Reconcile against the exact key that held the reservation (which encodes + # the reservation's minute), not a key recomputed at response time. This + # tracks actual usage even when the pre-call estimate was 0, and avoids a + # minute-boundary mismatch for calls that span into the next minute. Always + # clear the stash afterwards (even if an increment throws) so a retry or a + # duplicate success event can't re-process it. + try: + if usage_resolved: + if itpm_key is not None: + itpm_delta = billable_input - itpm_reserved + if itpm_delta != 0: + await dual_cache.async_increment_cache( + key=itpm_key, + value=itpm_delta, + ttl=RoutingArgsTTL, + parent_otel_span=parent_otel_span, + ) + + if otpm_key is not None: + otpm_delta = completion_tokens - otpm_reserved + if otpm_delta != 0: + await dual_cache.async_increment_cache( + key=otpm_key, + value=otpm_delta, + ttl=RoutingArgsTTL, + parent_otel_span=parent_otel_span, + ) + else: + verbose_router_logger.debug( + "[IO TOKEN LIMIT] usage missing; keeping reservation " + f"(itpm_reserved={itpm_reserved}, otpm_reserved={otpm_reserved})" + ) + finally: + _clear_reservation_from_kwargs(kwargs) + + verbose_router_logger.debug( + f"[IO TOKEN LIMIT] reconciled " + f"(usage_resolved={usage_resolved}, itpm_reserved={itpm_reserved}, " + f"billable_input={billable_input}, otpm_reserved={otpm_reserved}, output={completion_tokens})" + ) + + +def io_token_refund_failure( + dual_cache: DualCache, + kwargs: Any, +) -> None: + itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs) + if itpm_key is None and otpm_key is None: + return + if itpm_key is not None and itpm_reserved > 0: + dual_cache.increment_cache( + key=itpm_key, + value=-itpm_reserved, + ttl=RoutingArgsTTL, + ) + if otpm_key is not None and otpm_reserved > 0: + dual_cache.increment_cache( + key=otpm_key, + value=-otpm_reserved, + ttl=RoutingArgsTTL, + ) + _clear_reservation_from_kwargs(kwargs) + verbose_router_logger.debug(f"[IO TOKEN LIMIT] refunded ITPM={itpm_reserved} OTPM={otpm_reserved}") + + +def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Optional[dict[str, Any]]) -> None: + """ + Synchronously refund and clear any reservation a previous deployment + attempt stashed in ``kwargs``, before it's overwritten for the next + attempt (retry/fallback). + + ``set_io_token_rate_limit_request_kwargs`` strips reservation sentinels + from ``kwargs`` on every deployment pick (a security measure so a + caller-forged reservation can't be replayed). Without this refund, a + retry after a non-RateLimitError failure (e.g. an upstream 500) would + wipe deployment A's still-unreconciled reservation before its failure + event - which may be scheduled as a background task - gets a chance to + refund it, permanently stranding the reservation until its TTL expires + and causing false rate-limit errors for subsequent requests. + + ponytail: uses sync ``DualCache.increment_cache`` which issues a blocking + Redis INCR when a Redis backend is configured. This only triggers on + streaming mid-stream retries (non-streaming failures await their failure + handler before retrying, so the sentinels are already cleared). Upgrade + path: make ``_update_kwargs_with_deployment`` async and switch to + ``async_io_token_refund_failure`` — requires touching all callers. + """ + if not kwargs: + return + io_token_refund_failure(dual_cache, kwargs) + + +async def async_io_token_refund_failure( + dual_cache: DualCache, + kwargs: Any, + *, + parent_otel_span: Optional[Span] = None, +) -> None: + itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs) + if itpm_key is None and otpm_key is None: + return + if itpm_key is not None and itpm_reserved > 0: + await dual_cache.async_increment_cache( + key=itpm_key, + value=-itpm_reserved, + ttl=RoutingArgsTTL, + parent_otel_span=parent_otel_span, + ) + if otpm_key is not None and otpm_reserved > 0: + await dual_cache.async_increment_cache( + key=otpm_key, + value=-otpm_reserved, + ttl=RoutingArgsTTL, + parent_otel_span=parent_otel_span, + ) + _clear_reservation_from_kwargs(kwargs) + verbose_router_logger.debug(f"[IO TOKEN LIMIT] refunded ITPM={itpm_reserved} OTPM={otpm_reserved}") + + +def build_io_token_rate_limit_headers( + *, + itpm_limit: Optional[int], + otpm_limit: Optional[int], + current_itpm: Optional[int], + current_otpm: Optional[int], +) -> dict[str, int]: + headers: dict[str, int] = {} + reset = seconds_until_minute_reset() + if itpm_limit is not None: + usage = current_itpm or 0 + headers["x-ratelimit-limit-input-tokens"] = itpm_limit + headers["x-ratelimit-remaining-input-tokens"] = max(0, itpm_limit - usage) + headers["x-ratelimit-reset-input-tokens"] = reset + if otpm_limit is not None: + usage = current_otpm or 0 + headers["x-ratelimit-limit-output-tokens"] = otpm_limit + headers["x-ratelimit-remaining-output-tokens"] = max(0, otpm_limit - usage) + headers["x-ratelimit-reset-output-tokens"] = reset + return headers 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 0c6450b191f..5f907a70380 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 @@ -1,13 +1,16 @@ """ -Enforce TPM/RPM rate limits set on model deployments. +Enforce TPM/RPM or separate ITPM/OTPM 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 router_settings.optional_pre_call_checks: ["enforce_model_rate_limits"] -When enabled via `enforce_model_rate_limits: true` in litellm_settings, -requests that exceed the configured TPM/RPM limits will receive a 429 error. +- tpm/rpm: combined TPM + optional RPM (legacy) +- itpm/otpm: separate input/output tokens per minute + +When a deployment sets both itpm/otpm and tpm/rpm, both are enforced. A warning +is logged the first time such a deployment is seen. """ +import contextlib from typing import TYPE_CHECKING, Any, Dict, Optional, Union import httpx @@ -16,6 +19,17 @@ from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + ITPM_RESERVED_KEY, + async_io_token_pre_call_check, + async_io_token_reconcile_success, + async_io_token_refund_failure, + deployment_has_io_token_limits, + get_io_token_rate_limit_request_kwargs, + io_token_pre_call_check, + io_token_reconcile_success, + io_token_refund_failure, +) from litellm.types.router import RouterErrors from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_utc_datetime @@ -34,7 +48,7 @@ class RoutingArgs: class ModelRateLimitingCheck(CustomLogger): """ - Pre-call check that enforces TPM/RPM limits on model deployments. + Pre-call check that enforces TPM/RPM or ITPM/OTPM 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. @@ -45,6 +59,42 @@ class ModelRateLimitingCheck(CustomLogger): def __init__(self, dual_cache: DualCache): self.dual_cache = dual_cache + # model_ids already warned about itpm/otpm + tpm/rpm on the same deployment, + # so the warning is logged once per deployment rather than per request. + self._io_token_conflict_warned_ids: set[str] = set() + + def _warn_io_token_and_tpm_rpm_coexist_once(self, deployment: dict) -> None: + tpm_limit, rpm_limit = self._get_deployment_limits(deployment) + if tpm_limit is None and rpm_limit is None: + return + model_id = deployment.get("model_info", {}).get("id") + # Dedup per deployment id; if there is no id (degenerate config) don't + # collapse every such deployment onto one key - warn each time instead. + if model_id is not None: + if model_id in self._io_token_conflict_warned_ids: + return + self._io_token_conflict_warned_ids.add(str(model_id)) + verbose_router_logger.warning( + f"Deployment '{model_id}' configures itpm/otpm alongside tpm/rpm; " + "both limit types are enforced on this deployment" + ) + + def _refund_io_token_reservation_if_any(self) -> None: + request_kwargs = get_io_token_rate_limit_request_kwargs() + if request_kwargs is not None: + io_token_refund_failure(self.dual_cache, request_kwargs) + + async def _async_refund_io_token_reservation_if_any( + self, + parent_otel_span: Optional[Span] = None, + ) -> None: + request_kwargs = get_io_token_rate_limit_request_kwargs() + if request_kwargs is not None: + await async_io_token_refund_failure( + self.dual_cache, + request_kwargs, + parent_otel_span=parent_otel_span, + ) def _get_deployment_limits(self, deployment: Dict) -> tuple[Optional[int], Optional[int]]: """ @@ -93,6 +143,15 @@ def pre_call_check(self, deployment: Dict) -> Optional[Dict]: Raises RateLimitError if deployment exceeds TPM/RPM limits. """ try: + io_reservation_made = False + if deployment_has_io_token_limits(deployment): + self._warn_io_token_and_tpm_rpm_coexist_once(deployment) + io_token_pre_call_check( + self.dual_cache, + deployment, + ) + io_reservation_made = True + tpm_limit, rpm_limit = self._get_deployment_limits(deployment) # If no limits are set, allow the request @@ -149,6 +208,8 @@ def pre_call_check(self, deployment: Dict) -> Optional[Dict]: return deployment except litellm.RateLimitError: + if io_reservation_made: + self._refund_io_token_reservation_if_any() raise except Exception as e: verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.pre_call_check: {str(e)}") @@ -159,9 +220,19 @@ async def async_pre_call_check(self, deployment: Dict, parent_otel_span: Optiona """ Async pre-call check for model rate limits. - Raises RateLimitError if deployment exceeds TPM/RPM limits. + Raises RateLimitError if deployment exceeds TPM/RPM or ITPM/OTPM limits. """ try: + io_reservation_made = False + if deployment_has_io_token_limits(deployment): + self._warn_io_token_and_tpm_rpm_coexist_once(deployment) + await async_io_token_pre_call_check( + self.dual_cache, + deployment, + parent_otel_span=parent_otel_span, + ) + io_reservation_made = True + tpm_limit, rpm_limit = self._get_deployment_limits(deployment) # If no limits are set, allow the request @@ -225,6 +296,8 @@ async def async_pre_call_check(self, deployment: Dict, parent_otel_span: Optiona return deployment except litellm.RateLimitError: + if io_reservation_made: + await self._async_refund_io_token_reservation_if_any(parent_otel_span=parent_otel_span) raise except Exception as e: verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.async_pre_call_check: {str(e)}") @@ -232,14 +305,31 @@ async def async_pre_call_check(self, deployment: Dict, parent_otel_span: Optiona return deployment async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - """ - Track TPM usage after successful request. + from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + ) - 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") + + # IO token reconciliation works purely from the cache keys stashed in + # kwargs/metadata, so it must run before the model_id guard below + # (which only the TPM-tracking path needs). Otherwise a request whose + # standard_logging_object lacks model_id would never return its + # reservation, leaving the counter elevated until the TTL expires. + slo_metadata = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {} + kwargs_metadata = kwargs.get("metadata") or {} + if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata: + await async_io_token_reconcile_success( + self.dual_cache, + kwargs, + response_obj, + parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), + ) + # Fall through: a deployment can also configure tpm/rpm alongside + # itpm/otpm, and that path's pre-call check reads the tpm_key + # counter tracked below, so it must still be incremented here. + if standard_logging_object is None: return @@ -272,6 +362,19 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti except Exception as e: verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.async_log_success_event: {str(e)}") + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + ) + + # Never fail the primary logging pipeline over an io-token refund error. + with contextlib.suppress(Exception): + await async_io_token_refund_failure( + self.dual_cache, + kwargs, + parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), + ) + def log_success_event(self, kwargs, response_obj, start_time, end_time): """ Sync version of tracking TPM usage after successful request. @@ -279,6 +382,18 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): """ try: standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") + slo_metadata = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {} + kwargs_metadata = kwargs.get("metadata") or {} + if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata: + io_token_reconcile_success( + self.dual_cache, + kwargs, + response_obj, + ) + # Fall through: a deployment can also configure tpm/rpm alongside + # itpm/otpm, and that path's pre-call check reads the tpm_key + # counter tracked below, so it must still be incremented here. + if standard_logging_object is None: return @@ -304,3 +419,10 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): except Exception as e: verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.log_success_event: {str(e)}") + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + with contextlib.suppress(Exception): + io_token_refund_failure( + self.dual_cache, + kwargs, + ) diff --git a/litellm/types/router.py b/litellm/types/router.py index 0c3485deae7..4bac9358392 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -214,6 +214,8 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): custom_llm_provider: Optional[str] = None tpm: Optional[int] = None rpm: Optional[int] = None + itpm: Optional[int] = None + otpm: 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/ @@ -359,6 +361,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): custom_llm_provider: Optional[str] tpm: Optional[int] rpm: Optional[int] + itpm: Optional[int] + otpm: Optional[int] order: Optional[int] weight: Optional[int] max_parallel_requests: Optional[int] @@ -552,6 +556,8 @@ class ModelGroupInfo(BaseModel): ] = Field(default="chat") tpm: Optional[int] = None rpm: Optional[int] = None + itpm: Optional[int] = None + otpm: Optional[int] = None supports_parallel_function_calling: bool = Field(default=False) supports_vision: bool = Field(default=False) supports_web_search: bool = Field(default=False) @@ -749,6 +755,8 @@ class RoutingStrategy(enum.Enum): class RouterCacheEnum(enum.Enum): TPM = "global_router:{id}:{model}:tpm:{current_minute}" RPM = "global_router:{id}:{model}:rpm:{current_minute}" + ITPM = "global_router:{id}:{model}:itpm:{current_minute}" + OTPM = "global_router:{id}:{model}:otpm:{current_minute}" class GenericBudgetWindowDetails(BaseModel): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2e0eceaddd9..380621f88a8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3120,6 +3120,8 @@ class CustomPricingLiteLLMParams(BaseModel): "client", "rpm", "tpm", + "itpm", + "otpm", "max_parallel_requests", "input_cost_per_token", "output_cost_per_token", diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 674cf738b3a..fbdae0bcbf6 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -324,7 +324,7 @@ "limit": 883 }, "UP006": { - "limit": 12870 + "limit": 12869 }, "UP007": { "limit": 2571 diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 83d6d56df4f..848a6c28a57 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -523,6 +523,55 @@ async def test_deployment_callback_on_success(sync_mode): assert tpm_key is not None +@pytest.mark.asyncio +async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment(): + """ + An IO-limited deployment (itpm/otpm, no tpm/rpm) must still record TPM usage + in the router's routing counter so TPM-aware routing strategies see its real + load in mixed model groups; its itpm/otpm enforcement runs separately. + """ + import time + + model_list = [ + { + "model_name": "opus", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fake", + "itpm": 1000, + }, + "model_info": {"id": "io-100"}, + } + ] + router = Router(model_list=model_list) + + standard_logging_payload = create_standard_logging_payload() + standard_logging_payload["total_tokens"] = 100 + standard_logging_payload["model_id"] = "io-100" + kwargs = { + "litellm_params": { + "metadata": { + "deployment": "openai/gpt-4o-mini", + "model_group": "opus", + }, + "model_info": {"id": "io-100"}, + }, + "standard_logging_object": standard_logging_payload, + } + response = litellm.ModelResponse(model="openai/gpt-4o-mini", usage={"total_tokens": 100}) + + tpm_key = await router.deployment_callback_on_success( + kwargs=kwargs, + completion_response=response, + start_time=time.time(), + end_time=time.time(), + ) + + # The IO deployment is no longer skipped: its TPM routing counter is tracked. + assert tpm_key is not None + assert await router.cache.async_get_cache(key=tpm_key) == 100 + + @pytest.mark.asyncio async def test_deployment_callback_on_failure(model_list): """Test if the '_deployment_callback_on_failure' function is working correctly""" @@ -923,6 +972,227 @@ class _Resp(BaseModel): assert headers["x-ratelimit-limit-requests"] == 100 +@pytest.mark.asyncio +async def test_set_response_headers_in_flight_delta_only_adjusts_tpm_rpm(model_list): + """ + The in-flight replay applies only to the post-incremented TPM/RPM counters + (`x-ratelimit-remaining-tokens` / `-requests`). The ITPM/OTPM counters are + incremented at reservation time (pre-call), so the input/output token + headers already reflect this request and must pass through untouched. + """ + from pydantic import BaseModel + + class _Usage(BaseModel): + total_tokens: int = 30 + prompt_tokens: int = 20 + completion_tokens: int = 10 + + class _Resp(BaseModel): + usage: _Usage = _Usage() + _hidden_params: dict = {} + + router = Router(model_list=model_list) + router.get_remaining_model_group_usage = AsyncMock( + return_value={ + "x-ratelimit-remaining-tokens": 1000, + "x-ratelimit-remaining-requests": 100, + "x-ratelimit-remaining-input-tokens": 1000, + "x-ratelimit-remaining-output-tokens": 500, + } + ) + + resp = _Resp() + resp._hidden_params = {} + await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") + + headers = resp._hidden_params["additional_headers"] + # TPM/RPM headers replay the in-flight increment... + assert headers["x-ratelimit-remaining-tokens"] == 970 + assert headers["x-ratelimit-remaining-requests"] == 99 + # ...but the reservation-based input/output headers pass through unchanged. + assert headers["x-ratelimit-remaining-input-tokens"] == 1000 + assert headers["x-ratelimit-remaining-output-tokens"] == 500 + + +@pytest.mark.asyncio +async def test_get_model_group_io_token_usage_sums_across_deployments(): + """ + get_model_group_io_token_usage must sum ITPM/OTPM across every deployment + in the model group (not just the first), reading the same per-deployment + cache keys the pre-call reservation writes to. + """ + from litellm.types.router import RouterCacheEnum + from litellm.utils import get_utc_datetime + + router = Router( + model_list=[ + { + "model_name": "opus", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "itpm": 1000, + "otpm": 500, + }, + "model_info": {"id": "io-usage-dep-1"}, + }, + { + "model_name": "opus", + "litellm_params": { + "model": "openai/gpt-4o", + "itpm": 1000, + "otpm": 500, + }, + "model_info": {"id": "io-usage-dep-2"}, + }, + ] + ) + + minute = get_utc_datetime().strftime("%H-%M") + keys_and_values = [ + ( + RouterCacheEnum.ITPM.value.format( + id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute + ), + 30, + ), + ( + RouterCacheEnum.OTPM.value.format( + id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute + ), + 10, + ), + ( + RouterCacheEnum.ITPM.value.format( + id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute + ), + 70, + ), + ( + RouterCacheEnum.OTPM.value.format( + id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute + ), + 20, + ), + ] + for key, value in keys_and_values: + await router.cache.async_increment_cache(key=key, value=value, ttl=60) + + current_itpm, current_otpm = await router.get_model_group_io_token_usage("opus") + + assert current_itpm == 100 + assert current_otpm == 30 + + +@pytest.mark.asyncio +async def test_get_model_group_io_token_usage_no_deployments_returns_none(): + router = Router(model_list=[]) + current_itpm, current_otpm = await router.get_model_group_io_token_usage( + "nonexistent-group" + ) + assert current_itpm is None + assert current_otpm is None + + +@pytest.mark.asyncio +async def test_get_remaining_model_group_usage_merges_io_and_tpm_headers(model_list): + """ + A model group with both itpm/otpm and tpm/rpm limits must expose the + standard remaining-tokens/requests headers alongside the input/output token + headers, so clients and prometheus gauges relying on either still get data. + """ + from unittest.mock import Mock + + from litellm.types.router import ModelGroupInfo + + router = Router(model_list=model_list) + router._cached_get_model_group_info = Mock( + return_value=ModelGroupInfo( + model_group="gpt-3.5-turbo", + providers=["openai"], + itpm=2000, + otpm=1000, + tpm=5000, + rpm=50, + ) + ) + router.get_model_group_io_token_usage = AsyncMock(return_value=(100, 40)) + router.get_model_group_usage = AsyncMock(return_value=(500, 5)) + + headers = await router.get_remaining_model_group_usage("gpt-3.5-turbo") + + assert headers["x-ratelimit-remaining-input-tokens"] == 1900 + assert headers["x-ratelimit-remaining-output-tokens"] == 960 + assert headers["x-ratelimit-remaining-tokens"] == 4500 + assert headers["x-ratelimit-remaining-requests"] == 45 + + +@pytest.mark.asyncio +async def test_set_response_headers_native_input_token_header_does_not_suppress_router_headers(model_list): + """ + A provider that natively returns `x-ratelimit-remaining-input-tokens` must + not suppress the router's own remaining-tokens/requests headers for a + non-IO model group. + """ + from pydantic import BaseModel + + class _Usage(BaseModel): + total_tokens: int = 42 + + class _Resp(BaseModel): + usage: _Usage = _Usage() + _hidden_params: dict = {} + + router = Router(model_list=model_list) + router.get_remaining_model_group_usage = AsyncMock( + return_value={ + "x-ratelimit-remaining-tokens": 1000, + "x-ratelimit-remaining-requests": 100, + } + ) + + resp = _Resp() + resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-input-tokens": 5}} + await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") + + headers = resp._hidden_params["additional_headers"] + assert headers["x-ratelimit-remaining-tokens"] == 958 + assert headers["x-ratelimit-remaining-requests"] == 99 + # the provider's native header is left untouched + assert headers["x-ratelimit-remaining-input-tokens"] == 5 + + +@pytest.mark.asyncio +async def test_set_response_headers_native_token_header_does_not_suppress_io_headers(model_list): + from pydantic import BaseModel + + class _Usage(BaseModel): + total_tokens: int = 42 + + class _Resp(BaseModel): + usage: _Usage = _Usage() + _hidden_params: dict = {} + + router = Router(model_list=model_list) + router.get_remaining_model_group_usage = AsyncMock( + return_value={ + "x-ratelimit-remaining-tokens": 1000, + "x-ratelimit-remaining-requests": 100, + "x-ratelimit-remaining-input-tokens": 900, + "x-ratelimit-remaining-output-tokens": 450, + } + ) + + resp = _Resp() + resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-tokens": 5}} + await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") + + headers = resp._hidden_params["additional_headers"] + assert headers["x-ratelimit-remaining-tokens"] == 5 + assert headers["x-ratelimit-remaining-requests"] == 99 + assert headers["x-ratelimit-remaining-input-tokens"] == 900 + assert headers["x-ratelimit-remaining-output-tokens"] == 450 + + @pytest.mark.asyncio async def test_set_response_headers_handles_missing_usage(model_list): """ @@ -952,6 +1222,72 @@ class _Resp(BaseModel): assert headers["x-ratelimit-remaining-requests"] == 99 +@pytest.mark.asyncio +async def test_set_response_headers_dict_anthropic_messages_response(model_list): + """Anthropic /v1/messages returns a dict; IO rate-limit headers must attach.""" + router = Router(model_list=model_list) + router.get_remaining_model_group_usage = AsyncMock( + return_value={ + "x-ratelimit-limit-input-tokens": 25, + "x-ratelimit-remaining-input-tokens": 20, + "x-ratelimit-limit-output-tokens": 100, + "x-ratelimit-remaining-output-tokens": 95, + } + ) + + resp = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + "usage": {"input_tokens": 5, "output_tokens": 1}, + } + await router.set_response_headers(response=resp, model_group="io-itpm-strict") + + assert "_hidden_params" in resp + headers = resp["_hidden_params"]["additional_headers"] + assert headers["x-litellm-model-group"] == "io-itpm-strict" + assert headers["x-ratelimit-limit-input-tokens"] == 25 + assert headers["x-ratelimit-remaining-input-tokens"] == 20 + assert headers["x-ratelimit-remaining-output-tokens"] == 95 + + +@pytest.mark.asyncio +async def test_set_response_headers_wraps_bare_async_generator(model_list): + """ + Streaming responses that never go through Router.make_call's usual + object-based wrappers (e.g. the Anthropic /v1/messages -> Responses API + bridge, which yields a raw async generator with no `_hidden_params` slot) + must still get IO rate-limit headers attached via a thin wrapper. + """ + + async def _raw_generator(): + yield {"type": "message_start"} + yield {"type": "message_stop"} + + router = Router(model_list=model_list) + router.get_remaining_model_group_usage = AsyncMock( + return_value={ + "x-ratelimit-limit-input-tokens": 25, + "x-ratelimit-remaining-input-tokens": 20, + } + ) + + wrapped = await router.set_response_headers(response=_raw_generator(), model_group="io-itpm-strict") + + assert hasattr(wrapped, "_hidden_params") + headers = wrapped._hidden_params["additional_headers"] + assert headers["x-litellm-model-group"] == "io-itpm-strict" + assert headers["x-ratelimit-limit-input-tokens"] == 25 + assert headers["x-ratelimit-remaining-input-tokens"] == 20 + + from collections.abc import AsyncIterator + + assert isinstance(wrapped, AsyncIterator) + chunks = [chunk async for chunk in wrapped] + assert chunks == [{"type": "message_start"}, {"type": "message_stop"}] + + def test_get_all_deployments(model_list): """Test if the 'get_all_deployments' function is working correctly""" router = Router(model_list=model_list) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 1d0dafed171..c185b694e68 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -4090,7 +4090,7 @@ def _build_logging_obj(self, *, model_call_details, response_cost_calculator): logging_obj._on_deferred_stream_complete = None return logging_obj - async def _drive_non_streaming(self, *, monkeypatch, response, logging_obj, route_type): + async def _drive_non_streaming(self, *, monkeypatch, response, logging_obj, route_type, return_result=False): import litellm.proxy.common_request_processing as crp from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth @@ -4119,7 +4119,7 @@ async def fake_post_call_success_hook(data, user_api_key_dict, response): "_has_post_call_guardrails", return_value=False, ): - await processing_obj.base_process_llm_request( + result = await processing_obj.base_process_llm_request( request=MagicMock(spec=Request, headers={}), fastapi_response=fastapi_response, user_api_key_dict=RealUserAPIKeyAuth(api_key="sk-test"), @@ -4131,6 +4131,8 @@ async def fake_post_call_success_hook(data, user_api_key_dict, response): llm_router=None, skip_pre_call_logic=True, ) + if return_result: + return fastapi_response, result return fastapi_response @pytest.mark.asyncio @@ -4352,3 +4354,46 @@ async def test_object_response_zero_cost_drops_header_like_chat_completions(self assert "x-litellm-response-cost" not in fastapi_response.headers recompute.assert_not_called() + + @pytest.mark.asyncio + async def test_messages_typeddict_does_not_leak_hidden_params_into_response_body(self, monkeypatch): + """ + Router.set_response_headers now writes rate-limit headers onto dict-shaped + responses (e.g. Anthropic /v1/messages, whose AnthropicMessagesResponse is a + TypedDict) via response["_hidden_params"] = ... . Unlike a pydantic model's + private attribute, that key is indistinguishable from any other dict key and + would otherwise serialize verbatim into the client-facing JSON body, leaking + response_cost/model_id/api_base/fallback errors. base_process_llm_request + must strip it before returning the response to the endpoint layer. + """ + from litellm.types.utils import AnthropicMessagesResponse + + response = AnthropicMessagesResponse( + id="msg_1", + type="message", + role="assistant", + content=[{"type": "text", "text": "hi"}], + model="claude-haiku-4-5", + usage={"input_tokens": 10, "output_tokens": 5}, + ) + response["_hidden_params"] = { + "additional_headers": {"x-ratelimit-limit-input-tokens": "25"}, + "response_cost": 0.00123, + "model_id": "internal-deployment-id", + } + logging_obj = self._build_logging_obj( + model_call_details={"response_cost": 0.00123}, + response_cost_calculator=MagicMock(return_value=999.0), + ) + + fastapi_response, result = await self._drive_non_streaming( + monkeypatch=monkeypatch, + response=response, + logging_obj=logging_obj, + route_type="anthropic_messages", + return_result=True, + ) + + assert "_hidden_params" not in result + assert fastapi_response.headers["x-ratelimit-limit-input-tokens"] == "25" + assert fastapi_response.headers["x-litellm-response-cost"] == "0.00123" diff --git a/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py b/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py index 2aee0f0a4ef..3a0deeb13d8 100644 --- a/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py +++ b/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py @@ -140,3 +140,25 @@ def test_get_fallback_errors_from_headers_invalid_json_returns_empty(): def test_get_fallback_errors_from_headers_missing_key_returns_empty(): result = get_fallback_errors_from_headers({}) assert result == [] + + +def test_get_hidden_params_dict_with_dict_response(): + response = {"id": "msg_1", "usage": {"input_tokens": 1, "output_tokens": 2}} + assert get_hidden_params_dict(response) == {} + + hidden_params = get_hidden_params_dict(response, create=True) + assert hidden_params == {} + assert response["_hidden_params"] == {} + + response["_hidden_params"] = {"additional_headers": {"x-test": "1"}} + assert get_hidden_params_dict(response) == { + "additional_headers": {"x-test": "1"}, + } + + +def test_add_fallback_headers_to_dict_response(): + response = {"id": "msg_1"} + result = add_fallback_headers_to_response(response=response, attempted_fallbacks=1) + + assert result is response + assert response["_hidden_params"]["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 diff --git a/tests/test_litellm/test_router/test_io_token_rate_limits.py b/tests/test_litellm/test_router/test_io_token_rate_limits.py new file mode 100644 index 00000000000..939e3189596 --- /dev/null +++ b/tests/test_litellm/test_router/test_io_token_rate_limits.py @@ -0,0 +1,979 @@ +""" +Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). +""" + +import asyncio + +import pytest + +import litellm +from litellm import Router +from litellm.caching.dual_cache import DualCache +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + ITPM_CACHE_KEY, + ITPM_RESERVED_KEY, + OTPM_CACHE_KEY, + OTPM_RESERVED_KEY, + _reservation_value, + _resolve_max_tokens, + async_io_token_pre_call_check, + async_io_token_reconcile_success, + build_io_token_rate_limit_headers, + deployment_has_io_token_limits, + get_io_token_rate_limit_request_kwargs, + io_token_reconcile_success, + io_token_refund_failure, + refund_stale_reservation_before_retry, + set_io_token_rate_limit_request_kwargs, +) +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( + ModelRateLimitingCheck, +) +from litellm.types.utils import ModelResponse, Usage + + +class TestIOTokenRateLimitHelpers: + def test_deployment_has_io_token_limits(self): + assert deployment_has_io_token_limits({"litellm_params": {"itpm": 100, "otpm": 50}}) + assert not deployment_has_io_token_limits({"litellm_params": {"model": "x"}}) + + def test_reservation_value_minimal_when_estimate_fails(self): + # A failed/empty estimate (0) must reserve a minimal slot, not the + # entire limit - otherwise one request whose estimate failed fills + # the whole bucket and blocks every concurrent request until it + # completes and reconciles. + assert _reservation_value(0, 100) == 1 + assert _reservation_value(0, 1) == 1 + # A real non-zero estimate is reserved as-is. + assert _reservation_value(42, 100) == 42 + + def test_resolve_max_tokens_respects_explicit_zero(self): + deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}} + # An explicit max_tokens=0 is honored, not replaced by the model default. + assert _resolve_max_tokens({"max_tokens": 0}, deployment) == 0 + # max_completion_tokens is the fallback only when max_tokens is absent. + assert _resolve_max_tokens({"max_completion_tokens": 12}, deployment) == 12 + assert _resolve_max_tokens({"max_output_tokens": 9}, deployment) == 9 + + def test_build_io_token_rate_limit_headers(self): + headers = build_io_token_rate_limit_headers( + itpm_limit=200, + otpm_limit=40, + current_itpm=15, + current_otpm=4, + ) + assert headers["x-ratelimit-limit-input-tokens"] == 200 + assert headers["x-ratelimit-remaining-input-tokens"] == 185 + assert headers["x-ratelimit-limit-output-tokens"] == 40 + assert headers["x-ratelimit-remaining-output-tokens"] == 36 + + +class TestModelRateLimitingCheckIOTokens: + @pytest.mark.asyncio + async def test_itpm_reservation_and_reconcile(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 50, + }, + "model_info": {"id": "io-test-id"}, + "model_name": "opus", + } + + request_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + otpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + + kwargs = { + "standard_logging_object": { + "model_id": "io-test-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[ + { + "message": {"role": "assistant", "content": "hi"}, + "index": 0, + "finish_reason": "stop", + } + ], + usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), + ) + await check.async_log_success_event(kwargs, response, None, None) + + current_itpm = await dual_cache.async_get_cache(key=itpm_key) + current_otpm = await dual_cache.async_get_cache(key=otpm_key) + # ITPM tracks input tokens only (billable prompt tokens), not output. + assert current_itpm == 5 + assert current_otpm == 3 + + @pytest.mark.asyncio + async def test_itpm_limit_raises_429(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 5, + }, + "model_info": {"id": "io-limit-id"}, + "model_name": "opus", + } + + # ITPM enforces input tokens only; the prompt alone must exceed the limit, + # a large max_tokens must not contribute to the ITPM reservation. + set_io_token_rate_limit_request_kwargs( + { + "messages": [ + { + "role": "user", + "content": "hello world this is a longer prompt that exceeds the tiny itpm limit", + } + ], + "max_tokens": 10, + "metadata": {}, + } + ) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "ITPM limit=5" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_otpm_atomic_reservation_no_overshoot_under_concurrency(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + otpm_limit = 10 + max_tokens = 4 + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "otpm": otpm_limit, + }, + "model_info": {"id": "io-otpm-race-id"}, + "model_name": "opus", + } + + set_io_token_rate_limit_request_kwargs( + { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": max_tokens, + "metadata": {}, + } + ) + + async def _attempt(): + try: + await check.async_pre_call_check(deployment) + return True + except litellm.RateLimitError: + return False + + results = await asyncio.gather(*[_attempt() for _ in range(8)]) + successes = sum(1 for r in results if r) + + minute = get_utc_datetime().strftime("%H-%M") + otpm_key = f"global_router:io-otpm-race-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + current_otpm = await dual_cache.async_get_cache(key=otpm_key) + + # Atomic reservation must never let concurrent requests overshoot the limit. + assert current_otpm is not None + assert current_otpm <= otpm_limit + assert successes == otpm_limit // max_tokens + assert current_otpm == successes * max_tokens + + @pytest.mark.asyncio + async def test_itpm_estimate_failure_reserves_minimal_not_full_limit(self): + """ + When input-token estimation yields 0 (no messages/prompt/input field, + unsupported model, tokenizer error), the reservation must be a + minimal 1 token, not the entire itpm limit. Otherwise the first + request whose estimate fails fills the whole bucket and every + concurrent request is rejected until it completes - effectively + serializing traffic to the deployment. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_limit = 5 + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": itpm_limit, + }, + "model_info": {"id": "io-itpm-estimate-fail-id"}, + "model_name": "opus", + } + + # No messages/prompt/input field -> _estimate_input_tokens returns 0. + set_io_token_rate_limit_request_kwargs( + { + "max_tokens": 5, + "metadata": {}, + } + ) + + async def _attempt(): + try: + await check.async_pre_call_check(deployment) + return True + except litellm.RateLimitError: + return False + + results = await asyncio.gather(*[_attempt() for _ in range(8)]) + successes = sum(1 for r in results if r) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-itpm-estimate-fail-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + current_itpm = await dual_cache.async_get_cache(key=itpm_key) + + # A minimal 1-token reservation per request lets itpm_limit concurrent + # requests through, instead of a single request starving the rest. + assert current_itpm is not None + assert current_itpm <= itpm_limit + assert successes == itpm_limit + + @pytest.mark.asyncio + async def test_reservation_read_prefers_top_level_metadata_over_litellm_params(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-lp-id:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + + # Production kwargs commonly carry litellm_params.metadata; the stashed + # reservation lives in the top-level metadata and must still be found. + kwargs = { + "standard_logging_object": { + "model_id": "io-lp-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + "metadata": {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key}, + "litellm_params": {"metadata": {"user_api_key_hash": "abc123"}}, + } + await check.async_log_failure_event(kwargs, None, None, None) + + current = await dual_cache.async_get_cache(key=itpm_key) + assert current == 0 + + @pytest.mark.asyncio + async def test_reconcile_tracks_actual_usage_when_estimate_zero(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + }, + "model_info": {"id": "io-zero-est-id"}, + "model_name": "opus", + } + + request_kwargs = {"max_tokens": 5, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-zero-est-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + # A failed/zero estimate reserves a minimal 1 token, not the full + # itpm limit, so it doesn't starve concurrent requests. + assert await dual_cache.async_get_cache(key=itpm_key) == 1 + + kwargs = { + "standard_logging_object": { + "model_id": "io-zero-est-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7), + ) + await check.async_log_success_event(kwargs, response, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 7 + + @pytest.mark.asyncio + async def test_zero_estimate_reserves_minimal_capacity_before_reconcile(self): + """ + A zero/failed estimate reserves a minimal 1 token rather than the + full itpm limit, so up to itpm_limit such calls are allowed + concurrently instead of the first one claiming the entire bucket. + """ + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 2, + }, + "model_info": {"id": "io-zero-cap-id"}, + "model_name": "opus", + } + request_kwargs = {"max_tokens": 5, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + # Second zero-estimate call still fits within the itpm=2 limit. + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + await check.async_pre_call_check(deployment) + + # A third exceeds the limit and is rejected. + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + with pytest.raises(litellm.RateLimitError): + await check.async_pre_call_check(deployment) + + @pytest.mark.asyncio + async def test_explicit_zero_max_tokens_does_not_reserve_otpm(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "otpm": 5, + }, + "model_info": {"id": "io-zero-output-id"}, + "model_name": "opus", + } + zero_output_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 0, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(zero_output_kwargs) + await check.async_pre_call_check(deployment) + + zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] + assert zero_output_kwargs["metadata"][OTPM_RESERVED_KEY] == 0 + assert (await dual_cache.async_get_cache(key=zero_output_otpm_key) or 0) == 0 + + normal_output_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(normal_output_kwargs) + await check.async_pre_call_check(deployment) + + normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] + assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5 + + def test_sync_io_pre_call_reserves_and_reconciles(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 50, + }, + "model_info": {"id": "io-sync-id"}, + "model_name": "opus", + } + request_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + check.pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + otpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + kwargs = { + "standard_logging_object": { + "model_id": "io-sync-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), + ) + check.log_success_event(kwargs, response, None, None) + + assert dual_cache.get_cache(key=itpm_key) == 5 + assert dual_cache.get_cache(key=otpm_key) == 3 + + @pytest.mark.asyncio + async def test_reconcile_runs_via_success_event_without_model_id(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_key = "global_router:io-noid:bedrock_mantle/test:itpm:12-34" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + # standard_logging_object has no model_id (only the TPM path needs it); + # IO reconciliation must still run off the stashed cache key. + kwargs = { + "standard_logging_object": { + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + await check.async_log_success_event(kwargs, response, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 3 + + @pytest.mark.asyncio + async def test_failure_clears_reservation_so_retry_is_not_poisoned(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-first:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + + # Shared request metadata carrying the first (IO) deployment's reservation. + metadata = {ITPM_RESERVED_KEY: 8, ITPM_CACHE_KEY: itpm_key} + fail_kwargs = { + "metadata": metadata, + "standard_logging_object": { + "model_id": "io-first", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + } + await check.async_log_failure_event(fail_kwargs, None, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 0 + assert ITPM_RESERVED_KEY not in metadata + assert ITPM_CACHE_KEY not in metadata + + # Retry succeeds on a non-IO fallback deployment reusing the same metadata. + retry_kwargs = { + "metadata": metadata, + "standard_logging_object": { + "model_id": "non-io-second", + "hidden_params": {"litellm_model_name": "openai/gpt-4o-mini"}, + "metadata": {}, + "total_tokens": 12, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=6, completion_tokens=6, total_tokens=12), + ) + await check.async_log_success_event(retry_kwargs, response, None, None) + + # The first deployment's ITPM counter is not driven negative... + assert await dual_cache.async_get_cache(key=itpm_key) == 0 + # ...and the non-IO deployment's TPM usage is tracked normally. + tpm_key = f"non-io-second:openai/gpt-4o-mini:tpm:{minute}" + assert await dual_cache.async_get_cache(key=tpm_key) == 12 + + @pytest.mark.asyncio + async def test_stale_reservation_refunded_before_retry_overwrites_it(self): + """ + A retry reuses the same mutable kwargs dict for the next deployment. + If deployment A's failure event hasn't run yet (e.g. it was scheduled + as a background task) when the retry calls + set_io_token_rate_limit_request_kwargs for deployment B, the router + must first synchronously refund + clear A's reservation via + refund_stale_reservation_before_retry - otherwise A's counter stays + elevated by the reservation until its TTL expires, and the + now-orphaned sentinels must not leak into B's accounting either. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + minute = get_utc_datetime().strftime("%H-%M") + itpm_key_a = f"global_router:io-retry-a:bedrock_mantle/test-a:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key_a, value=9, ttl=60) + + # Deployment A's still-unreconciled reservation, stashed on the shared + # kwargs dict the retry loop reuses. + shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_CACHE_KEY: itpm_key_a}} + + # Router calls this before overwriting kwargs for deployment B's attempt - + # simulating the fix landing ahead of set_io_token_rate_limit_request_kwargs. + refund_stale_reservation_before_retry(dual_cache, shared_kwargs) + + # A's reservation is refunded immediately, not left stranded for a + # background failure task that may run arbitrarily later (or never, + # if the sentinels get cleared out from under it first). + assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + assert ITPM_RESERVED_KEY not in shared_kwargs["metadata"] + assert ITPM_CACHE_KEY not in shared_kwargs["metadata"] + + # A's own (now-late) failure event finds nothing left to refund and + # is a safe no-op, since the sentinels were already cleared above. + io_token_refund_failure(dual_cache, shared_kwargs) + assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + + # The retry proceeds to stash deployment B's own reservation on the + # same dict; it starts clean, unaffected by A's cleared sentinels. + set_io_token_rate_limit_request_kwargs(shared_kwargs) + itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" + shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 + shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b + await dual_cache.async_increment_cache(key=itpm_key_b, value=4, ttl=60) + assert await dual_cache.async_get_cache(key=itpm_key_b) == 4 + + @pytest.mark.asyncio + async def test_client_supplied_reservation_keys_are_stripped(self): + # metadata is caller-controlled; the server-only reservation sentinels + # must be removed before the router captures the request kwargs. + forged = { + "metadata": {ITPM_RESERVED_KEY: 999999, ITPM_CACHE_KEY: "attacker:key:itpm:00-00"}, + "litellm_metadata": {OTPM_RESERVED_KEY: 7}, + "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, + } + set_io_token_rate_limit_request_kwargs(forged) + stored = get_io_token_rate_limit_request_kwargs() + + assert ITPM_RESERVED_KEY not in stored["metadata"] + assert ITPM_CACHE_KEY not in stored["metadata"] + assert OTPM_RESERVED_KEY not in stored["litellm_metadata"] + assert OTPM_CACHE_KEY not in stored["litellm_params"]["metadata"] + + @pytest.mark.asyncio + async def test_forged_reservation_cannot_decrement_counter(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + victim_key = "global_router:victim:model:itpm:00-00" + await dual_cache.async_increment_cache(key=victim_key, value=100, ttl=60) + + # A caller forges a reservation pointing at another deployment's counter. + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 100, ITPM_CACHE_KEY: victim_key}, + "standard_logging_object": { + "model_id": "m", + "hidden_params": {"litellm_model_name": "model"}, + "metadata": {}, + "total_tokens": 2, + }, + } + # The router sanitizes the request kwargs before the call runs. + set_io_token_rate_limit_request_kwargs(kwargs) + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + await check.async_log_success_event(kwargs, response, None, None) + + # The forged reservation was stripped, so the victim counter is untouched. + assert await dual_cache.async_get_cache(key=victim_key) == 100 + + @pytest.mark.asyncio + async def test_otpm_reservation_error_rolls_back_itpm(self): + from litellm.utils import get_utc_datetime + + class _OtpmFailCache(DualCache): + async def async_increment_cache(self, key, **kwargs): + if ":otpm:" in key: + raise RuntimeError("transient cache error") + return await super().async_increment_cache(key=key, **kwargs) + + dual_cache = _OtpmFailCache() + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 1000, + "otpm": 1000, + }, + "model_info": {"id": "io-rollback-id"}, + "model_name": "opus", + } + set_io_token_rate_limit_request_kwargs( + { + "messages": [{"role": "user", "content": "hello world"}], + "max_tokens": 5, + "metadata": {}, + } + ) + + with pytest.raises(RuntimeError): + await async_io_token_pre_call_check(dual_cache, deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-rollback-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + # A transient OTPM error must release the ITPM reservation, not leak it. + assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + + @pytest.mark.asyncio + async def test_reconcile_clears_stash_even_when_increment_errors(self): + class _ItpmFailCache(DualCache): + async def async_increment_cache(self, key, **kwargs): + if ":itpm:" in key: + raise RuntimeError("transient cache error") + return await super().async_increment_cache(key=key, **kwargs) + + dual_cache = _ItpmFailCache() + metadata = {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: "global_router:x:model:itpm:00-00"} + kwargs = {"metadata": metadata} + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + with pytest.raises(RuntimeError): + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + # The stash is cleared even though reconciliation raised, so a duplicate + # success event can't re-process it. + assert ITPM_RESERVED_KEY not in metadata + assert ITPM_CACHE_KEY not in metadata + + @pytest.mark.asyncio + async def test_io_conflict_warning_not_collapsed_for_missing_model_id(self, caplog): + import logging + + check = ModelRateLimitingCheck(dual_cache=DualCache()) + deployment = { + "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100, "tpm": 1000}, + "model_info": {}, + } + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + check._warn_io_token_and_tpm_rpm_coexist_once(deployment) + check._warn_io_token_and_tpm_rpm_coexist_once(deployment) + + warnings = [r for r in caplog.records if "both limit types are enforced" in r.message] + # id-less deployments are not collapsed onto a single dedup key. + assert len(warnings) == 2 + + @pytest.mark.asyncio + async def test_missing_deployment_id_skips_io_reservation(self): + dual_cache = DualCache() + deployment = { + "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100}, + "model_info": {}, # no id -> cannot build a per-deployment cache key + "model_name": "opus", + } + request_kwargs = { + "messages": [{"role": "user", "content": "hello world"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + + result = await async_io_token_pre_call_check(dual_cache, deployment) + + assert result is deployment + # No reservation is stashed, so nothing lands in a shared None:None bucket. + assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] + + @pytest.mark.asyncio + async def test_reconcile_uses_reservation_minute_key(self): + dual_cache = DualCache() + # Reservation was made on a fixed minute key; a call that finishes in a + # later minute must reconcile against that same key, never a key built + # from the response-time minute. + itpm_key = "global_router:io-min-id:bedrock_mantle/test:itpm:99-99" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}} + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4), + ) + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 4 + + @pytest.mark.asyncio + async def test_reconcile_missing_usage_keeps_reservation(self): + dual_cache = DualCache() + itpm_key = "global_router:io-missing-usage:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-missing-usage:bedrock_mantle/test:otpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 8, + OTPM_RESERVED_KEY: 5, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + ) + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 8 + assert await dual_cache.async_get_cache(key=otpm_key) == 5 + assert ITPM_RESERVED_KEY not in kwargs["metadata"] + + @pytest.mark.asyncio + async def test_reconcile_total_tokens_only_keeps_reservation(self): + """ + A response usage object with only total_tokens (no prompt/completion + breakdown) can't be split into input/output, so it must be treated the + same as missing usage: keep the reservation instead of resolving to + (0, 0) and refunding it in full. + """ + dual_cache = DualCache() + itpm_key = "global_router:io-total-only:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-total-only:bedrock_mantle/test:otpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 8, + OTPM_RESERVED_KEY: 5, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = {"type": "message", "usage": {"total_tokens": 13}} + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 8 + assert await dual_cache.async_get_cache(key=otpm_key) == 5 + + def test_reconcile_standard_logging_total_tokens_only_keeps_reservation(self): + dual_cache = DualCache() + itpm_key = "global_router:io-slo-total-only:bedrock_mantle/test:itpm:00-00" + dual_cache.set_cache(key=itpm_key, value=10, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": {"total_tokens": 4}, + } + response = {"type": "message", "role": "assistant", "content": []} + + io_token_reconcile_success(dual_cache, kwargs, response) + + assert dual_cache.get_cache(key=itpm_key) == 10 + + @pytest.mark.asyncio + async def test_reconcile_falls_back_to_standard_logging_object(self): + dual_cache = DualCache() + itpm_key = "global_router:io-slo-fallback:bedrock_mantle/test:itpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "prompt_tokens": 4, + "completion_tokens": 0, + "total_tokens": 4, + }, + } + response = {"type": "message", "role": "assistant", "content": []} + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 4 + + def test_sync_reconcile_anthropic_dict_usage(self): + dual_cache = DualCache() + itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00" + dual_cache.set_cache(key=itpm_key, value=6, ttl=60) + dual_cache.set_cache(key=otpm_key, value=4, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 6, + OTPM_RESERVED_KEY: 4, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = { + "type": "message", + "usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1}, + } + + io_token_reconcile_success(dual_cache, kwargs, response) + + assert dual_cache.get_cache(key=itpm_key) == 2 + assert dual_cache.get_cache(key=otpm_key) == 2 + + @pytest.mark.asyncio + async def test_io_and_tpm_rpm_limits_both_enforced_with_warning(self, caplog): + import logging + + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-mixed-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + deployment = { + "litellm_params": { + "model": deployment_name, + "itpm": 100, + "rpm": 1, + }, + "model_info": {"id": model_id}, + "model_name": "opus", + } + + minute = get_utc_datetime().strftime("%H-%M") + rpm_key = f"{model_id}:{deployment_name}:rpm:{minute}" + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + await dual_cache.async_increment_cache(key=rpm_key, value=5, ttl=60) + + request_kwargs = { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + with pytest.raises(litellm.RateLimitError): + await check.async_pre_call_check(deployment) + + assert await dual_cache.async_get_cache(key=rpm_key) == 6 + assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] + assert any("both limit types are enforced" in record.message for record in caplog.records) + + @pytest.mark.asyncio + async def test_io_success_still_tracks_tpm_for_mixed_deployment(self): + """ + A deployment with itpm/otpm AND tpm/rpm must have BOTH counters updated on + success, otherwise the tpm_key the pre-call check reads is never written + and the tpm_limit can never be enforced. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-tpm-mixed-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=5, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "model_id": model_id, + "total_tokens": 7, + "hidden_params": {"litellm_model_name": deployment_name}, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + await check.async_log_success_event(kwargs, response, None, None) + + # ITPM reconciled down from the 5-token reservation to actual usage (3). + assert await dual_cache.async_get_cache(key=itpm_key) == 3 + # TPM tracking must still run so the tpm/rpm pre-call path can enforce it. + assert await dual_cache.async_get_cache(key=tpm_key) == 7 + + def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-tpm-mixed-sync-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" + dual_cache.set_cache(key=itpm_key, value=5, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "model_id": model_id, + "total_tokens": 7, + "hidden_params": {"litellm_model_name": deployment_name}, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + check.log_success_event(kwargs, response, None, None) + + assert dual_cache.get_cache(key=itpm_key) == 3 + assert dual_cache.get_cache(key=tpm_key) == 7 + + @pytest.mark.asyncio + async def test_failure_refunds_itpm_reservation(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-refund-id:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + + reservation = {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key} + kwargs = { + "standard_logging_object": { + "model_id": "io-refund-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": dict(reservation), + }, + "metadata": dict(reservation), + } + await check.async_log_failure_event(kwargs, None, None, None) + + current = await dual_cache.async_get_cache(key=itpm_key) + assert current == 0 + + +class TestRouterIOTokenIntegration: + @pytest.mark.asyncio + async def test_model_group_info_aggregates_io_limits(self): + router = Router( + model_list=[ + { + "model_name": "opus", + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 20, + }, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + info = router.get_model_group_info("opus") + assert info is not None + assert info.itpm == 100 + assert info.otpm == 20 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index bfd1efc99f3..496a0462ebe 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25430,6 +25430,8 @@ export interface components { input_cost_per_video_per_second_above_15s_interval?: number | null; /** Input Cost Per Video Per Second Above 8S Interval */ input_cost_per_video_per_second_above_8s_interval?: number | null; + /** Itpm */ + itpm?: number | null; /** Litellm Credential Name */ litellm_credential_name?: string | null; /** Litellm Trace Id */ @@ -25465,6 +25467,8 @@ export interface components { ocr_cost_per_page?: number | null; /** Organization */ organization?: string | null; + /** Otpm */ + otpm?: number | null; /** Output Cost Per Audio Per Second */ output_cost_per_audio_per_second?: number | null; /** Output Cost Per Audio Token */ @@ -27261,6 +27265,8 @@ export interface components { * @default false */ is_public_model_group: boolean; + /** Itpm */ + itpm?: number | null; /** Max Input Tokens */ max_input_tokens?: number | null; /** Max Output Tokens */ @@ -27272,6 +27278,8 @@ export interface components { mode: string | ("chat" | "embedding" | "completion" | "image_generation" | "audio_transcription" | "rerank" | "moderations") | null; /** Model Group */ model_group: string; + /** Otpm */ + otpm?: number | null; /** Output Cost Per Token */ output_cost_per_token?: number | null; /** Providers */ @@ -33216,6 +33224,8 @@ export interface components { input_cost_per_video_per_second_above_15s_interval?: number | null; /** Input Cost Per Video Per Second Above 8S Interval */ input_cost_per_video_per_second_above_8s_interval?: number | null; + /** Itpm */ + itpm?: number | null; /** Litellm Credential Name */ litellm_credential_name?: string | null; /** Litellm Trace Id */ @@ -33251,6 +33261,8 @@ export interface components { ocr_cost_per_page?: number | null; /** Organization */ organization?: string | null; + /** Otpm */ + otpm?: number | null; /** Output Cost Per Audio Per Second */ output_cost_per_audio_per_second?: number | null; /** Output Cost Per Audio Token */