feat(router): infer custom classifier tiers from model mappings (#48987)

* feat(router): infer custom classifier tiers from model mappings

* fix(router): preserve custom tier floors on classifier fallback

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
moe-berri 2026-10-29 11:30:23 -07:00 • committed by GitHub
parent 5dcc3a1bbf
commit a1bd5feebe
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 192 additions and 27 deletions

View file

@ -19,6 +19,7 @@ from litellm.router_strategy.complexity_router.config import (
COMPLEXITY_ROUTER_CONFIG_KEYS,
DEFAULT_JEV_INSTRUCTIONS,
LLM_CLASSIFIER_TYPES,
ComplexityTier,
configured_local_heuristic,
normalize_classifier_config_aliases,
)
@ -195,12 +196,12 @@ def uses_heuristic_v2_classifier(complexity_router_config: object) -> bool:
def defines_custom_tiers(complexity_router_config: object) -> bool:
"""Whether this complexity config replaces the built-in tier ladder with operator-defined tier_definitions.
Mirrors the SQL spelling on the capability record: only an actual array claims the capability,
so an explicit JSON null or a malformed value does not.
"""
return isinstance(_mapping(complexity_router_config).get("tier_definitions"), (list, tuple))
"""Whether this router uses explicitly defined or custom Python classifier tiers."""
config: Final = _mapping(complexity_router_config)
return isinstance(config.get("tier_definitions"), (list, tuple)) or (
config.get("classifier_type") == "custom"
and any(name not in ComplexityTier.__members__ for name in _mapping(config.get("tiers")))
)
OPERATOR_CLASSIFIER_PROMPT_FIELDS: Final = ("classification_prompt", "classification_examples")
@ -289,6 +290,8 @@ _OPENSOURCE_CLASSIFIER_CONFIG_SQL: Final = (
"COALESCE({config} -> 'opensource_classifier_config', {config} -> 'jev_classifier_config')"
)
_BUILT_IN_TIER_KEYS_SQL: Final = ", ".join(f"'{tier.value}'" for tier in ComplexityTier)
CUSTOMIZATION_CAPABILITY: Final = GatedAutoRouterCapability(
key="tier_or_classifier_prompt",
subject="with operator-defined tier_definitions or an operator-written classifier prompt",
@ -299,6 +302,9 @@ CUSTOMIZATION_CAPABILITY: Final = GatedAutoRouterCapability(
uses=uses_custom_tier_or_classifier_prompt,
sql_config_predicate=(
"jsonb_typeof({config} -> 'tier_definitions') = 'array' OR "
"({config} ->> 'classifier_type' = 'custom' AND "
"jsonb_typeof({config} -> 'tiers') = 'object' AND "
f"({{config}} -> 'tiers') - ARRAY[{_BUILT_IN_TIER_KEYS_SQL}]::text[] <> jsonb_build_object()) OR "
f"({{config}} ->> 'classifier_type' IN ({_LLM_CLASSIFIER_TYPES_SQL}) AND ("
"{config} -> 'classifier_llm_config' ->> 'system_prompt' IS NOT NULL OR "
f"{_OPERATOR_PROMPT_FIELDS_SQL})) OR "

View file

@ -2607,9 +2607,10 @@ class ComplexityRouter(CustomLogger):
"""The classifier-failed outcome for classifier_fallback='default_model'.
The outcome still carries a tier because ClassificationOutcome requires one, so it reports
the tier whose pool holds default_model, and MEDIUM when no pool does. Nothing about the
request produced that tier, so the pre-routing hook never logs it as the request's tier: it
routes this cause straight to default_model rather than picking from the tier's pool, since
the tier whose pool holds default_model, or the first custom tier (MEDIUM for built-in
tiers) when no pool does. Nothing about the request produced that tier, so the
pre-routing hook never logs it as the request's tier: it routes this cause straight to
default_model rather than picking from the tier's pool, since
a pool with several models would otherwise land somewhere else and the point of this
fallback is a known destination when classification failed.
@ -2622,14 +2623,11 @@ class ComplexityRouter(CustomLogger):
"""
default_model: Final = self.config.default_model
pools: Final = self._tier_pools()
tier: Final = next(
(
candidate
for candidate in self.config.active_tier_severity_order()
if default_model in pools.get(candidate.value, ())
),
ComplexityTier.MEDIUM,
name: Final = next(
(candidate for candidate in self.config.tier_names() if default_model in pools.get(candidate, ())),
self.config.tier_names()[0] if self.config.has_custom_tiers else ComplexityTier.MEDIUM.value,
)
tier: Final = name if self.config.has_custom_tiers else ComplexityTier(name)
return ClassificationOutcome(
tier=tier, score=None, signals=("classifier-failed:default-model",), cause="default_model_fallback"
)
@ -3415,8 +3413,8 @@ class ComplexityRouter(CustomLogger):
def _active_tier_severity(self, tier: ComplexityTier | str) -> int:
"""Position of a tier in the active severity order: TIER_SEVERITY_ORDER for the built-in
set, tier_definitions list order (ascending) for a custom set -- the same order
keyword_tier_rules resolve severity against."""
set, tier_definitions list order or inferred mapping order (ascending) for a custom set --
the same order keyword_tier_rules resolve severity against."""
return self.config.tier_names().index(_tier_name(tier))
def _matched_plan_mode_signal(
@ -4961,7 +4959,8 @@ class ComplexityRouter(CustomLogger):
fallback_model: Final = self.config.default_model if not self.config.plugins else None
# A sentinel-carrying request skips the failure exit below, whether or not the floor
# moved the tier: default_model carries no tier guarantee (its placeholder tier is the
# pool that holds it, or MEDIUM when none does), so a placeholder at or above the floor
# pool that holds it, or the first custom tier / MEDIUM when none does), so a placeholder
# at or above the floor
# would otherwise route a plan-mode request to a model the floor cannot vouch for. The
# clamped tier's pool is the destination the floor can guarantee.
if outcome.cause == "default_model_fallback" and fallback_model is not None and plan_mode_sentinel is None:

View file

@ -2040,8 +2040,10 @@ class ComplexityRouterConfig(LiteLLMBaseModel):
@property
def has_custom_tiers(self) -> bool:
"""True when the operator replaced the built-in tier set via tier_definitions."""
return self.tier_definitions is not None
"""Whether routing uses names outside the built-in complexity ladder."""
return self.tier_definitions is not None or (
self.classifier_type == "custom" and any(name not in ComplexityTier.__members__ for name in self.tiers)
)
@property
def uses_llm_classifier(self) -> bool:
@ -2059,6 +2061,8 @@ class ComplexityRouterConfig(LiteLLMBaseModel):
"""The active tier names: the defined names, or the built-in set in severity order."""
if self.tier_definitions is not None:
return tuple(definition.name for definition in self.tier_definitions)
if self.has_custom_tiers:
return tuple(self.tiers)
return tuple(tier.value for tier in self.active_tier_severity_order())
def classifier_wire_labels(self) -> tuple[str, ...]:
@ -2069,7 +2073,7 @@ class ComplexityRouterConfig(LiteLLMBaseModel):
def resolve_classified_tier(self, label: str) -> ComplexityTier | str | None:
"""Resolve a classifier reply to the active tier it names, or None when it names none."""
if self.tier_definitions is None:
if not self.has_custom_tiers:
return self.tier_for_label(label)
folded: Final = label.strip().casefold()
return next((name for name in self.tier_names() if name.casefold() == folded), None)
@ -2138,7 +2142,7 @@ class ComplexityRouterConfig(LiteLLMBaseModel):
"set replaces",
),
(
self.classifier_fallback == "default_model",
self.classifier_fallback == "default_model" and self.classifier_type != "custom",
"classifier_fallback 'default_model' cannot be combined with tier_definitions: fallback_tier "
"is where a custom-tier router routes when the classifier fails",
),
@ -2183,17 +2187,19 @@ class ComplexityRouterConfig(LiteLLMBaseModel):
@model_validator(mode="after")
def _validate_tier_definitions(self) -> "ComplexityRouterConfig":
if self.tier_definitions is None:
if self.tier_definitions is None and not self.has_custom_tiers:
if self.fallback_tier is not None:
raise ValueError("fallback_tier requires tier_definitions")
for message in self._built_in_opening_conflicts():
raise ValueError(message)
return self
names: Final = tuple(definition.name for definition in self.tier_definitions)
if not 2 <= len(names) <= MAX_TIER_DEFINITIONS:
names: Final = self.tier_names()
if self.tier_definitions is not None and not 2 <= len(names) <= MAX_TIER_DEFINITIONS:
raise ValueError(
f"tier_definitions must define between 2 and {MAX_TIER_DEFINITIONS} tiers, got {len(names)}"
)
if any(not name.strip() or name != name.strip() or "\n" in name or "\r" in name for name in names):
raise ValueError("tier names must be non-empty and contain no surrounding whitespace or newlines")
folded: Final = tuple(name.casefold() for name in names)
duplicated: Final = tuple(
sorted(frozenset(name for name, fold in zip(names, folded) if folded.count(fold) > 1))
@ -2221,9 +2227,11 @@ class ComplexityRouterConfig(LiteLLMBaseModel):
f"tiers must map every defined tier to at least one model; empty: {', '.join(empty_pools)}"
)
if self.fallback_tier is None:
if self.classifier_type == "custom" and self.classifier_fallback == "default_model":
return self
raise ValueError(
"fallback_tier is required with tier_definitions: it is where requests route when the "
"LLM classifier fails"
"LLM classifier fails. Custom classifiers can instead use classifier_fallback='default_model'"
)
if self.fallback_tier not in defined:
raise ValueError(

View file

@ -772,3 +772,21 @@ def test_deployed_decision_classifier_is_a_model_access_dependency() -> None:
("cheap", "tier"),
("decision-deployment", "classifier"),
)
@pytest.mark.parametrize(
"classifier, tiers, expected",
[
("custom", {"quick": "fast", "careful": "capable"}, True),
("custom", {"SIMPLE": "fast", "SECURITY": "capable"}, True),
("custom", {"SIMPLE": "fast", "REASONING": "capable"}, False),
("heuristic", {"quick": "fast"}, False),
],
)
def test_inferred_custom_tiers_claim_existing_gateway_capability(
classifier: str,
tiers: dict[str, str],
expected: bool,
) -> None:
capabilities: Final = claimed_capabilities({"classifier_type": classifier, "tiers": tiers})
assert any(capability.key == "tier_or_classifier_prompt" for capability in capabilities) is expected

View file

@ -17937,3 +17937,137 @@ class TestNonReasoningTier:
"complex",
"reasoning",
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"verdict, expected", [("quick", "fast"), (" CAREFUL ", "capable"), (None, "backup"), ("unknown", "backup")]
)
async def test_custom_classifier_infers_tiers_and_routes_fallback(verdict: str | None, expected: str) -> None:
router: Final = ComplexityRouter(
model_name="auto",
litellm_router_instance=litellm.Router(model_list=[]),
default_model="backup",
complexity_router_config={
"classifier_type": "custom",
"classifier_plugin": _FixedTierClassifier(verdict),
"classifier_fallback": "default_model",
"tiers": {"quick": "fast", "careful": "capable"},
},
)
result: Final = await router.async_pre_routing_hook(
model="auto", request_kwargs={}, messages=[{"role": "user", "content": "hello"}]
)
assert result is not None
assert result.model == expected
@pytest.mark.asyncio
async def test_custom_classifier_can_extend_builtin_tiers() -> None:
router: Final = ComplexityRouter(
model_name="auto",
litellm_router_instance=litellm.Router(model_list=[]),
default_model="backup",
complexity_router_config={
"classifier_type": "custom",
"classifier_plugin": _FixedTierClassifier("SECURITY"),
"classifier_fallback": "default_model",
"tiers": {"SIMPLE": "fast", "REASONING": "capable", "SECURITY": "security-model"},
},
)
result: Final = await router.async_pre_routing_hook(
model="auto", request_kwargs={}, messages=[{"role": "user", "content": "LITELLM ESCALATE"}]
)
assert result is not None
assert result.model == "security-model"
@pytest.mark.parametrize(
"tiers", [{"": "fast"}, {" quick": "fast"}, {"quick": "fast", "QUICK": "capable"}, {"quick": []}]
)
def test_inferred_custom_tiers_reject_unusable_routes(tiers: dict[str, str | list[str]]) -> None:
with pytest.raises(ValidationError):
ComplexityRouterConfig(
classifier_type="custom",
classifier_plugin=_FixedTierClassifier("quick"),
classifier_fallback="default_model",
tiers=tiers,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("explicit_fallback", [None, "quick"])
async def test_defined_custom_tiers_allow_default_model_fallback(explicit_fallback: str | None) -> None:
router: Final = ComplexityRouter(
model_name="auto",
litellm_router_instance=litellm.Router(model_list=[]),
default_model="backup",
complexity_router_config={
"classifier_type": "custom",
"classifier_plugin": _RaisingClassifier(),
"classifier_fallback": "default_model",
"tiers": {"quick": "fast", "careful": "capable"},
"tier_definitions": [
{"name": "quick", "description": "Routine work"},
{"name": "careful", "description": "Difficult work"},
],
"fallback_tier": explicit_fallback,
},
)
result: Final = await router.async_pre_routing_hook(
model="auto", request_kwargs={}, messages=[{"role": "user", "content": "hello"}]
)
assert result is not None
assert result.model == ("fast" if explicit_fallback else "backup")
@pytest.mark.parametrize("setting", ["adaptive", "session_affinity", "stall_escalation_enabled"])
def test_inferred_custom_tiers_reject_order_dependent_features(setting: str) -> None:
with pytest.raises(ValidationError, match="severity order"):
ComplexityRouterConfig.model_validate(
{
"classifier_type": "custom",
"classifier_plugin": _FixedTierClassifier("quick"),
"classifier_fallback": "default_model",
"tiers": {"quick": "fast", "careful": "capable"},
setting: True,
}
)
@pytest.mark.asyncio
@pytest.mark.parametrize("defined", [False, True])
@pytest.mark.parametrize("plan_mode", [False, True])
@pytest.mark.parametrize(
"default_model,planned_model",
[("backup", "capable"), ("fast", "capable"), ("capable", "capable"), ("premium-model", "premium-model")],
)
async def test_custom_classifier_failure_preserves_plan_floor_and_direct_default(
defined: bool, plan_mode: bool, default_model: str, planned_model: str
) -> None:
definitions: Final = (
[{"name": name, "description": name} for name in ("quick", "careful", "premium")] if defined else None
)
router: Final = ComplexityRouter(
model_name="auto",
litellm_router_instance=litellm.Router(model_list=[]),
default_model=default_model,
complexity_router_config={
"classifier_type": "custom",
"classifier_plugin": _RaisingClassifier(),
"classifier_fallback": "default_model",
"tiers": {"quick": "fast", "careful": "capable", "premium": "premium-model"},
"tier_definitions": definitions,
"plan_mode_min_tier": "careful",
},
)
prompt: Final = "<system-reminder>Plan mode is active</system-reminder> hello" if plan_mode else "hello"
result: Final = await router.async_pre_routing_hook(
model="auto", request_kwargs={}, messages=[{"role": "user", "content": prompt}]
)
assert result is not None
assert result.model == (planned_model if plan_mode else default_model)
assert result.routing_decision is not None
assert result.routing_decision.get("tier") == (
("premium" if planned_model == "premium-model" else "careful") if plan_mode else None
)