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 <krrishdholakia@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-08-20 01:39:02 +00:00 committed by GitHub
parent 7b6e16cfd3
commit f5cfa84220
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 480 additions and 15 deletions

View file

@ -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(

View file

@ -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:

View file

@ -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,
),
)

View file

@ -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:

View file

@ -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)

View file

@ -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",
}
)

View file

@ -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"}

View file

@ -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