mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
refactor(complexity_router): resolve a tier once, for both config load and routing
Three review rounds landed three P1s on the same validator, each one a boolean
added to close one more way it disagreed with request-time resolution: first that
an unservable default_tier was accepted at all, then that the plugin path never
consults default_model, now that a tier present with an empty pool is not the same
as a tier absent. They are one defect. `_validate_default_tier_is_servable` was a
second, hand-maintained model of what `get_model_for_tier` does, so it drifted from
it one case at a time.
`ComplexityRouterConfig.resolve_tier` is now the single answer to "what may serve
this tier", returning models or the reason there are none. The validator asks it
and so does selection, so there is no precedence left to re-derive and nothing to
drift; the validator is four lines with no conditions of its own.
Resolution keys off a tier's models rather than its key being present, which is the
third finding fixed where it lives rather than mirrored into config: `{MEDIUM: []}`
and a tiers map with no MEDIUM both say the tier has no models, so both fall
through to default_model. Mirroring the old asymmetry into the validator instead
would have made it part of the config contract.
The deployment-level complexity_router_default_model now goes in before validation
rather than being assigned onto the validated model afterwards. router.py always
derives one from the MEDIUM-then-SIMPLE tier, so an explicit default_tier outside
`tiers` is servable on every proxy deployment; validating before it was applied
failed those configs at startup for a gap routing did not have.
Drops _pick_from_tier_value and the hand-rolled plugin-path raise, both subsumed.
Tests pin the invariant rather than the instances: a config the validator accepts
is one the default tier can be served from, across the tier's own models,
default_model at either level, and an empty pool falling through.
This commit is contained in:
parent
be7b23ce0f
commit
0d194b881e
4 changed files with 183 additions and 95 deletions
|
|
@ -142,7 +142,9 @@ complexity_router_config:
|
|||
default_tier: MEDIUM # SIMPLE | MEDIUM | COMPLEX | REASONING
|
||||
```
|
||||
|
||||
Set `default_tier: SIMPLE` to keep unmatched traffic on the cheapest tier, which is how the router behaved before this setting existed. A `default_tier` you set explicitly has to have a model behind it, either its own non-empty entry in `tiers` or a `default_model`; the config is rejected at load time otherwise, rather than failing on the first unmatched request.
|
||||
Set `default_tier: SIMPLE` to keep unmatched traffic on the cheapest tier, which is how the router behaved before this setting existed.
|
||||
|
||||
`default_tier` names which tier no-signal traffic is classified as; it is resolved to a model by the same chain as any classified tier, so it can be served by its own entry in `tiers`, by `default_model`, or by the MEDIUM tier. A `default_tier` you set explicitly has to be servable by one of those, and the config is rejected at load time when it is not, rather than failing on the first unmatched request. With routing plugins configured that chain stops at the tier's own models, since a model the plugins never vetted must not serve, so the tier itself has to name one.
|
||||
|
||||
The check is on the individual dimensions, not on the weighted score, because contributions cancel. `"hi, quick python question"` scores zero with three dimensions firing (short prompt, a simple indicator, a code keyword); it has real evidence of being simple and stays SIMPLE. Prompts shorter than the `simple` token threshold or longer than the `complex` one also fire `tokenCount`, so they score normally and are outside this path.
|
||||
|
||||
|
|
|
|||
|
|
@ -46,6 +46,8 @@ from .config import (
|
|||
TIER_SEVERITY_ORDER,
|
||||
ComplexityRouterConfig,
|
||||
ComplexityTier,
|
||||
TierModels,
|
||||
TierUnservable,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -87,6 +89,24 @@ def _append_custom_keywords(base_keywords: list[str], custom_keywords: list[str]
|
|||
return [*base_keywords, *deduped_custom.values()]
|
||||
|
||||
|
||||
def _servable_models(config: ComplexityRouterConfig, tier: ComplexityTier) -> tuple[str, ...]:
|
||||
"""The models that may serve this tier, raising the resolver's own reason when none may.
|
||||
|
||||
The one place a `TierUnservable` becomes an exception, so every caller reports the same
|
||||
gap the same way and none of them re-derives when a tier is servable.
|
||||
"""
|
||||
match config.resolve_tier(tier):
|
||||
case TierModels(models=models):
|
||||
return models
|
||||
case TierUnservable() as unservable:
|
||||
raise ValueError(unservable.describe())
|
||||
|
||||
|
||||
def _pick_model(models: tuple[str, ...]) -> str:
|
||||
"""One model out of a resolved, non-empty pool; a single pin never consults the RNG."""
|
||||
return models[0] if len(models) == 1 else random.choice(models)
|
||||
|
||||
|
||||
# Metadata keys that carry only the parent request's budget reservation state. These
|
||||
# must not reach internal sub-calls (classifier, embedding): the reservation belongs to
|
||||
# the routed completion being decided on, not to the sub-call itself, and forwarding it
|
||||
|
|
@ -328,15 +348,15 @@ class ComplexityRouter(CustomLogger):
|
|||
self.model_name = model_name
|
||||
self.litellm_router_instance = litellm_router_instance
|
||||
|
||||
# Parse config - always create a new instance to avoid singleton mutation
|
||||
if complexity_router_config:
|
||||
self.config = ComplexityRouterConfig(**complexity_router_config)
|
||||
else:
|
||||
self.config = ComplexityRouterConfig()
|
||||
|
||||
# Override default_model if provided
|
||||
# Parse config - always create a new instance to avoid singleton mutation.
|
||||
# The deployment-level default_model goes in before validation, not onto the
|
||||
# validated model afterwards, so the config that gets checked is the config that
|
||||
# routes; assigning it later left validation judging a default_model that requests
|
||||
# would never see.
|
||||
config_fields: dict[str, Any] = dict(complexity_router_config or {})
|
||||
if default_model:
|
||||
self.config.default_model = default_model
|
||||
config_fields["default_model"] = default_model
|
||||
self.config = ComplexityRouterConfig(**config_fields)
|
||||
|
||||
# Build effective keyword lists (use config overrides or defaults)
|
||||
self.code_keywords = self.config.code_keywords or DEFAULT_CODE_KEYWORDS
|
||||
|
|
@ -833,30 +853,10 @@ class ComplexityRouter(CustomLogger):
|
|||
Returns:
|
||||
The model name configured for that tier.
|
||||
"""
|
||||
tier_key = tier.value if isinstance(tier, ComplexityTier) else tier
|
||||
return _pick_model(_servable_models(self.config, tier))
|
||||
|
||||
if tier_key in self.config.tiers:
|
||||
return self._pick_from_tier_value(self.config.tiers[tier_key], tier_key)
|
||||
|
||||
if self.config.default_model:
|
||||
return self.config.default_model
|
||||
|
||||
medium_key = ComplexityTier.MEDIUM.value
|
||||
if medium_key in self.config.tiers:
|
||||
return self._pick_from_tier_value(self.config.tiers[medium_key], medium_key)
|
||||
|
||||
raise ValueError(f"No model configured for tier {tier_key} and no default_model set")
|
||||
|
||||
@staticmethod
|
||||
def _pick_from_tier_value(model: str | list[str], tier_key: str) -> str:
|
||||
if isinstance(model, str):
|
||||
return model
|
||||
if not model:
|
||||
raise ValueError(f"Empty model pool for tier {tier_key}")
|
||||
return random.choice(model)
|
||||
|
||||
def _tier_pools(self) -> dict[str, list[str]]:
|
||||
return {tier: (models if isinstance(models, list) else [models]) for tier, models in self.config.tiers.items()}
|
||||
def _tier_pools(self) -> dict[str, tuple[str, ...]]:
|
||||
return {tier: self.config.models_for(tier) for tier in self.config.tiers}
|
||||
|
||||
async def _pick_model_for_tier(
|
||||
self,
|
||||
|
|
@ -871,20 +871,11 @@ class ComplexityRouter(CustomLogger):
|
|||
from litellm.types.router import RoutingContext
|
||||
|
||||
tier_key = tier.value
|
||||
candidates = list(self._tier_pools().get(tier_key, []))
|
||||
if not candidates:
|
||||
# Distinct from the plugin denial below: nothing was filtered out, the tier was
|
||||
# never given models. Saying "after routing-plugin filtering" would send an
|
||||
# operator to read plugin code for what is a gap in `tiers`.
|
||||
raise ValueError(
|
||||
f"Tier {tier_key} has no models configured. Routing plugins are configured, so "
|
||||
f"default_model is not consulted: a model the plugins never vetted must not serve"
|
||||
)
|
||||
metadata_key = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata"
|
||||
context = RoutingContext(
|
||||
raw_messages=raw_messages or [],
|
||||
structured_messages=resolved_messages or [],
|
||||
candidate_models=candidates,
|
||||
candidate_models=list(_servable_models(self.config, tier)),
|
||||
metadata=request_kwargs.get(metadata_key) or {},
|
||||
)
|
||||
for plugin in self.config.plugins:
|
||||
|
|
@ -897,7 +888,7 @@ class ComplexityRouter(CustomLogger):
|
|||
# silently bypassed. Raise instead, matching the Router-level plugin
|
||||
# pipeline's own fail-closed behavior for the same situation.
|
||||
raise ValueError(f"No candidate models left for tier {tier_key} after routing-plugin filtering")
|
||||
return self._pick_from_tier_value(context.candidate_models, tier_key)
|
||||
return _pick_model(tuple(context.candidate_models))
|
||||
|
||||
def _ensure_adaptive_router(self) -> Any | None:
|
||||
if not self.config.adaptive:
|
||||
|
|
|
|||
|
|
@ -5,8 +5,9 @@ Contains default keyword lists, weights, tier boundaries, and configuration clas
|
|||
All values are configurable via proxy config.yaml.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Literal
|
||||
from typing import Literal, assert_never
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
|
||||
|
|
@ -253,6 +254,43 @@ class ClassifierLLMConfig(BaseModel):
|
|||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TierModels:
|
||||
"""The models that may serve a tier."""
|
||||
|
||||
models: tuple[str, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TierUnservable:
|
||||
"""Why no configured model may serve a tier."""
|
||||
|
||||
tier: str
|
||||
configured_tiers: tuple[str, ...]
|
||||
reason: Literal["nothing_configured", "plugins_require_own_pool"]
|
||||
|
||||
def describe(self) -> str:
|
||||
match self.reason:
|
||||
case "plugins_require_own_pool":
|
||||
why = (
|
||||
"routing plugins are configured, so only this tier's own models may serve it: "
|
||||
"a model the plugins never vetted must not serve"
|
||||
)
|
||||
remedy = "give it models in tiers, or name a tier that has them"
|
||||
case "nothing_configured":
|
||||
why = "it has no models, and neither default_model nor the MEDIUM tier supplies one"
|
||||
remedy = "give it models in tiers, or set default_model"
|
||||
case _:
|
||||
assert_never(self.reason)
|
||||
return (
|
||||
f"No model can serve tier {self.tier}: {why}. "
|
||||
f"Configured tiers: {', '.join(self.configured_tiers) or 'none'}. To fix, {remedy}"
|
||||
)
|
||||
|
||||
|
||||
TierResolution = TierModels | TierUnservable
|
||||
|
||||
|
||||
class ComplexityRouterConfig(BaseModel):
|
||||
"""Configuration for the ComplexityRouter."""
|
||||
|
||||
|
|
@ -276,8 +314,8 @@ class ComplexityRouterConfig(BaseModel):
|
|||
description=(
|
||||
"Tier used when no scoring dimension fires, i.e. the prompt matched no keyword, "
|
||||
"pattern or token-count threshold and the scorer has no evidence either way. "
|
||||
"When set explicitly it must have a model behind it: a non-empty entry in `tiers`, "
|
||||
"or `default_model`. "
|
||||
"It resolves to a model the same way a classified tier does, and when set explicitly "
|
||||
"it is rejected at load unless that resolution finds one. "
|
||||
"Set to SIMPLE to restore the previous behavior of treating unmatched traffic as simple"
|
||||
),
|
||||
)
|
||||
|
|
@ -472,33 +510,48 @@ class ComplexityRouterConfig(BaseModel):
|
|||
raise ValueError("classifier_llm_config is required when classifier_type is 'llm'")
|
||||
return self
|
||||
|
||||
def models_for(self, tier: ComplexityTier | str) -> tuple[str, ...]:
|
||||
"""The models this tier itself names; absent, an empty pool and an empty pin all mean none."""
|
||||
configured = self.tiers.get(tier.value if isinstance(tier, ComplexityTier) else tier)
|
||||
if isinstance(configured, str):
|
||||
return (configured,) if configured else ()
|
||||
return tuple(configured or ())
|
||||
|
||||
def resolve_tier(self, tier: ComplexityTier | str) -> TierResolution:
|
||||
"""Which models may serve this tier, or why none may.
|
||||
|
||||
Config validation and request-time selection both ask here, so a config that loads
|
||||
is a config that routes. A second hand-kept copy of this precedence drifts from it
|
||||
one case at a time, until an accepted config raises on its first request.
|
||||
"""
|
||||
tier_key = tier.value if isinstance(tier, ComplexityTier) else tier
|
||||
configured = tuple(sorted(self.tiers))
|
||||
own = self.models_for(tier_key)
|
||||
if own:
|
||||
return TierModels(own)
|
||||
if self.plugins:
|
||||
return TierUnservable(tier_key, configured, "plugins_require_own_pool")
|
||||
if self.default_model:
|
||||
return TierModels((self.default_model,))
|
||||
medium = self.models_for(ComplexityTier.MEDIUM)
|
||||
if medium:
|
||||
return TierModels(medium)
|
||||
return TierUnservable(tier_key, configured, "nothing_configured")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_default_tier_is_servable(self) -> "ComplexityRouterConfig":
|
||||
"""Reject an explicit `default_tier` that nothing can serve.
|
||||
|
||||
Left implicit it is not checked, so a partial `tiers` map keeps loading and MEDIUM
|
||||
resolves through the same chain every other tier does.
|
||||
"""
|
||||
if "default_tier" not in self.model_fields_set:
|
||||
return self
|
||||
if self.tiers.get(self.default_tier.value):
|
||||
return self
|
||||
# default_model rescues this only without plugins. The plugin path never falls back
|
||||
# to it, since a model the plugins did not vet must not serve, so accepting it here
|
||||
# would validate a config whose every no-signal request fails at routing time.
|
||||
if self.default_model and not self.plugins:
|
||||
return self
|
||||
remedy = (
|
||||
"Add it to tiers, or name a tier that is configured"
|
||||
if self.plugins
|
||||
else "Add it to tiers, name a tier that is configured, or set default_model in complexity_router_config"
|
||||
)
|
||||
because = (
|
||||
"routing plugins are configured, so default_model is not consulted: a model the plugins "
|
||||
"never vetted must not serve"
|
||||
if self.plugins
|
||||
else "the deployment-level complexity_router_default_model does not count here, because "
|
||||
"falling through to it would serve every no-signal request from a model this tier never names"
|
||||
)
|
||||
raise ValueError(
|
||||
f"default_tier {self.default_tier.value} is not a non-empty entry in tiers "
|
||||
f"({sorted(self.tiers)}). {remedy}; {because}"
|
||||
)
|
||||
match self.resolve_tier(self.default_tier):
|
||||
case TierModels():
|
||||
return self
|
||||
case TierUnservable() as unservable:
|
||||
raise ValueError(f"default_tier {self.default_tier.value} is unservable. {unservable.describe()}")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_adaptive_pools(self) -> "ComplexityRouterConfig":
|
||||
|
|
|
|||
|
|
@ -347,7 +347,7 @@ class TestModelSelection:
|
|||
"default_model": "mid",
|
||||
},
|
||||
)
|
||||
pool = ["cheap", "premium"]
|
||||
pool = ("cheap", "premium")
|
||||
with patch(
|
||||
"litellm.router_strategy.complexity_router.complexity_router.random.choice",
|
||||
return_value="premium",
|
||||
|
|
@ -356,16 +356,30 @@ class TestModelSelection:
|
|||
choice.assert_called_once_with(pool)
|
||||
assert router.get_model_for_tier(ComplexityTier.MEDIUM) == "mid"
|
||||
|
||||
def test_get_model_for_tier_empty_pool_raises(self, mock_router_instance):
|
||||
@pytest.mark.parametrize("no_models", [[], ""], ids=["empty_pool", "empty_pin"])
|
||||
def test_a_tier_with_no_models_falls_through_like_an_absent_one(self, mock_router_instance, no_models):
|
||||
"""`{"SIMPLE": []}` and a `tiers` map with no SIMPLE key say the same thing, so they
|
||||
have to resolve the same way. Selecting on the key being present sent the empty one
|
||||
into a pool it then refused to pick from, raising where the absent one served
|
||||
default_model."""
|
||||
absent, empty = (
|
||||
ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config={"tiers": tiers, "default_model": "mid"},
|
||||
)
|
||||
for tiers in ({}, {"SIMPLE": no_models})
|
||||
)
|
||||
assert empty.get_model_for_tier(ComplexityTier.SIMPLE) == absent.get_model_for_tier(ComplexityTier.SIMPLE)
|
||||
assert empty.get_model_for_tier(ComplexityTier.SIMPLE) == "mid"
|
||||
|
||||
def test_a_tier_with_no_models_and_nothing_to_fall_through_to_raises(self, mock_router_instance):
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config={
|
||||
"tiers": {"SIMPLE": []},
|
||||
"default_model": "mid",
|
||||
},
|
||||
complexity_router_config={"tiers": {"SIMPLE": []}},
|
||||
)
|
||||
with pytest.raises(ValueError, match="Empty model pool for tier SIMPLE"):
|
||||
with pytest.raises(ValueError, match="No model can serve tier SIMPLE"):
|
||||
router.get_model_for_tier(ComplexityTier.SIMPLE)
|
||||
|
||||
|
||||
|
|
@ -1911,7 +1925,7 @@ class TestAdaptiveSoftFloors:
|
|||
"default_model": "mid",
|
||||
},
|
||||
)
|
||||
pool = ["cheap", "premium"]
|
||||
pool = ("cheap", "premium")
|
||||
with patch(
|
||||
"litellm.router_strategy.complexity_router.complexity_router.random.choice",
|
||||
return_value="premium",
|
||||
|
|
@ -4869,29 +4883,57 @@ class TestNoSignalDefaultTier:
|
|||
with pytest.raises(ValidationError):
|
||||
_router_with_default_tier(mock_router_instance, "CHEAPEST")
|
||||
|
||||
def test_default_tier_without_a_model_is_rejected_at_config_time(self, mock_router_instance):
|
||||
"""An explicit default_tier naming a tier with no model would only surface as a routing
|
||||
def test_default_tier_with_no_model_anywhere_is_rejected_at_config_time(self, mock_router_instance):
|
||||
"""An explicit default_tier that nothing can serve would only surface as a routing
|
||||
failure on the first no-signal request, so reject the config instead."""
|
||||
with pytest.raises(ValidationError, match="default_tier COMPLEX is not a non-empty entry in tiers"):
|
||||
with pytest.raises(ValidationError, match="default_tier COMPLEX is unservable"):
|
||||
ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config={
|
||||
"tiers": {"SIMPLE": "simple-model", "MEDIUM": "medium-model"},
|
||||
"default_tier": "COMPLEX",
|
||||
},
|
||||
complexity_router_config={"tiers": {"COMPLEX": []}, "default_tier": "COMPLEX"},
|
||||
)
|
||||
|
||||
def test_default_tier_with_an_empty_pool_is_rejected(self, mock_router_instance):
|
||||
with pytest.raises(ValidationError, match="default_tier MEDIUM is not a non-empty entry in tiers"):
|
||||
ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config={
|
||||
"tiers": {"SIMPLE": "simple-model", "MEDIUM": []},
|
||||
"default_tier": "MEDIUM",
|
||||
},
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"config_overrides, deployment_default_model, expected",
|
||||
[
|
||||
({"tiers": {"SIMPLE": "s", "MEDIUM": "m"}}, None, "m"),
|
||||
({"tiers": {"SIMPLE": "s", "MEDIUM": "m"}, "default_model": "f"}, None, "f"),
|
||||
({"tiers": {"SIMPLE": "s", "MEDIUM": "m"}}, "d", "d"),
|
||||
({"tiers": {"SIMPLE": "s", "COMPLEX": []}, "default_model": "f"}, None, "f"),
|
||||
],
|
||||
ids=["medium_tier", "config_default_model", "deployment_default_model", "empty_pool_falls_through"],
|
||||
)
|
||||
def test_a_default_tier_the_config_accepts_is_one_a_request_can_land_on(
|
||||
self, mock_router_instance, config_overrides, deployment_default_model, expected
|
||||
):
|
||||
"""The invariant the validator exists to hold: loading and routing agree.
|
||||
|
||||
default_tier is not a special tier; it names which tier no-signal traffic is
|
||||
classified as, and that tier resolves through the same chain as any classified
|
||||
one. Every accepted config here therefore has to serve, and a validator that
|
||||
re-derives that chain by hand instead of asking `resolve_tier` drifts from it
|
||||
one case at a time, accepting configs whose first no-signal request raises."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config={"default_tier": "COMPLEX", **config_overrides},
|
||||
default_model=deployment_default_model,
|
||||
)
|
||||
assert router.get_model_for_tier(ComplexityTier.COMPLEX) == expected
|
||||
|
||||
def test_the_deployment_level_default_model_is_visible_to_validation(self, mock_router_instance):
|
||||
"""router.py always hands ComplexityRouter a derived complexity_router_default_model,
|
||||
so a default_tier outside `tiers` is servable on every proxy deployment. Validating
|
||||
before that value was applied failed those configs at startup for a gap that routing
|
||||
did not have."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config={"tiers": {"SIMPLE": "s"}, "default_tier": "COMPLEX"},
|
||||
default_model="derived-from-tiers",
|
||||
)
|
||||
assert router.config.default_model == "derived-from-tiers"
|
||||
assert router.get_model_for_tier(ComplexityTier.COMPLEX) == "derived-from-tiers"
|
||||
|
||||
def test_default_model_does_not_rescue_the_default_tier_when_plugins_are_configured(self, mock_router_instance):
|
||||
"""The plugin path builds candidates from the tier pool alone and never consults
|
||||
|
|
@ -4934,7 +4976,7 @@ class TestNoSignalDefaultTier:
|
|||
"session_affinity": False,
|
||||
},
|
||||
)
|
||||
with pytest.raises(ValueError, match="Tier MEDIUM has no models configured"):
|
||||
with pytest.raises(ValueError, match="No model can serve tier MEDIUM: routing plugins are configured"):
|
||||
await router.async_pre_routing_hook(
|
||||
model="test-model",
|
||||
request_kwargs={},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue