From b2bbc3192350b1c160ac47f2823595b9ec39fe77 Mon Sep 17 00:00:00 2001 From: chuan932 Date: Thu, 29 Oct 2026 20:19:23 +0800 Subject: [PATCH] feat(router): clamp completion token budgets with per-deployment max_output_tokens_cap Add an opt-in litellm_params.max_output_tokens_cap (int) per deployment. At deployment time, any caller-requested output budget the request already carries (max_tokens / max_completion_tokens / max_output_tokens) above the cap is lowered to the cap. Smaller budgets are respected and requests without a budget are untouched: a default remains the job of default_litellm_params, which the cap now also bounds. Works for both dict and pydantic deployment params; invalid cap values are ignored. Operators fronting self-hosted OpenAI-compatible engines need a hard output ceiling per deployment to bound queue latency and memory; router defaults only fill in missing parameters, they cannot lower a caller budget. Registers the new param in the proxy credential-registry classification tables (NotSecret), as its guard test requires of every LiteLLM_Params field with a credential-shaped name. Tests: tests/unit/test_router/test_max_output_tokens_cap.py (12 passing). Signed-off-by: Chao Huan --- litellm/router.py | 34 ++++- litellm/types/router.py | 1 + .../proxy/test_credential_slot_registry.py | 2 + .../test_router/test_max_output_tokens_cap.py | 128 ++++++++++++++++++ 4 files changed, 164 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_router/test_max_output_tokens_cap.py diff --git a/litellm/router.py b/litellm/router.py index 3809a5d6555..54c95a48a6d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4931,6 +4931,36 @@ class Router: if "tool_choice" not in kwargs and dep_params.get("tool_choice") is not None: kwargs["tool_choice"] = dep_params["tool_choice"] + @staticmethod + def _apply_max_output_tokens_cap_to_request( + deployment: Mapping[str, object], + kwargs: dict[str, object], # mutable-ok: the shared request mapping is updated in place + ) -> None: + """ + Lower caller-requested completion token budgets to the deployment's + max_output_tokens_cap, when set (see litellm_params.max_output_tokens_cap). + + Only lowers what the request already carries: a smaller caller budget is + respected, and requests without an output budget are untouched (a default + belongs in router default_litellm_params). Applies to the OpenAI token + budget fields so provider translation can normalize it afterwards. + """ + dep_params_raw: Final[object] = deployment.get("litellm_params", {}) or {} + dep_params: Final = ( + dep_params_raw + if isinstance(dep_params_raw, Mapping) + else dep_params_raw.model_dump(exclude_none=True) + ) + cap: Final[object] = dep_params.get("max_output_tokens_cap") + if not isinstance(cap, int) or isinstance(cap, bool) or cap <= 0: + return + for name in ("max_tokens", "max_completion_tokens", "max_output_tokens"): + requested = kwargs.get(name) + if isinstance(requested, bool) or not isinstance(requested, int): + continue + if requested > cap: + kwargs[name] = cap # rebind-ok: the clamped budget rides the shared request mapping to the provider call + def _update_kwargs_with_deployment( self, deployment: dict, @@ -4939,10 +4969,11 @@ class Router: splits_stream_timeout: bool = False, ) -> None: """ - 5 jobs: + 6 jobs: - Adds selected deployment, model_info and api_base to kwargs["metadata"] (used for logging) - Adds default litellm params to kwargs, if set. - Merges tools from deployment with request (proxy-configured tools + request tools). + - Lowers caller-requested completion token budgets to the deployment's max_output_tokens_cap, when set. - On a fallback hop, drops the encrypted reasoning this deployment cannot decrypt, keeping its summary. - On a Responses continuation, sends the partial response as prefill where this deployment takes one. """ @@ -5141,6 +5172,7 @@ class Router: ): kwargs[STREAM_USAGE_INJECTED_FIELD] = True self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name) + Router._apply_max_output_tokens_cap_to_request(deployment=deployment, kwargs=kwargs) hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs) if _is_fallback_hop(hop_kwargs): EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( diff --git a/litellm/types/router.py b/litellm/types/router.py index 60c9e596588..f8972b25fb2 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -717,6 +717,7 @@ class LiteLLM_Params(GenericLiteLLMParams): def __init__(self, *, model: str, **data: object) -> None: ... # kwargs-ok: provider fields vary by model model: str + max_output_tokens_cap: int | None = None model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) def __contains__(self, key) -> bool: diff --git a/tests/unit/proxy/test_credential_slot_registry.py b/tests/unit/proxy/test_credential_slot_registry.py index 6332038feb8..61d1f3b6c7d 100644 --- a/tests/unit/proxy/test_credential_slot_registry.py +++ b/tests/unit/proxy/test_credential_slot_registry.py @@ -128,6 +128,7 @@ DEPLOYMENT_PARAM_CLASSIFICATION: Final[Mapping[str, Classification]] = MappingPr "litellm_credential_name": NotSecret("name of a credentials table entry, not a credential"), "default_api_key_tpm_limit": NotSecret("rate limit number"), "default_api_key_rpm_limit": NotSecret("rate limit number"), + "max_output_tokens_cap": NotSecret("max completion-token budget, not a credential"), "valkey_password": Unplanted(), "anthropic_identity_token": Unplanted(), "anthropic_identity_token_file": Unplanted(), @@ -156,6 +157,7 @@ REQUEST_BODY_PARAM_CLASSIFICATION: Final[Mapping[str, Classification]] = Mapping "litellm_credential_name": NotSecret("name of a credentials table entry, not a credential"), "default_api_key_tpm_limit": NotSecret("rate limit number"), "default_api_key_rpm_limit": NotSecret("rate limit number"), + "max_output_tokens_cap": NotSecret("max completion-token budget, not a credential"), } ) diff --git a/tests/unit/test_router/test_max_output_tokens_cap.py b/tests/unit/test_router/test_max_output_tokens_cap.py new file mode 100644 index 00000000000..72650173d21 --- /dev/null +++ b/tests/unit/test_router/test_max_output_tokens_cap.py @@ -0,0 +1,128 @@ +""" +Unit tests for per-deployment litellm_params.max_output_tokens_cap. + +The cap only lowers output-token budgets the caller already requested; smaller +caller budgets win, and requests without any budget field are left alone (a +default budget belongs in router default_litellm_params). +""" + +import pytest + +from litellm import Router +from litellm.types.router import LiteLLM_Params + +MODEL_LIST = [ + { + "model_name": "capped-model", + "litellm_params": { + "model": "openai/capped-model", + "api_base": "http://localhost:9911/v1", + "api_key": "sk-not-real", + "max_output_tokens_cap": 500, + }, + }, + { + "model_name": "uncapped-model", + "litellm_params": { + "model": "openai/uncapped-model", + "api_base": "http://localhost:9911/v1", + "api_key": "sk-not-real", + }, + }, +] + +BUDGET_FIELDS = ("max_tokens", "max_completion_tokens", "max_output_tokens") + + +def _deployment(model_name: str) -> dict: + for entry in MODEL_LIST: + if entry["model_name"] == model_name: + return entry + raise ValueError(model_name) + + +def _apply_capped(kwargs: dict) -> dict: + Router._apply_max_output_tokens_cap_to_request(deployment=_deployment("capped-model"), kwargs=kwargs) + return kwargs + + +@pytest.mark.parametrize("field", BUDGET_FIELDS) +def test_cap_lowers_oversized_budget(field): + kwargs = _apply_capped({field: 999_999}) + assert kwargs[field] == 500 + + +@pytest.mark.parametrize("field", BUDGET_FIELDS) +def test_cap_keeps_smaller_budget(field): + kwargs = _apply_capped({field: 200}) + assert kwargs[field] == 200 + + +def test_cap_lowers_every_oversized_budget_in_one_request(): + kwargs = _apply_capped( + { + "max_tokens": 999_999, + "max_completion_tokens": 50_000, + "max_output_tokens": 10_000, + } + ) + assert kwargs["max_tokens"] == 500 + assert kwargs["max_completion_tokens"] == 500 + assert kwargs["max_output_tokens"] == 500 + + +def test_request_without_budget_is_untouched(): + kwargs = _apply_capped({"messages": [{"role": "user", "content": "hi"}]}) + assert "max_tokens" not in kwargs + assert "max_completion_tokens" not in kwargs + assert "max_output_tokens" not in kwargs + + +def test_uncapped_deployment_is_untouched(): + kwargs = {"max_tokens": 999_999} + Router._apply_max_output_tokens_cap_to_request(deployment=_deployment("uncapped-model"), kwargs=kwargs) + assert kwargs["max_tokens"] == 999_999 + + +def test_non_int_cap_values_are_ignored(): + for bad_cap in ("500", 500.0, True, 0, -1, None): + deployment = { + "model_name": "capped-model", + "litellm_params": {"model": "openai/capped-model", "max_output_tokens_cap": bad_cap}, + } + kwargs = {"max_tokens": 999_999} + Router._apply_max_output_tokens_cap_to_request(deployment=deployment, kwargs=kwargs) + assert kwargs["max_tokens"] == 999_999, bad_cap + + +def test_pydantic_deployment_params_are_supported(): + deployment = { + "model_name": "capped-model", + "litellm_params": LiteLLM_Params( + model="openai/capped-model", + api_base="http://localhost:9911/v1", + api_key="sk-not-real", + max_output_tokens_cap=500, + ), + } + kwargs = {"max_tokens": 999_999} + Router._apply_max_output_tokens_cap_to_request(deployment=deployment, kwargs=kwargs) + assert kwargs["max_tokens"] == 500 + + +def test_cap_applies_after_default_litellm_params_at_deployment_time(): + """ + End to end through Router._update_kwargs_with_deployment: a default + max_tokens from router default_litellm_params must also be clamped, so the + cap is a hard ceiling for every path that adds a budget. + """ + router = Router( + model_list=MODEL_LIST, + default_litellm_params={"max_tokens": 999_999}, + ) + kwargs = { + "model": "capped-model", + "messages": [{"role": "user", "content": "hi"}], + } + router._update_kwargs_with_deployment(deployment=_deployment("capped-model"), kwargs=kwargs) + assert kwargs["max_tokens"] == 500 # default applied, then capped