diff --git a/litellm/router.py b/litellm/router.py index 20af2e5649e..f6f31e7bec3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5106,6 +5106,36 @@ class Router(DeprecatedRouterMembers): 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, @@ -5114,10 +5144,11 @@ class Router(DeprecatedRouterMembers): 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. """ @@ -5316,6 +5347,7 @@ class Router(DeprecatedRouterMembers): ): 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 f927ae0475b..bcc79009b19 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -724,6 +724,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