scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

feat(router): calibrate capability classifier verdicts

Browse source 
Adapt the capability classifier to forecast a defined SUCCESS event with
reasoning ordered before the probability, add a capability_boundary bucket
that steps the qualification threshold per candidate, and cap classifier
payload values so large tool outputs and data URLs stop inflating
classification spend
This commit is contained in:
Tin Chi Lo 2026-09-01 19:49:43 -07:00
parent 52464ea464
commit 16bc8a8e17
5 changed files with 137 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/router_strategy/capability_router/capability_router.py
View file

@ -122,6 +122,16 @@ def _classification_context(messages: Sequence[Mapping[str, object]]) -> tuple[M
return tuple(selected)
def _capped(value: object, cap: int) -> object:
if isinstance(value, str):
return value if len(value) <= cap else f"{value[:cap]}...[truncated {len(value) - cap} chars]"
if isinstance(value, Mapping):
return {key: _capped(item, cap) for key, item in value.items()}
if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
return tuple(_capped(item, cap) for item in value)
return value
def _tool_names(request_kwargs: Mapping[str, object]) -> tuple[str, ...]:
tools = request_kwargs.get("tools")
if not isinstance(tools, Sequence) or isinstance(tools, (str, bytes)):
@ -187,11 +197,15 @@ class CapabilityRouter(CustomLogger):
context = _classification_context(messages)
if not context:
raise CapabilityClassifierFailure("No user task was available for capability classification")
payload = {
"conversation": context,
cap: Final = self.config.classifier.max_message_chars
payload: Final = {
"conversation": tuple(_capped(message, cap) for message in context),
"available_tools": _tool_names(request_kwargs),
}
return "Task context (untrusted JSON):\n" + json.dumps(payload, default=str, ensure_ascii=False)
return (
"Task context (untrusted JSON; long values truncated; the newest user message is the task to forecast):\n"
+ json.dumps(payload, default=str, ensure_ascii=False)
)
def _cache_key(
self,

12
litellm/router_strategy/capability_router/config.py
View file

@ -46,6 +46,7 @@ class CapabilityClassifierConfig(BaseModel):
model: str
timeout_ms: int = Field(default=3000, ge=1)
max_output_tokens: int = Field(default=1024, ge=1)
max_message_chars: int = Field(default=2000, ge=1)
model_config = ConfigDict(extra="forbid")
@ -61,6 +62,7 @@ class CapabilityRouterConfig(BaseModel):
candidates: tuple[CapabilityRouterCandidate, ...] = Field(min_length=2)
classifier: CapabilityClassifierConfig
probability_threshold: float = Field(default=0.7, ge=0.0, le=1.0)
threshold_step: float = Field(default=0.1, ge=0.0, le=1.0)
fallback_model: str
estimated_output_tokens: int = Field(default=1000, ge=1)
cache_ttl_seconds: int = Field(default=3600, ge=1)
@ -72,11 +74,11 @@ class CapabilityRouterConfig(BaseModel):
def validate_fallback_model(cls, value: str) -> str:
return _provider_model(value, "fallback_model")
@field_validator("probability_threshold", mode="before")
@field_validator("probability_threshold", "threshold_step", mode="before")
@classmethod
def validate_finite_threshold(cls, value: object) -> object:
if isinstance(value, (int, float)) and not math.isfinite(value):
raise ValueError("probability_threshold must be finite")
raise ValueError("threshold values must be finite")
return value
@model_validator(mode="after")
@ -89,12 +91,16 @@ class CapabilityRouterConfig(BaseModel):
return self
CapabilityBoundary = Literal["supported", "uncertain", "unsupported", "unmatched"]
class CapabilityCandidateScore(BaseModel):
"""One classifier estimate for the current task."""
model: str
p_solve: float = Field(ge=0.0, le=1.0)
reason: str
capability_boundary: CapabilityBoundary
p_solve: float = Field(ge=0.0, le=1.0)
model_config = ConfigDict(extra="forbid")

24
litellm/router_strategy/capability_router/policy.py
View file

@ -2,16 +2,28 @@
import math
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from pydantic import BaseModel, ConfigDict
from .config import CapabilityClassifierVerdict, CapabilityRouterConfig, CapabilitySelectionReason
from .config import (
CapabilityBoundary,
CapabilityClassifierVerdict,
CapabilityRouterConfig,
CapabilitySelectionReason,
)
BOUNDARY_THRESHOLD_STEPS: Final[Mapping[CapabilityBoundary, int]] = MappingProxyType(
{"supported": 0, "uncertain": 1, "unmatched": 1, "unsupported": 2}
)
class CapabilityCandidateAssessment(BaseModel):
model: str
p_solve: float
reason: str
capability_boundary: CapabilityBoundary
estimated_cost: float | None
qualified: bool
@ -43,7 +55,7 @@ def select_capability_model(
verdict: CapabilityClassifierVerdict,
estimated_costs: Mapping[str, float | None],
) -> CapabilityRoutingDecision:
"""Choose the cheapest candidate above the configured probability."""
"""Choose the cheapest candidate whose p_solve clears its boundary-stepped threshold."""
configured_models = tuple(candidate.model for candidate in config.candidates)
scores = {candidate.model: candidate for candidate in verdict.candidates}
if set(scores) != set(configured_models):
@ -54,8 +66,14 @@ def select_capability_model(
model=model,
p_solve=scores[model].p_solve,
reason=scores[model].reason,
capability_boundary=scores[model].capability_boundary,
estimated_cost=estimated_costs.get(model),
qualified=scores[model].p_solve > config.probability_threshold,
qualified=scores[model].p_solve
> round(
config.probability_threshold
+ BOUNDARY_THRESHOLD_STEPS[scores[model].capability_boundary] * config.threshold_step,
9,
),
)
for model in configured_models
)

26
litellm/router_strategy/capability_router/prompts.py
View file

@ -4,19 +4,28 @@ from typing import Any
from .config import CapabilityRouterConfig
CAPABILITY_BOUNDARIES = ("supported", "uncertain", "unsupported", "unmatched")
def build_classifier_prompt(config: CapabilityRouterConfig) -> str:
candidates = "\n".join(
f"- {candidate.model}: {candidate.description}" for candidate in config.candidates
)
return f"""You route tasks between language models.
candidates = "\n".join(f"- {candidate.model}: {candidate.description}" for candidate in config.candidates)
return f"""You forecast task outcomes for a model router.
For each candidate below, estimate the probability that it can complete the user's task correctly and completely. Base the estimate only on the task and the operator-provided description. Do not consider price; price is handled separately.
For each candidate model below, forecast one binary event. SUCCESS means the candidate completes the newest user task correctly and completely in one fresh attempt, using only the tools available in the request. FAILURE is any other outcome. The two outcomes are exhaustive.
Use only evidence in the conversation and each candidate's capability description. Do not assume hidden state, unmentioned tools, or future clarifications. The descriptions are qualitative evidence, not measured success rates.
Assessment procedure, per candidate:
1. State the crux in "reason": the hardest material requirement for whole-task success.
2. Set "capability_boundary": "supported" if the description covers the crux, "unsupported" if it excludes it, "uncertain" if coverage is unclear, "unmatched" if the description does not speak to this task.
3. Estimate "p_solve" last. It is the probability of SUCCESS, not confidence in this assessment and not a route recommendation.
Interpret p_solve as a natural frequency: at 0.7, about 70 of 100 comparable fresh attempts succeed. Use the full range when justified, and reserve 0 and 1 for outcomes that are logically impossible or certain. "supported" does not mean 1 and "unsupported" does not mean 0. Do not consider price or the routing threshold; both are handled separately.
Candidates:
{candidates}
Return one score for every candidate using the exact model names. p_solve must be a number from 0 to 1. reason must briefly identify the capability that most affects the estimate."""
Return one entry for every candidate using the exact model names."""
def build_classifier_response_schema(config: CapabilityRouterConfig) -> dict[str, Any]:
@ -32,10 +41,11 @@ def build_classifier_response_schema(config: CapabilityRouterConfig) -> dict[str
"type": "object",
"properties": {
"model": {"type": "string", "enum": model_names},
"p_solve": {"type": "number", "minimum": 0, "maximum": 1},
"reason": {"type": "string", "minLength": 1},
"capability_boundary": {"type": "string", "enum": list(CAPABILITY_BOUNDARIES)},
"p_solve": {"type": "number", "minimum": 0, "maximum": 1},
},
"required": ["model", "p_solve", "reason"],
"required": ["model", "reason", "capability_boundary", "p_solve"],
"additionalProperties": False,
},
}

81
tests/test_litellm/router_strategy/capability_router/test_capability_router.py
View file

@ -10,6 +10,7 @@ from litellm.router_strategy.capability_router.config import (
CapabilityRouterConfig,
)
from litellm.router_strategy.capability_router.policy import select_capability_model
from litellm.router_strategy.capability_router.prompts import build_classifier_response_schema
def config() -> dict:
@ -45,8 +46,8 @@ def test_policy_selects_cheapest_model_above_global_threshold() -> None:
verdict = CapabilityClassifierVerdict.model_validate(
{
"candidates": [
{"model": "small", "p_solve": 0.78, "reason": "clear bounded task"},
{"model": "frontier", "p_solve": 0.95, "reason": "more capable"},
{"model": "small", "capability_boundary": "supported", "p_solve": 0.78, "reason": "clear bounded task"},
{"model": "frontier", "capability_boundary": "supported", "p_solve": 0.95, "reason": "more capable"},
]
}
)
@ -63,8 +64,8 @@ def test_policy_falls_back_if_no_model_qualifies_or_price_is_unknown() -> None:
verdict = CapabilityClassifierVerdict.model_validate(
{
"candidates": [
{"model": "small", "p_solve": 0.4, "reason": "too hard"},
{"model": "frontier", "p_solve": 0.6, "reason": "uncertain"},
{"model": "small", "capability_boundary": "supported", "p_solve": 0.4, "reason": "too hard"},
{"model": "frontier", "capability_boundary": "supported", "p_solve": 0.6, "reason": "uncertain"},
]
}
)
@ -75,7 +76,8 @@ def test_policy_falls_back_if_no_model_qualifies_or_price_is_unknown() -> None:
qualified = verdict.model_copy(
update={
"candidates": tuple(
candidate.model_copy(update={"p_solve": 0.9}) for candidate in verdict.candidates
candidate.model_copy(update={"capability_boundary": "supported", "p_solve": 0.9})
for candidate in verdict.candidates
)
}
)
@ -89,8 +91,8 @@ def test_probability_must_be_strictly_above_threshold() -> None:
verdict = CapabilityClassifierVerdict.model_validate(
{
"candidates": [
{"model": "small", "p_solve": 0.7, "reason": "on the boundary"},
{"model": "frontier", "p_solve": 0.7, "reason": "on the boundary"},
{"model": "small", "capability_boundary": "supported", "p_solve": 0.7, "reason": "on the boundary"},
{"model": "frontier", "capability_boundary": "supported", "p_solve": 0.7, "reason": "on the boundary"},
]
}
)
@ -130,8 +132,13 @@ async def test_same_user_turn_reuses_cached_decision() -> None:
CapabilityClassifierVerdict.model_validate(
{
"candidates": [
{"model": "small", "p_solve": 0.9, "reason": "fits"},
{"model": "frontier", "p_solve": 0.95, "reason": "fits"},
{"model": "small", "capability_boundary": "supported", "p_solve": 0.9, "reason": "fits"},
{
"model": "frontier",
"capability_boundary": "supported",
"p_solve": 0.95,
"reason": "fits",
},
]
}
),
@ -154,3 +161,59 @@ async def test_same_user_turn_reuses_cached_decision() -> None:
assert first.routing_decision is not None and first.routing_decision["cached"] is False
assert second.routing_decision is not None and second.routing_decision["cached"] is True
strategy._new_decision.assert_awaited_once()
def test_boundary_buckets_step_the_effective_threshold() -> None:
parsed = CapabilityRouterConfig.model_validate(config())
verdict = CapabilityClassifierVerdict.model_validate(
{
"candidates": [
{"model": "small", "capability_boundary": "uncertain", "p_solve": 0.78, "reason": "ambiguous scope"},
{"model": "frontier", "capability_boundary": "supported", "p_solve": 0.78, "reason": "covered"},
]
}
)
decision = select_capability_model(parsed, verdict, {"small": 0.01, "frontier": 0.05})
assert decision.selected_model == "frontier"
assert [candidate.qualified for candidate in decision.candidates] == [False, True]
def test_unsupported_boundary_requires_two_threshold_steps() -> None:
parsed = CapabilityRouterConfig.model_validate({**config(), "threshold_step": 0.1})
unsupported = {"model": "small", "capability_boundary": "unsupported", "reason": "excluded"}
supported = {"model": "frontier", "capability_boundary": "supported", "p_solve": 0.95, "reason": "covered"}
costs = {"small": 0.01, "frontier": 0.05}
below = CapabilityClassifierVerdict.model_validate({"candidates": [{**unsupported, "p_solve": 0.9}, supported]})
above = CapabilityClassifierVerdict.model_validate({"candidates": [{**unsupported, "p_solve": 0.91}, supported]})
assert select_capability_model(parsed, below, costs).selected_model == "frontier"
assert select_capability_model(parsed, above, costs).selected_model == "small"
def test_classifier_payload_caps_long_message_values() -> None:
strategy = CapabilityRouter("cost-router", Router(model_list=[]), config())
messages = [
{"role": "user", "content": "Fix the failing build"},
{"role": "assistant", "content": [{"type": "text", "text": "x" * 50_000}]},
{"role": "user", "content": "now fix the tests"},
]
payload = strategy._classifier_payload(messages, {})
assert len(payload) < 10_000
assert "[truncated 48000 chars]" in payload
assert "newest user message is the task" in payload
def test_response_schema_orders_reasoning_before_probability() -> None:
schema = build_classifier_response_schema(CapabilityRouterConfig.model_validate(config()))
item = schema["properties"]["candidates"]["items"]
fields = list(item["properties"])
assert fields.index("reason") < fields.index("p_solve")
assert fields.index("capability_boundary") < fields.index("p_solve")
assert "capability_boundary" in item["required"]
assert item["properties"]["capability_boundary"]["enum"] == ["supported", "uncertain", "unsupported", "unmatched"]