mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
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:
parent
7b6e16cfd3
commit
f5cfa84220
8 changed files with 480 additions and 15 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
17
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
17
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue