From f5cfa842207c1980ddda98817cc50606cb86b5d5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 20 Aug 2026 01:39:02 +0000 Subject: [PATCH] feat(router): allow per-tier litellm_params in complexity autorouter config (#37064) * feat(router): support complexity tier request params Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): make complexity tier params immutable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(router): simplify complexity tier overlays Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): preserve plain tier config round trips Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): mask tier params in routing decisions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 4 + .../complexity_router/README.md | 15 + .../complexity_router/complexity_router.py | 79 ++++- .../complexity_router/config.py | 96 +++++- litellm/types/router.py | 2 + litellm/types/utils.py | 2 + .../router_strategy/test_complexity_router.py | 280 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 17 ++ 8 files changed, 480 insertions(+), 15 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b25b4f92467..ef04423e3ab 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -11191,6 +11191,8 @@ class Router: if pre_routing_hook_response is not None: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages + if pre_routing_hook_response.litellm_params: + request_kwargs.update(pre_routing_hook_response.litellm_params) ######################################################### # Resolve the strategy and logger AFTER the pre-routing hook, since @@ -11300,6 +11302,8 @@ class Router: if pre_routing_hook_response is not None: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages + if pre_routing_hook_response.litellm_params: + request_kwargs.update(pre_routing_hook_response.litellm_params) # 2. Get healthy deployments healthy_deployments: Final = await self.async_get_healthy_deployments( diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 65f1029bf55..cf7bde93360 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -53,6 +53,21 @@ model_list: REASONING: o1-preview ``` +Each tier can also use a model entry with request parameter overrides. A tier value may be +a model string, a single object, or a list mixing strings and objects. Object entries must +contain a model name and may contain any LiteLLM request parameters. The model name must +still resolve to a deployment in `model_list`; this configuration does not create one + +```yaml + tiers: + COMPLEX: opus + REASONING: + - model_name: opus + litellm_params: + reasoning_effort: xhigh + - abc +``` + ### Renaming the tiers `tier_labels` puts your own vocabulary on the four tiers: diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 5d6e13c7fc0..0cb50cf3a3d 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -30,6 +30,7 @@ from litellm.constants import EMPTY_MAPPING, RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata +from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, @@ -663,6 +664,35 @@ class ClassificationOutcome(NamedTuple): classifier_cost: float | None = None +class _SessionAffinityPin(NamedTuple): + model: str + tier: ComplexityTier | None + + +def _parse_session_affinity_pin(value: object) -> _SessionAffinityPin | None: + if isinstance(value, str): + return _SessionAffinityPin(model=value, tier=None) + parts: Final[tuple[object, object] | None] = ( + (value.get("model"), value.get("tier")) + if isinstance(value, Mapping) + else (value[0], value[1]) + if isinstance(value, (list, tuple)) and len(value) == 2 + else None + ) + if parts is None: + return None + model, tier_value = parts + if not isinstance(model, str): + return None + tier: Final = ComplexityTier(tier_value) if isinstance(tier_value, str) else None + return _SessionAffinityPin(model=model, tier=tier) + + +def _session_affinity_cache_value(model: str, tier: ComplexityTier | str | None) -> Mapping[str, str | None]: + tier_value: Final = _tier_name(tier) if tier is not None else None + return {"model": model, "tier": tier_value} # mutable-ok: cache requires JSON mapping + + class ComplexityRouter(CustomLogger): """ Complexity router that classifies requests and routes to appropriate models. @@ -1078,6 +1108,7 @@ class ComplexityRouter(CustomLogger): classifier_model: str | None = None, classifier_cost: float | None = None, conversation_continuing: bool = True, + tier_litellm_params: Mapping[str, object] | None = None, ) -> StandardLoggingRoutingDecision: """Assemble the per-request provenance record for this router's decision. @@ -1127,6 +1158,10 @@ class ComplexityRouter(CustomLogger): decision["classifier_model"] = classifier_model if classifier_cost is not None: decision["classifier_cost"] = classifier_cost + if tier_litellm_params: + masked_tier_litellm_params: Final = mask_credentials_in_payload(tier_litellm_params) + if isinstance(masked_tier_litellm_params, Mapping): + decision["tier_litellm_params"] = masked_tier_litellm_params return decision async def aclassify( @@ -1457,6 +1492,13 @@ 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]: + if tier is None: + return MappingProxyType({}) + entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ()) + entry: Final = next((candidate for candidate in entries if candidate.model_name == model), None) + return entry.litellm_params if entry is not None else MappingProxyType({}) + @staticmethod def _pick_from_tier_value(model: str | list[str], tier_key: str) -> str: if isinstance(model, str): @@ -2068,9 +2110,10 @@ class ComplexityRouter(CustomLogger): cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None if cache_key is not None: - pinned_model: Final = await self.litellm_router_instance.cache.async_get_cache(key=cache_key) - if isinstance(pinned_model, str): - routed_model: str | None = pinned_model + pinned_value: Final = await self.litellm_router_instance.cache.async_get_cache(key=cache_key) + pinned_pin: Final = _parse_session_affinity_pin(pinned_value) + if pinned_pin is not None: + routed_model: str | None = pinned_pin.model pin_escalation_keyword: str | None = None if self.escalation_keywords: user_message: Final = ( @@ -2079,16 +2122,21 @@ class ComplexityRouter(CustomLogger): if user_message is not None: pin_escalation_keyword = self._matched_escalation_keyword(user_message) if pin_escalation_keyword is not None: - routed_model = self._escalated_pin(pinned_model) + routed_model = self._escalated_pin(pinned_pin.model) if routed_model is not None: - escalated: Final = routed_model != pinned_model + escalated: Final = routed_model != pinned_pin.model + resolved_pin_tier: Final = ( + pinned_pin.tier + if not escalated and pinned_pin.tier is not None + else self._tier_for_model(routed_model) + ) # The floor outranks the pin because plan mode is a transient state of the # session, not a request to move it: the turns carrying the sentinel route at # the floor, and the stored pin deliberately keeps the session's own model so # the first turn after plan mode exits auto-routes exactly as it would have. # Escalation is the opposite on purpose -- an explicit ask to re-pin higher. pin_plan_sentinel: Final = self._matched_plan_mode_signal(request_kwargs, resolved_messages) - pinned_tier: Final = self._tier_for_model(routed_model) if pin_plan_sentinel is not None else None + pinned_tier: Final = resolved_pin_tier if pin_plan_sentinel is not None else None plan_floored: Final = ( pinned_tier is not None and self._apply_plan_mode_floor(pinned_tier) != pinned_tier ) @@ -2099,7 +2147,7 @@ class ComplexityRouter(CustomLogger): # pin mid-conversation just because it outlives the original write. await self.litellm_router_instance.cache.async_set_cache( key=cache_key, - value=session_model, + value=_session_affinity_cache_value(session_model, resolved_pin_tier), ttl=self.config.session_affinity_ttl_seconds, ) if self.config.adaptive: @@ -2118,19 +2166,23 @@ class ComplexityRouter(CustomLogger): verbose_router_logger.info( "ComplexityRouter: routing decision cause=%s, routed_model=%s", cause, routed_model ) + routed_pin_tier: Final = self._tier_for_model(routed_model) if plan_floored else resolved_pin_tier + session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model) has_original_messages: Final = messages is not None and len(messages) > 0 return self._with_session_deployment_affinity( PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, + litellm_params=session_tier_litellm_params, routing_decision=self._build_routing_decision( routed_model=routed_model, cause=cause, - tier=self._tier_for_model(routed_model), + tier=routed_pin_tier, matched_keyword=pin_plan_sentinel if plan_floored else None, escalation_keyword=pin_escalation_keyword, escalated=escalated, conversation_continuing=conversation_continuing, + tier_litellm_params=session_tier_litellm_params, ), ) ) @@ -2157,7 +2209,10 @@ class ComplexityRouter(CustomLogger): if pinnable and cache_key is not None and response is not None: await self.litellm_router_instance.cache.async_set_cache( key=cache_key, - value=response.model, + value=_session_affinity_cache_value( + response.model, + response.routing_decision.get("tier") if response.routing_decision is not None else None, + ), ttl=self.config.session_affinity_ttl_seconds, ) return self._with_session_deployment_affinity(response) @@ -2271,6 +2326,7 @@ 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_cause: Final[RoutingDecisionCause] = ( "plan_mode" if keyword_plan_floored @@ -2286,6 +2342,7 @@ class ComplexityRouter(CustomLogger): return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, + litellm_params=keyword_tier_litellm_params, routing_decision=self._build_routing_decision( routed_model=routed_model, conversation_continuing=conversation_continuing, @@ -2294,6 +2351,7 @@ class ComplexityRouter(CustomLogger): matched_keyword=plan_mode_sentinel if keyword_plan_floored else override.matched_keyword, escalation_keyword=escalation_keyword, escalated=keyword_escalated, + tier_litellm_params=keyword_tier_litellm_params, ), ) @@ -2380,6 +2438,7 @@ class ComplexityRouter(CustomLogger): routed_model, ) + tier_litellm_params: Final = self._litellm_params_for_model(tier, routed_model) classifier_model: Final = ( self.config.classifier_llm_config.model if outcome.cause == "llm_classifier" and self.config.classifier_llm_config is not None @@ -2405,6 +2464,7 @@ class ComplexityRouter(CustomLogger): return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, + litellm_params=tier_litellm_params, routing_decision=self._build_routing_decision( routed_model=routed_model, conversation_continuing=conversation_continuing, @@ -2417,5 +2477,6 @@ class ComplexityRouter(CustomLogger): escalated=escalated, classifier_model=classifier_model, classifier_cost=outcome.classifier_cost, + tier_litellm_params=tier_litellm_params, ), ) diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index e82ae991100..73f1378e5f7 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -5,10 +5,12 @@ Contains default keyword lists, weights, tier boundaries, and configuration clas All values are configurable via proxy config.yaml. """ +from collections.abc import Mapping from enum import Enum -from typing import Final, Literal +from types import MappingProxyType +from typing import Annotated, Final, Literal -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin @@ -159,6 +161,44 @@ class ReminderMarkerPair(BaseModel): return self +class ComplexityTierModel(BaseModel): + model_config = ConfigDict(frozen=True) + + model_name: str + litellm_params: Annotated[Mapping[str, object], SkipValidation()] = Field( + default_factory=lambda: MappingProxyType({}) + ) + + @field_validator("litellm_params", mode="before") + @classmethod + def _freeze_litellm_params(cls, value: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType(dict(value)) + + @field_serializer("litellm_params") + def _serialize_litellm_params(self, value: Mapping[str, object]) -> Mapping[str, object]: + return dict(value) # mutable-ok: Pydantic JSON serialization requires a concrete mapping + + +def _normalize_tier_entries( + raw_value: object, + tier: str, +) -> tuple[str | list[str], tuple[ComplexityTierModel, ...]]: + raw_entries: Final = raw_value if isinstance(raw_value, (list, tuple)) else (raw_value,) + entries: Final = tuple( + ComplexityTierModel(model_name=entry) if isinstance(entry, str) else ComplexityTierModel.model_validate(entry) + for entry in raw_entries + ) + model_names: Final = tuple(entry.model_name for entry in entries) + if len(model_names) != len(frozenset(model_names)): + raise ValueError(f"tier {tier} contains duplicate model_name values; each pool entry needs distinct parameters") + normalized: Final = ( + entries[0].model_name + if not isinstance(raw_value, (list, tuple)) + else list(model_names) # mutable-ok: config.tiers must preserve its existing list contract + ) + return normalized, entries + + # ─── Default Keyword Lists ─── # Note: Keywords should be full words/phrases to avoid substring false positives. # The matching logic uses word boundary detection for single-word keywords. @@ -425,6 +465,9 @@ class ComplexityRouterConfig(BaseModel): "A list is randomly picked from when adaptive=False, and used as a soft-floor home pool when adaptive=True" ), ) + tier_model_configs: Mapping[str, tuple[ComplexityTierModel, ...]] = Field( + default_factory=dict, + ) tier_definitions: tuple[TierDefinition, ...] | None = Field( default=None, @@ -777,6 +820,55 @@ class ComplexityRouterConfig(BaseModel): coerced[key] = item return coerced + @model_validator(mode="before") + @classmethod + def _normalize_tier_model_configs(cls, value: object) -> object: + if not isinstance(value, dict): + return value + raw_tiers: Final = value.get("tiers") + if not isinstance(raw_tiers, dict): + return value + existing_configs: Final = value.get("tier_model_configs") + normalized_entries: Final = MappingProxyType( + {tier: _normalize_tier_entries(raw_value, tier) for tier, raw_value in raw_tiers.items()} + ) + normalized_tiers: Final = MappingProxyType( + {tier: normalized for tier, (normalized, _) in normalized_entries.items()} + ) + incoming_params: Final = ( + MappingProxyType( + { + (tier, entry.model_name): entry.litellm_params + for tier, entries in existing_configs.items() + for entry in (ComplexityTierModel.model_validate(item) for item in entries) + } + ) + if isinstance(existing_configs, dict) + else MappingProxyType({}) + ) + tier_model_configs: Final = MappingProxyType( + { + tier: tuple( + entry.model_copy( + update=MappingProxyType( + { + "litellm_params": incoming_params.get((tier, entry.model_name), entry.litellm_params), + } + ) + ) + for entry in entries + ) + for tier, (_, entries) in normalized_entries.items() + if any(entry.litellm_params for entry in entries) + or (isinstance(existing_configs, dict) and tier in existing_configs) + } + ) + return { # mutable-ok: Pydantic before-validator requires a concrete mapping + **value, + "tiers": normalized_tiers, + "tier_model_configs": tier_model_configs, + } + @field_validator("escalation_keywords") @classmethod def _normalize_escalation_keywords(cls, value: list[str] | None) -> list[str] | None: diff --git a/litellm/types/router.py b/litellm/types/router.py index 7d1dd1358d5..99a4603ae49 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -4,6 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc import datetime import enum +from collections.abc import Mapping from dataclasses import dataclass from typing import Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints @@ -897,6 +898,7 @@ class PreRoutingHookResponse(BaseModel): messages: list[dict[str, Any]] | None routing_decision: StandardLoggingRoutingDecision | None = None session_affinity_ttl_seconds: int | None = None + litellm_params: Mapping[str, object] | None = None _PreRoutingStrategyT_co = TypeVar("_PreRoutingStrategyT_co", covariant=True) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 72ca68b574c..96b9343353d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2843,6 +2843,7 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): conversation_continuing: bool savings_baseline_model: str savings_baseline_deployment_id: str + tier_litellm_params: Mapping[str, object] # writable-ok: Pydantic warns on ReadOnly TypedDict fields # Fields whose values quote the caller's prompt. Dropped when an operator turns message @@ -2868,6 +2869,7 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset( "conversation_continuing", "savings_baseline_model", "savings_baseline_deployment_id", + "tier_litellm_params", } ) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 7586bbf551e..006fb79b08f 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2166,6 +2166,38 @@ class TestRouterPreRoutingAliasOverrides: assert request_kwargs["drop_params"] is True assert request_kwargs["cache_control_injection_points"] == [{"location": "message", "role": "system"}] + @pytest.mark.asyncio + async def test_tier_litellm_params_are_applied_before_deployment_selection(self): + router = Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": { + "SIMPLE": { + "model_name": "gpt-4o-mini", + "litellm_params": {"reasoning_effort": "xhigh"}, + } + } + }, + }, + }, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, + ] + ) + request_kwargs: Dict = {"reasoning_effort": "low"} + + deployment = await router.async_get_available_deployment( + model="smart-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert deployment["model_name"] == "gpt-4o-mini" + assert request_kwargs["reasoning_effort"] == "xhigh" + @pytest.mark.asyncio async def test_alias_custom_pricing_is_not_applied_to_request_kwargs(self): """Custom pricing on the alias prices the alias, not the tier deployment @@ -3851,7 +3883,7 @@ class TestSessionAffinity: cache.async_set_cache.assert_called_once() call_kwargs = cache.async_set_cache.call_args.kwargs assert call_kwargs["ttl"] == 120 - assert call_kwargs["value"] == "gpt-4o-mini" + assert call_kwargs["value"] == {"model": "gpt-4o-mini", "tier": "SIMPLE"} @pytest.mark.asyncio async def test_ttl_refreshed_on_cache_hit(self, mock_router_instance, basic_config): @@ -3875,7 +3907,7 @@ class TestSessionAffinity: assert result.model == "o1-preview" cache.async_set_cache.assert_called_once() call_kwargs = cache.async_set_cache.call_args.kwargs - assert call_kwargs["value"] == "o1-preview" + assert call_kwargs["value"] == {"model": "o1-preview", "tier": "REASONING"} assert call_kwargs["ttl"] == 90 @pytest.mark.asyncio @@ -5545,12 +5577,14 @@ class TestRedactedLoggingDropsPromptText: "tier_boundaries": {"simple_medium": 0.15, "medium_complex": 0.35, "complex_reasoning": 0.6}, "classifier_model": "claude-haiku", "escalated": True, + "tier_litellm_params": {"reasoning_effort": "xhigh"}, "signals": ["code (python)"], "matched_keyword": "deploy to k8s", "escalation_keyword": "LITELLM ESCALATE", } kept = Router._redact_prompt_text_if_needed(request_kwargs={}, routing_decision=full) assert set(full) - set(kept) == {"signals", "matched_keyword", "escalation_keyword"} + assert kept["tier_litellm_params"] == {"reasoning_effort": "xhigh"} @pytest.mark.asyncio async def test_redaction_via_request_header_is_honored(self): @@ -7273,8 +7307,6 @@ class TestClassificationRubrics: }, ) assert config.classifier_llm_config.system_prompt == "Grade the data sensitivity of the request." - - def _custom_tier_config(**overrides) -> Dict: """A valid operator-defined tier set: two built-in names plus one custom tier.""" return { @@ -8123,3 +8155,243 @@ class TestPlanModeTierFloor: assert result.model == "gpt-4o" assert result.routing_decision is not None assert result.routing_decision["tier"] == "MEDIUM" +def test_tier_model_params_are_normalized_without_changing_model_pools(): + config = ComplexityRouterConfig( + tiers={ + "SIMPLE": "mini", + "REASONING": [ + {"model_name": "opus", "litellm_params": {"reasoning_effort": "xhigh"}}, + "abc", + ], + } + ) + + assert config.tiers == {"SIMPLE": "mini", "REASONING": ["opus", "abc"]} + assert config.tier_model_configs["REASONING"][0].litellm_params == {"reasoning_effort": "xhigh"} + rebuilt = ComplexityRouterConfig.model_validate(config.model_dump()) + assert rebuilt.tier_model_configs["REASONING"][0].litellm_params == {"reasoning_effort": "xhigh"} + + +def test_tier_model_params_accept_a_single_object(): + config = ComplexityRouterConfig( + tiers={"REASONING": {"model_name": "opus", "litellm_params": {"thinking": {"type": "enabled"}}}} + ) + + assert config.tiers == {"REASONING": "opus"} + assert config.tier_model_configs["REASONING"][0].model_name == "opus" + + +@pytest.mark.parametrize( + "tiers", + [ + {"REASONING": [{"litellm_params": {"reasoning_effort": "xhigh"}}]}, + ], +) +def test_tier_model_params_reject_malformed_entries(tiers): + with pytest.raises(ValidationError): + ComplexityRouterConfig(tiers=tiers) + + +def test_tier_model_params_reject_duplicate_models(): + with pytest.raises(ValidationError, match="duplicate model_name"): + ComplexityRouterConfig( + tiers={ + "REASONING": [ + {"model_name": "opus", "litellm_params": {"reasoning_effort": "xhigh"}}, + {"model_name": "opus", "litellm_params": {"reasoning_effort": "low"}}, + ] + } + ) + + +def test_non_adaptive_empty_tier_pool_remains_valid(): + config = ComplexityRouterConfig(tiers={"SIMPLE": []}) + assert config.tiers == {"SIMPLE": []} + + +def test_adaptive_empty_tier_pool_is_rejected(): + with pytest.raises(ValidationError, match="adaptive=True"): + ComplexityRouterConfig(adaptive=True, tiers={"SIMPLE": []}) + + +def test_tier_model_params_are_used_by_pools_and_savings_baseline(mock_router_instance): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": { + "SIMPLE": "mini", + "REASONING": [{"model_name": "opus", "litellm_params": {"reasoning_effort": "xhigh"}}, "abc"], + } + }, + ) + + 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"} + + +@pytest.mark.asyncio +async def test_tier_model_params_reach_the_hook_response_and_override_client_values(mock_router_instance): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": { + "REASONING": { + "model_name": "opus", + "litellm_params": {"reasoning_effort": "xhigh", "max_tokens": 512}, + } + }, + "keyword_tier_rules": [{"keywords": ["reason carefully"], "tier": "REASONING"}], + }, + ) + request_kwargs = {"reasoning_effort": "low", "metadata": {}} + + response = await router.async_pre_routing_hook( + model="test-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "reason carefully about this"}], + ) + + assert response is not None + assert response.litellm_params == {"reasoning_effort": "xhigh", "max_tokens": 512} + assert response.routing_decision is not None + assert response.routing_decision["tier_litellm_params"] == response.litellm_params + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["classification", "keyword", "session"]) +async def test_tier_params_mask_credentials_in_routing_decision(route, mock_router_instance): + params = {"reasoning_effort": "xhigh", "api_key": "secret-tier-key"} + config = { + "tiers": { + tier.value: {"model_name": "opus", "litellm_params": params} + for tier in ComplexityTier + }, + "keyword_tier_rules": [{"keywords": ["reason carefully"], "tier": "REASONING"}] + if route == "keyword" + else None, + "session_affinity": route == "session", + } + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=config, + ) + request_kwargs = {"metadata": {"session_id": "masked-params-session"}} + if route == "session": + mock_router_instance.cache = DualCache() + await mock_router_instance.cache.async_set_cache( + key=router._get_session_affinity_cache_key("masked-params-session", request_kwargs), + value={"model": "opus", "tier": "REASONING"}, + ) + message = "reason carefully about this" if route == "keyword" else "hello" + + response = await router.async_pre_routing_hook( + model="test-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": message}], + ) + + assert response is not None + assert response.litellm_params == params + assert response.routing_decision is not None + assert response.routing_decision["tier_litellm_params"] == { + "reasoning_effort": "xhigh", + "api_key": "secr*******-key", + } + + +@pytest.mark.asyncio +async def test_session_pin_outside_tiers_does_not_inherit_medium_params(mock_router_instance): + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": { + "SIMPLE": "mini", + "MEDIUM": {"model_name": "medium", "litellm_params": {"reasoning_effort": "low"}}, + }, + "session_affinity": True, + "default_model": "orphan", + }, + ) + request_kwargs = {"metadata": {"session_id": "orphan-session"}} + await mock_router_instance.cache.async_set_cache( + key=router._get_session_affinity_cache_key("orphan-session", request_kwargs), + value="orphan", + ) + + response = await router.async_pre_routing_hook( + model="test-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hello"}], + ) + + assert response is not None + assert response.model == "orphan" + assert response.litellm_params == {} + + +@pytest.mark.asyncio +async def test_session_pin_uses_recorded_tier_when_model_is_in_multiple_tiers(mock_router_instance): + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": { + "SIMPLE": {"model_name": "shared", "litellm_params": {"reasoning_effort": "low"}}, + "REASONING": {"model_name": "shared", "litellm_params": {"reasoning_effort": "xhigh"}}, + }, + "session_affinity": True, + }, + ) + request_kwargs = {"metadata": {"session_id": "shared-session"}} + await mock_router_instance.cache.async_set_cache( + key=router._get_session_affinity_cache_key("shared-session", request_kwargs), + value={"model": "shared", "tier": "SIMPLE"}, + ) + + response = await router.async_pre_routing_hook( + model="test-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hello"}], + ) + + assert response is not None + assert response.litellm_params == {"reasoning_effort": "low"} + assert response.routing_decision is not None + assert response.routing_decision["tier"] == "SIMPLE" + + +@pytest.mark.asyncio +async def test_session_pin_survives_json_list_round_trip(mock_router_instance): + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value=["shared", "SIMPLE"]) + mock_router_instance.cache = cache + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": { + "SIMPLE": {"model_name": "shared", "litellm_params": {"reasoning_effort": "low"}}, + "REASONING": {"model_name": "shared", "litellm_params": {"reasoning_effort": "xhigh"}}, + }, + "session_affinity": True, + }, + ) + request_kwargs = {"metadata": {"session_id": "json-round-trip-session"}} + + response = await router.async_pre_routing_hook( + model="test-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hello"}], + ) + + assert response is not None + assert response.model == "shared" + assert response.litellm_params == {"reasoning_effort": "low"} + assert cache.async_set_cache.call_args.kwargs["value"] == {"model": "shared", "tier": "SIMPLE"} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0dc67d567bb..fed15a2ba41 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23800,6 +23800,15 @@ export interface components { * @enum {string} */ ComplexityTier: "SIMPLE" | "MEDIUM" | "COMPLEX" | "REASONING"; + /** ComplexityTierModel */ + ComplexityTierModel: { + /** Litellm Params */ + litellm_params?: { + [key: string]: unknown; + }; + /** Model Name */ + model_name: string; + }; /** * ComplianceCheckRequest * @description Request payload for compliance check endpoints. @@ -32556,6 +32565,10 @@ export interface components { tier_labels?: { [key: string]: string; }; + /** Tier Model Configs */ + tier_model_configs?: { + [key: string]: components["schemas"]["ComplexityTierModel"][]; + }; /** * Tiers * @description Mapping of complexity tiers to a model or model pool. A list is randomly picked from when adaptive=False, and used as a soft-floor home pool when adaptive=True @@ -33540,6 +33553,10 @@ export interface components { tier_boundaries?: components["schemas"]["StandardLoggingRoutingDecisionTierBoundaries"]; /** Tier Label */ tier_label?: string; + /** Tier Litellm Params */ + tier_litellm_params?: { + [key: string]: unknown; + }; }; /** * StandardLoggingRoutingDecisionTierBoundaries