mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
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:
parent
5dcc3a1bbf
commit
a1bd5feebe
5 changed files with 192 additions and 27 deletions
|
|
@ -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 "
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue