diff --git a/litellm/router.py b/litellm/router.py index 0d4045e6ca8..c8153b596d6 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1122,6 +1122,27 @@ def as_output_cap(value: object) -> int | None: return cap if cap >= 0 else None +def stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str, int] | None: + """None means unstamped; an empty mapping means the original caller supplied no limit.""" + stamped: Final = next( + ( + bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) + for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) + if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket + ), + None, + ) + if not isinstance(stamped, dict): + return None + return MappingProxyType( + { + carrier: cap + for carrier, value in stamped.items() + if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None + } + ) + + def _exact_base_model_max_input_tokens(base_model: str) -> int | None: try: model_info: Final = litellm.get_model_info(model=base_model) @@ -15235,9 +15256,7 @@ class Router: request_kwargs: dict, responses_call: bool, ) -> bool: - """Apply a routing strategy's per-tier litellm_params on top of the request and report - whether they pinned an output ceiling, so the caller can hand the request its own ceiling - back on a routing pass that pins none.""" + """Apply per-tier litellm_params and report whether they pinned an output ceiling.""" if not tier_litellm_params: return False reject_api_key_secret_reference(tier_litellm_params) @@ -15250,34 +15269,26 @@ class Router: return not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(surface_tier_params) @staticmethod - def _restore_client_ceiling_no_tier_pins(request_kwargs: MutableMapping[str, object]) -> None: - """A model-group fallback re-enters routing with the kwargs an earlier auto-router pass - already rewrote, so a ceiling sized for that pass's tier would ride onto a group no tier - chose. When this pass pins none, hand the request back exactly the carriers the caller - sent, which the first pinning pass stamped. The stamp lives in a metadata bucket a - caller can also write, so the proxy strips the key at ingestion and this read takes - nothing but the three ceiling carriers as integers: no other key ever reaches kwargs.""" - stamped: Final = next( - ( - bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) - for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) - if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket - ), - None, - ) - if not isinstance(stamped, dict): + def _restore_client_output_ceiling(request_kwargs: MutableMapping[str, object]) -> None: + """Restore caller intent before routing, so neither a previous tier pin nor alias defaults + become the next attempt's budget. The proxy strips caller-written stamps at ingestion.""" + callers_ceiling: Final = stamped_client_ceiling(request_kwargs) + if callers_ceiling is None: return - callers_ceiling: Final = MappingProxyType( - { - carrier: cap - for carrier, value in stamped.items() - if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None - } - ) for carrier in OUTPUT_TOKEN_CEILING_PARAMS: request_kwargs.pop(carrier, None) request_kwargs.update(callers_ceiling) + @staticmethod + def _stamp_client_output_ceiling( + request_kwargs: dict[str, object], # mutable-ok: preserve request-owned metadata across routing attempts + ) -> None: + _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) + metadata_bucket.setdefault( + CLIENT_OUTPUT_CEILING_METADATA_KEY, + {carrier: request_kwargs[carrier] for carrier in OUTPUT_TOKEN_CEILING_PARAMS if carrier in request_kwargs}, + ) + @staticmethod def _drop_client_carriers_a_tier_pin_supersedes( request_kwargs: dict[str, object], @@ -15294,15 +15305,7 @@ class Router: one setting under three names, and a provider handed two of them either rejects the request or picks one by iteration order.""" if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_litellm_params): - _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) - metadata_bucket.setdefault( - CLIENT_OUTPUT_CEILING_METADATA_KEY, - { - carrier: request_kwargs[carrier] - for carrier in OUTPUT_TOKEN_CEILING_PARAMS - if carrier in request_kwargs - }, - ) + Router._stamp_client_output_ceiling(request_kwargs) for carrier in OUTPUT_TOKEN_CEILING_PARAMS: request_kwargs.pop(carrier, None) if "reasoning_effort" not in tier_litellm_params: @@ -15374,14 +15377,12 @@ class Router: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages record_pre_routing_selection(request_kwargs, model) - tier_pins_ceiling: Final = self._pin_tier_params_onto_request( + self._pin_tier_params_onto_request( model=model, tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, request_kwargs=request_kwargs, responses_call=responses_call, ) - if not tier_pins_ceiling: - self._restore_client_ceiling_no_tier_pins(request_kwargs) ######################################################### # Resolve the strategy and logger AFTER the pre-routing hook, since @@ -15539,14 +15540,12 @@ class Router: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages record_pre_routing_selection(request_kwargs, model) - tier_pins_ceiling: Final = self._pin_tier_params_onto_request( + self._pin_tier_params_onto_request( model=model, tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, request_kwargs=request_kwargs, responses_call=responses_call, ) - if not tier_pins_ceiling: - self._restore_client_ceiling_no_tier_pins(request_kwargs) # 2. Get healthy deployments healthy_deployments: Final = await self.async_get_healthy_deployments( @@ -15896,6 +15895,7 @@ class Router: spend metadata is stamped before routing and the response carries the tier group the strategy picked. """ + self._restore_client_output_ceiling(request_kwargs) self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) requested_registered_model_name: Final = self.get_model_from_alias(model=model) or model registered_model_name: Final = await self._resolve_claude_code_session_router( @@ -16062,7 +16062,8 @@ class Router: # them) until `_update_kwargs_with_deployment` drops any the selected # deployment sets itself (its own `aws_region_name` beats the marker's). # Per-tier `litellm_params` on the hook response are deliberate overrides - # the caller applies on top, so those keys are never forwarded here. + # the caller applies on top, so those keys and equivalent output-limit + # names are never tracked as alias defaults here. marker_params: Final = ( self._forwardable_alias_marker_params( model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs @@ -16073,9 +16074,16 @@ class Router: tier_param_keys: Final = ( tuple(pre_routing_hook_response.litellm_params or ()) if pre_routing_hook_response is not None else () ) - newly_forwarded: Final = tuple( - (key, value) for key, value in marker_params if key not in request_kwargs and key not in tier_param_keys + overridden_marker_keys: Final = frozenset(tier_param_keys) | ( + OUTPUT_TOKEN_CEILING_PARAMS if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_param_keys) else frozenset() ) + newly_forwarded: Final = tuple( + (key, value) + for key, value in marker_params + if key not in request_kwargs and key not in overridden_marker_keys + ) + if pre_routing_hook_response is not None: + self._stamp_client_output_ceiling(request_kwargs) request_kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, None) request_kwargs.update(newly_forwarded) if newly_forwarded: diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 7567a0ef37e..08d39841b61 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -107,6 +107,20 @@ Session affinity, user-turn classification, adaptive routing, routing plugins, c When cache costs change the model, the routing decision reports `cause: prompt_cache_cost`. Its signals include the original model, classification cause and both estimated costs +### Output token limits + +Set `max_tokens_from_tier_model` under `complexity_router_config` to choose how a routed request's output limit is determined: + +| Setting | Effective output limit | +| --- | --- | +| `true` (default) | Replace the caller's limit with the selected model group's ceiling | +| `"clamp"` | Use the smaller of the caller's positive limit and the selected model group's ceiling | +| `false` | Keep the caller's limit | + +Clamp reads `max_tokens`, `max_completion_tokens` and `max_output_tokens`. If several are supplied, it uses the smallest positive value. With no positive caller limit, it uses the model ceiling. For example, a caller requesting 100 tokens keeps that limit on a model with a 32,768-token ceiling; a request for 100,000 is capped at 32,768. Rerouting uses the original caller budget, and cache-aware cost estimates use the same effective limit. + +All modes preserve explicit output limits in a tier's own `litellm_params`. An unknown ceiling leaves the caller's value unchanged. In the dashboard, select **Respect caller budget** under **Advanced settings > Compatibility > Output token limit** to enable clamp. + ### Capability forecasting Set `classifier_type: capability` to use diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 12968bf7d4a..f5fd4ae9c98 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -23,6 +23,7 @@ import random import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from functools import partial from itertools import accumulate, chain, islice, takewhile from threading import Lock from types import MappingProxyType @@ -161,6 +162,17 @@ def _tier_name(tier: ComplexityTier | str) -> str: return tier.value if isinstance(tier, ComplexityTier) else tier +def _callers_output_cap(request_kwargs: Mapping[str, object]) -> int | None: + """The smallest positive output ceiling the caller sent under any of its three names, read + from the stamp when an earlier routing pass already rewrote the kwargs.""" + from litellm.router import as_output_cap, stamped_client_ceiling + + stamped: Final = stamped_client_ceiling(request_kwargs) + carriers: Final = request_kwargs if stamped is None else stamped + caps: Final = tuple(as_output_cap(carriers.get(carrier)) for carrier in OUTPUT_TOKEN_CEILING_PARAMS) + return min((cap for cap in caps if cap is not None and cap > 0), default=None) + + def _built_in_tier_or_none(tier_name: str) -> ComplexityTier | None: """The built-in tier a `tiers` key names, or None when the key is an operator-defined name.""" return ComplexityTier.__members__.get(tier_name) @@ -3028,7 +3040,9 @@ class ComplexityRouter(CustomLogger): raise ValueError(f"No model configured for tier {tier_key} and no default_model set") - def _litellm_params_for_model(self, tier: ComplexityTier | str | None, model: str) -> Mapping[str, object]: + def _litellm_params_for_model( + self, tier: ComplexityTier | str | None, model: str, request_kwargs: Mapping[str, object] + ) -> Mapping[str, object]: entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ()) if tier is not None else () entry: Final = next((candidate for candidate in entries if candidate.model_name == model), None) explicit: Final = entry.litellm_params if entry is not None else MappingProxyType({}) @@ -3037,7 +3051,12 @@ class ComplexityRouter(CustomLogger): ceiling: Final = self._group_output_ceiling(model) if ceiling is None: return explicit - return MappingProxyType({**explicit, "max_tokens": ceiling}) + callers_cap: Final = ( + _callers_output_cap(request_kwargs) if self.config.max_tokens_from_tier_model == "clamp" else None + ) + return MappingProxyType( + {**explicit, "max_tokens": ceiling if callers_cap is None else min(callers_cap, ceiling)} + ) @staticmethod def _pick_from_tier_value(model: str | Sequence[str], tier_key: str) -> str: @@ -3768,7 +3787,7 @@ class ComplexityRouter(CustomLogger): response: PreRoutingHookResponse, messages: list[dict[str, object]] | None, # mutable-ok: forwarded verbatim to the list-typed re-pick resolved_messages: Sequence[Mapping[str, object]] | None, - request_kwargs: dict, # mutable-ok: same shape the hook receives + request_kwargs: dict[str, object], # mutable-ok: same shape the hook receives context_fit: _RequestContextFit | None = None, retained_pin: _SessionAffinityPin | None = None, ) -> PreRoutingHookResponse: @@ -3865,7 +3884,7 @@ class ComplexityRouter(CustomLogger): conversation_continuing=bool(decision.get("conversation_continuing", True)) if decision is not None else True, - tier_litellm_params=self._litellm_params_for_model(new_tier, new_model), + tier_litellm_params=self._litellm_params_for_model(new_tier, new_model, request_kwargs), context_escalation_original_tier=( decision.get("context_escalation_original_tier") if decision is not None else None ), @@ -3876,7 +3895,7 @@ class ComplexityRouter(CustomLogger): return HookResponse( model=new_model, messages=response.messages, - litellm_params=self._litellm_params_for_model(new_tier, new_model), + litellm_params=self._litellm_params_for_model(new_tier, new_model, request_kwargs), routing_decision=new_decision, ) @@ -3904,7 +3923,7 @@ class ComplexityRouter(CustomLogger): model_name: str, messages: list[dict[str, Any]] | None, # mutable-ok: forwarded verbatim to the router's own probe input: str | list | None, # mutable-ok: mirrors the owner's own input parameter, which this forwards verbatim - request_kwargs: dict, # mutable-ok: same shape the hook receives + request_kwargs: dict[str, object], # mutable-ok: same shape the hook receives ) -> bool: """Whether the router would find a deployment for this group ON THIS REQUEST. @@ -3962,7 +3981,7 @@ class ComplexityRouter(CustomLogger): messages: list[dict[str, object]] | None, # mutable-ok: forwarded verbatim to the list-typed re-pick input: str | list | None, # mutable-ok: mirrors the owner's own input parameter, which this forwards verbatim resolved_messages: Sequence[Mapping[str, object]] | None, - request_kwargs: dict, # mutable-ok: same shape the hook receives + request_kwargs: dict[str, object], # mutable-ok: same shape the hook receives context_fit: _RequestContextFit | None = None, retained_pin: _SessionAffinityPin | None = None, ) -> PreRoutingHookResponse: @@ -4045,14 +4064,14 @@ class ComplexityRouter(CustomLogger): classifier_model=decision.get("classifier_model"), classifier_cost=decision.get("classifier_cost"), conversation_continuing=bool(decision.get("conversation_continuing", True)), - tier_litellm_params=self._litellm_params_for_model(candidate_tier, new_model), + tier_litellm_params=self._litellm_params_for_model(candidate_tier, new_model, request_kwargs), context_escalation_original_tier=decision.get("context_escalation_original_tier"), previous_decision=decision, ) return response.model_copy( update={ "model": new_model, - "litellm_params": self._litellm_params_for_model(candidate_tier, new_model), + "litellm_params": self._litellm_params_for_model(candidate_tier, new_model, request_kwargs), "routing_decision": new_decision, } ) @@ -4090,14 +4109,14 @@ class ComplexityRouter(CustomLogger): classifier_model=decision.get("classifier_model"), classifier_cost=decision.get("classifier_cost"), conversation_continuing=bool(decision.get("conversation_continuing", True)), - tier_litellm_params=self._litellm_params_for_model(None, default_model), + tier_litellm_params=self._litellm_params_for_model(None, default_model, request_kwargs), context_escalation_original_tier=decision.get("context_escalation_original_tier"), previous_decision=decision, ) return response.model_copy( update={ "model": default_model, - "litellm_params": self._litellm_params_for_model(None, default_model), + "litellm_params": self._litellm_params_for_model(None, default_model, request_kwargs), "routing_decision": default_decision, } ) @@ -4580,7 +4599,9 @@ class ComplexityRouter(CustomLogger): if pin_placement is not None and pin_context_original_tier is not None else pin_source_tier ) - session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model) + session_tier_litellm_params: Final = self._litellm_params_for_model( + routed_pin_tier, routed_model, request_kwargs + ) has_original_messages: Final = messages is not None and len(messages) > 0 return self._with_session_deployment_affinity( await self._gate_response_health( @@ -4678,14 +4699,14 @@ class ComplexityRouter(CustomLogger): choice: Final = await choose_cached_model( router=self.litellm_router_instance, config=self.config, - params_for_model=self._litellm_params_for_model, + params_for_model=partial(self._litellm_params_for_model, request_kwargs=request_kwargs), response=response, request_kwargs=request_kwargs, messages=messages, ) if choice is None or not context_fit.accepts(choice.model): return response - params: Final = self._litellm_params_for_model(choice.tier, choice.model) + params: Final = self._litellm_params_for_model(choice.tier, choice.model, request_kwargs) decision: Final[StandardLoggingRoutingDecision] = { **response.routing_decision, "routed_model": choice.model, @@ -4714,7 +4735,7 @@ class ComplexityRouter(CustomLogger): async def _classify_and_route( self, model: str, - request_kwargs: dict, + request_kwargs: dict[str, object], messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, @@ -4772,7 +4793,7 @@ class ComplexityRouter(CustomLogger): ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs ) fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM - default_tier_params: Final = self._litellm_params_for_model(fallback_tier, routed_model) + default_tier_params: Final = self._litellm_params_for_model(fallback_tier, routed_model, request_kwargs) return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, @@ -4810,7 +4831,7 @@ class ComplexityRouter(CustomLogger): _tier_name(plan_floor), routed_model, ) - plan_tier_params: Final = self._litellm_params_for_model(plan_floor, routed_model) + plan_tier_params: Final = self._litellm_params_for_model(plan_floor, routed_model, request_kwargs) return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, @@ -4839,7 +4860,9 @@ class ComplexityRouter(CustomLogger): ) keyword_plan_floored: Final = routed_tier != escalated_tier routed_model = await self._pick_model_for_tier(routed_tier, messages, resolved_messages, request_kwargs) - keyword_tier_litellm_params: Final = self._litellm_params_for_model(routed_tier, routed_model) + keyword_tier_litellm_params: Final = self._litellm_params_for_model( + routed_tier, routed_model, request_kwargs + ) keyword_cause: Final[RoutingDecisionCause] = ( "plan_mode" if keyword_plan_floored @@ -4919,7 +4942,7 @@ class ComplexityRouter(CustomLogger): outcome.signals, fallback_model, ) - fallback_tier_params: Final = self._litellm_params_for_model(None, fallback_model) + fallback_tier_params: Final = self._litellm_params_for_model(None, fallback_model, request_kwargs) return PreRoutingHookResponse( model=fallback_model, messages=messages if has_original_messages else None, @@ -5001,7 +5024,7 @@ class ComplexityRouter(CustomLogger): routed_model, ) - tier_litellm_params: Final = self._litellm_params_for_model(tier, routed_model) + tier_litellm_params: Final = self._litellm_params_for_model(tier, routed_model, request_kwargs) classifier_model: Final = ( f"{outcome.jev_verdict.provider}/{outcome.jev_verdict.model}" if outcome.cause == "jev_classifier" and outcome.jev_verdict is not None diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index ec960a19f79..020f2d33742 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -1462,7 +1462,7 @@ class ComplexityRouterConfig(LiteLLMBaseModel): "wording the built-ins don't cover, or after a client release changes its strings." ), ) - max_tokens_from_tier_model: bool = Field( + max_tokens_from_tier_model: bool | Literal["clamp"] = Field( default=True, description=( "Set max_tokens on every routed request to the output ceiling of the tier model it " @@ -1473,7 +1473,8 @@ class ComplexityRouterConfig(LiteLLMBaseModel): "from each deployment's model_info and then the model cost map; a tier model with a " "deployment whose ceiling is unknown keeps the caller's value. A max_tokens, " "max_completion_tokens or max_output_tokens in the tier's own litellm_params still " - "wins. Set false to forward the caller's value unchanged." + 'wins. Set "clamp" to keep a smaller positive caller value and cap a larger one at ' + "the ceiling. Set false to forward the caller's value unchanged." ), ) route_housekeeping_to_cheapest_tier: bool = Field( diff --git a/tests/unit/proxy/common_utils/test_cache_aware_routing.py b/tests/unit/proxy/common_utils/test_cache_aware_routing.py index 24fc1f22a6a..4eb22b2e34f 100644 --- a/tests/unit/proxy/common_utils/test_cache_aware_routing.py +++ b/tests/unit/proxy/common_utils/test_cache_aware_routing.py @@ -1,5 +1,6 @@ from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from functools import partial from typing import Final import pytest @@ -148,7 +149,7 @@ async def _select( return await select_cached_model( router=router, config=config, - params_for_model=complexity._litellm_params_for_model, + params_for_model=partial(complexity._litellm_params_for_model, request_kwargs=request_kwargs), response=response, body=body, request_kwargs=request_kwargs, @@ -234,7 +235,9 @@ async def test_disabled_setting_does_not_access_prediction_services() -> None: await choose_cached_model( router=_router(), config=config, - params_for_model=ComplexityRouter("smart", _router(), config.model_dump())._litellm_params_for_model, + params_for_model=partial( + ComplexityRouter("smart", _router(), config.model_dump())._litellm_params_for_model, request_kwargs={} + ), response=_response(), request_kwargs={}, messages=None, @@ -348,6 +351,7 @@ def test_repeated_model_in_multiple_tiers_is_only_considered_once() -> None: (True, "custom_endpoint", "cheap"), (True, "compaction", "cheap"), (True, "guardrail", "cheap"), + (True, "clamp", "strong"), ], ) async def test_router_applies_opt_in_and_preserves_failure_semantics( @@ -368,7 +372,9 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( from litellm.router_strategy.complexity_router.context_compaction import initialize_compaction_state config: Final = _config( - cache_aware_routing=enabled, cache_aware_routing_timeout_ms=1 if behavior == "deadline" else 2000 + cache_aware_routing=enabled, + cache_aware_routing_timeout_ms=1 if behavior == "deadline" else 2000, + max_tokens_from_tier_model="clamp" if behavior == "clamp" else True, ) models: Final = _router( strong_output_rate=0.000048 if behavior == "tier_cost" else 0.000015, @@ -423,6 +429,7 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( body: Final = { **_body(), **({"max_tokens": 1000} if behavior in ("tier_cost", "tier_context") else {}), + **({"max_tokens": 100} if behavior == "clamp" else {}), **({"thinking": {"type": "enabled", "budget_tokens": 10000}} if behavior == "unsupported_shape" else {}), } kwargs: Final = { @@ -437,6 +444,7 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( else {} ), **({"guardrails": ["test-guardrail"]} if behavior == "guardrail" else {}), + **({"max_tokens": 100} if behavior == "clamp" else {}), } if expected is None: with pytest.raises(asyncio.CancelledError): @@ -458,6 +466,10 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( if enabled and behavior == "success": assert requests.qsize() == 4 assert response.routing_decision is not None + if behavior == "clamp": + assert "cache-aware:estimated-cost=0.01950000;original-cost=0.06390000" in ( + response.routing_decision["signals"] or () + ), "both arms must be estimated at the caller's 100 output tokens, not the 1024 default" assert response.routing_decision["cause"] == ( "prompt_cache_cost" if expected == "strong" else "heuristic_scorer" ) diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index 2d27a025ffa..78dc2eaf314 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -36,11 +36,12 @@ from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( + CLIENT_OUTPUT_CEILING_METADATA_KEY, OUTPUT_TOKEN_CEILING_PARAMS, RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, ) -from litellm.router import as_output_cap +from litellm.router import as_output_cap, stamped_client_ceiling from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -13733,7 +13734,7 @@ def test_tier_model_params_are_used_by_pools_and_savings_baseline(mock_router_in assert router._tier_pools() == {"SIMPLE": ["mini"], "REASONING": ["opus", "abc"]} assert router._hardest_tier_models() == ("opus", "abc") - assert router._litellm_params_for_model(ComplexityTier.REASONING, "opus") == {"reasoning_effort": "xhigh"} + assert router._litellm_params_for_model(ComplexityTier.REASONING, "opus", {}) == {"reasoning_effort": "xhigh"} @pytest.mark.asyncio @@ -16917,9 +16918,10 @@ class TestMaxTokensFromTierModel: @staticmethod def _router( tier_litellm_params: dict | None = None, - max_tokens_from_tier_model: bool | None = None, + max_tokens_from_tier_model: bool | Literal["clamp"] | None = None, simple_deployments: list[dict] | None = None, extra_config: dict | None = None, + marker_params: Mapping[str, object] | None = None, ) -> Router: simple_tier: dict = {"model_name": "small"} if tier_litellm_params: @@ -16934,7 +16936,11 @@ class TestMaxTokensFromTierModel: model_list=[ { "model_name": "smart-router", - "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": config}, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": config, + **(marker_params or {}), + }, }, *(simple_deployments or [TestMaxTokensFromTierModel.SMALL]), { @@ -17016,6 +17022,167 @@ class TestMaxTokensFromTierModel: assert sent["max_tokens"] == 8192 + @pytest.mark.asyncio + @pytest.mark.parametrize( + "prompt, client_kwargs, expected", + [ + (COMPLEX_PROMPT, {"max_tokens": 2000}, ("anthropic/claude-sonnet-5", 2000)), + ("hi", {"max_tokens": 100000}, ("anthropic/claude-haiku-4-5", 8192)), + (COMPLEX_PROMPT, {}, ("anthropic/claude-sonnet-5", 64000)), + (COMPLEX_PROMPT, {"max_tokens": 0}, ("anthropic/claude-sonnet-5", 64000)), + (COMPLEX_PROMPT, {"max_completion_tokens": 2000}, ("anthropic/claude-sonnet-5", 2000)), + (COMPLEX_PROMPT, {"max_tokens": 3000, "max_completion_tokens": 2000}, ("anthropic/claude-sonnet-5", 2000)), + ], + ids=["below-ceiling", "above-ceiling", "omitted", "zero", "other-carrier", "smallest-carrier"], + ) + async def test_clamp_keeps_the_smaller_of_the_client_value_and_the_ceiling( + self, prompt: str, client_kwargs: dict[str, int], expected: tuple[str, int] + ) -> None: + sent: Final = await self._routed(self._router(max_tokens_from_tier_model="clamp"), prompt, **client_kwargs) + + assert (sent["model"], sent["max_tokens"]) == expected + assert not (OUTPUT_TOKEN_CEILING_PARAMS - {"max_tokens"}) & sent.keys() + + @pytest.mark.asyncio + async def test_clamp_on_the_responses_surface_keeps_its_own_name(self) -> None: + router: Final = self._router(max_tokens_from_tier_model="clamp") + + sent: Final = await self._routed_responses(router, self.COMPLEX_PROMPT, max_output_tokens=2000) + + assert (sent["model"], sent["max_output_tokens"]) == ("anthropic/claude-sonnet-5", 2000) + assert "max_tokens" not in sent + + @pytest.mark.asyncio + @pytest.mark.parametrize("tier_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + @pytest.mark.parametrize("marker_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + @pytest.mark.parametrize("surface", ["chat", "responses"]) + @pytest.mark.parametrize("caller_cap", [None, 100]) + async def test_clamp_leaves_the_operators_own_tier_ceiling_in_charge( + self, + tier_carrier: str, + marker_carrier: str, + surface: str, + caller_cap: int | None, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + surface_key: Final = "max_output_tokens" if surface == "responses" else "max_tokens" + caller_key: Final = "max_tokens" if surface == "responses" else "max_output_tokens" + request_caps: Final = {caller_key: caller_cap} if caller_cap is not None else {} + router: Final = self._router( + tier_litellm_params={tier_carrier: 4321}, + max_tokens_from_tier_model="clamp", + marker_params={marker_carrier: 2000}, + simple_deployments=[{ + "model_name": "small", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "test-only", + "api_base": "https://tier.test/v1", + surface_key: 9000, + }, + "model_info": {"max_output_tokens": 16384}, + }], + ) + + with respx.mock(assert_all_mocked=True) as upstream: + upstream.post(host="tier.test").mock(side_effect=TestHealthFallbackDispatch._http_response) + if surface == "responses": + await router.aresponses(model="smart-router", input="hi", **request_caps) + else: + await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": "hi"}], **request_caps + ) + + assert len(upstream.calls) == 1 + sent: Final = json.loads(upstream.calls[0].request.content) + wire_key: Final = "max_output_tokens" if surface == "responses" else "max_completion_tokens" + assert {key: value for key, value in sent.items() if key in OUTPUT_TOKEN_CEILING_PARAMS} == {wire_key: 4321} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "callers_original, expected", + [({"max_tokens": 2000}, 2000), ({}, 64000)], + ids=["caller-sent-a-value", "caller-sent-nothing"], + ) + async def test_clamp_on_a_rerouted_request_uses_the_callers_original_value( + self, callers_original: dict[str, int], expected: int + ) -> None: + """A retry re-enters routing with kwargs an earlier pass already rewrote; the stamp holds + what the caller actually sent.""" + router: Final = self._router(max_tokens_from_tier_model="clamp") + rewritten: Final[dict[str, object]] = { + "max_tokens": 8192, + "metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: callers_original}, + } + + sent: Final = await self._routed(router, self.COMPLEX_PROMPT, **rewritten) + + assert sent["max_tokens"] == expected + + @pytest.mark.asyncio + @pytest.mark.parametrize("marker_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + @pytest.mark.parametrize("caller_cap", [None, 3000]) + @pytest.mark.parametrize("tier_cap", [None, 4321]) + @pytest.mark.parametrize("known_ceiling", [True, False]) + async def test_rerouting_preserves_caller_budget_without_adopting_marker_defaults( + self, marker_carrier: str, caller_cap: int | None, tier_cap: int | None, known_ceiling: bool + ) -> None: + router: Final = self._router( + max_tokens_from_tier_model="clamp", + marker_params={marker_carrier: 2000}, + tier_litellm_params={"max_tokens": tier_cap} if tier_cap else None, + simple_deployments=None if known_ceiling else [{ + "model_name": "small", + "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, + }], + ) + request: Final[dict[str, object]] = {} if caller_cap is None else {"max_output_tokens": caller_cap} + await router.async_get_available_deployment( + model="smart-router", request_kwargs=request, messages=[{"role": "user", "content": "hi"}] + ) + first: Final = {key: request[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in request} + await router.async_get_available_deployment( + model="smart-router", request_kwargs=request, + messages=[{"role": "user", "content": self.COMPLEX_PROMPT}], + ) + + expected_first: Final = ( + {"max_tokens": tier_cap or caller_cap or 8192} + if known_ceiling or tier_cap else + {marker_carrier: 2000, **({"max_output_tokens": caller_cap} if caller_cap else {})} + ) + assert first == expected_first + assert request["max_tokens"] == (caller_cap or 64000) + assert not (OUTPUT_TOKEN_CEILING_PARAMS - {"max_tokens"}) & request.keys() + + await router.async_get_available_deployment( + model="smart-router", request_kwargs=request, messages=[{"role": "user", "content": "hi"}] + ) + assert {key: request[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in request} == first + + @pytest.mark.asyncio + @pytest.mark.parametrize("mode", [False, "clamp"]) + @pytest.mark.parametrize("marker_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + async def test_unpinned_output_limits_still_use_marker_defaults( + self, mode: bool | Literal["clamp"], marker_carrier: str + ) -> None: + router: Final = self._router( + max_tokens_from_tier_model=mode, + marker_params={marker_carrier: 2000}, + simple_deployments=[{ + "model_name": "small", + "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, + }], + ) + sent: Final = await self._routed(router) + + assert {key: sent[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in sent} == {marker_carrier: 2000} + + def test_an_unknown_mode_is_rejected(self) -> None: + with pytest.raises(ValidationError): + ComplexityRouterConfig(tiers={"SIMPLE": "small"}, max_tokens_from_tier_model="sometimes") + @pytest.mark.asyncio async def test_a_tier_model_with_an_unknown_ceiling_keeps_the_client_value(self): unmapped: dict = {"model_name": "small", "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}} @@ -17112,7 +17279,7 @@ class TestMaxTokensFromTierModel: complexity_router_config={"tiers": {"SIMPLE": "small"}, "default_model": "big"}, ) - assert dict(strategy._litellm_params_for_model(None, "big")) == {"max_tokens": 64000} + assert dict(strategy._litellm_params_for_model(None, "big", {})) == {"max_tokens": 64000} @pytest.mark.asyncio async def test_a_fallback_into_a_plain_group_gets_the_callers_ceiling_back(self): @@ -17228,29 +17395,52 @@ class TestMaxTokensFromTierModel: def test_a_client_cap_is_read_as_an_integer_or_ignored(self, value, expected): assert as_output_cap(value) == expected + @pytest.mark.parametrize( + "request_kwargs, expected", + [ + ({"max_tokens": 8192}, None), + ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {}}}, {}), + ({"litellm_metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_tokens": "2000"}}}, {"max_tokens": 2000}), + ( + {"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_output_tokens": 5, "api_key": "x", "max_tokens": True}}}, + {"max_output_tokens": 5}, + ), + ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: "not-a-dict"}}, None), + ], + ids=["no-stamp", "empty-stamp", "litellm-metadata-bucket", "only-integer-carriers", "malformed-stamp"], + ) + def test_the_stamp_reads_back_only_the_callers_integer_carriers( + self, request_kwargs: dict[str, object], expected: dict[str, int] | None + ) -> None: + stamped: Final = stamped_client_ceiling(request_kwargs) + + assert (None if stamped is None else dict(stamped)) == expected + def test_restoring_the_callers_ceiling_reads_the_stamp_and_replaces_every_carrier(self): stamped: dict = {"max_output_tokens": 500, "metadata": {"_client_output_ceiling": {"max_tokens": 8192}}} - Router._restore_client_ceiling_no_tier_pins(stamped) + Router._restore_client_output_ceiling(stamped) assert {k: v for k, v in stamped.items() if k != "metadata"} == {"max_tokens": 8192} coerced: dict = { "max_tokens": 64000, "metadata": {"_client_output_ceiling": {"max_tokens": "8192", "max_completion_tokens": 100.0}}, } - Router._restore_client_ceiling_no_tier_pins(coerced) + Router._restore_client_output_ceiling(coerced) assert {k: v for k, v in coerced.items() if k != "metadata"} == { "max_tokens": 8192, "max_completion_tokens": 100, } unstamped: dict = {"max_tokens": 64000, "metadata": {}} - Router._restore_client_ceiling_no_tier_pins(unstamped) + Router._restore_client_output_ceiling(unstamped) assert unstamped["max_tokens"] == 64000 @pytest.mark.asyncio async def test_pinning_stamps_the_callers_carriers_once(self): router = self._router() request_kwargs: dict = {"max_completion_tokens": 8192} + router._stamp_client_output_ceiling(request_kwargs) + assert stamped_client_ceiling(request_kwargs) == {"max_completion_tokens": 8192} first = router._pin_tier_params_onto_request( model="big", tier_litellm_params={"max_tokens": 64000}, request_kwargs=request_kwargs, responses_call=False diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 72a08c53267..b59aa051e38 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -452,7 +452,7 @@ export interface ComplexityRouterConfigValue { route_housekeeping_to_cheapest_tier?: boolean; housekeeping_patterns?: string[]; reminder_markers?: ReminderMarkerPair[]; - max_tokens_from_tier_model?: boolean; + max_tokens_from_tier_model?: boolean | "clamp"; classifier_plugin_timeout_ms?: number; } diff --git a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx index 6e7b1489a4c..f596ecfdc62 100644 --- a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx @@ -1,7 +1,28 @@ import { Switch } from "@/components/ui/switch"; +import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import React from "react"; import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +const OUTPUT_TOKEN_POLICIES = [ + { + value: true, + label: "Use model ceiling", + description: "Replace the caller's output limit with the selected model's ceiling. This is the default.", + }, + { + value: "clamp", + label: "Respect caller budget", + description: + "Use the smaller of the caller's output limit and the model ceiling. Use the ceiling when no limit is sent.", + }, + { + value: false, + label: "Pass through caller limit", + description: "Forward the caller's output limit unchanged.", + }, +] as const; + const ResponseFormatControls: React.FC<{ value: ComplexityRouterConfigValue; onChange: (value: ComplexityRouterConfigValue) => void; @@ -18,18 +39,31 @@ const ResponseFormatControls: React.FC<{ Return the resolved underlying model name in responses instead of the autorouter alias. -