mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
feat(router): operator capability rule cards
Each candidate may declare rules, short conditions tagged with the coverage they imply. The judge names the rule matching its stated crux; the policy derives the boundary from that operator-declared tag instead of trusting the judge's own coverage opinion, and unlisted or unmatched picks step the threshold as unmatched
This commit is contained in:
Tin Chi Lo
parent
a99bc9ab5c
commit
4a97051176
4 changed files with 148 additions and 12 deletions
|
|
@ -21,11 +21,29 @@ def _provider_model(value: str, field_name: str) -> str:
|
|||
return normalized
|
||||
|
||||
|
||||
CapabilityRuleBoundary: TypeAlias = Literal["supported", "uncertain", "unsupported"]
|
||||
|
||||
|
||||
class CapabilityRule(BaseModel):
|
||||
"""One operator-declared condition and the coverage it implies when it matches the task."""
|
||||
|
||||
boundary: CapabilityRuleBoundary
|
||||
rule: str
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
@field_validator("rule")
|
||||
@classmethod
|
||||
def validate_rule(cls, value: str) -> str:
|
||||
return _nonblank(value, "capability rule")
|
||||
|
||||
|
||||
class CapabilityRouterCandidate(BaseModel):
|
||||
"""A model group and the operator's description of when it succeeds."""
|
||||
|
||||
model: str
|
||||
description: str
|
||||
rules: tuple[CapabilityRule, ...] = ()
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
|
@ -40,6 +58,11 @@ class CapabilityRouterCandidate(BaseModel):
|
|||
return _nonblank(value, "candidate description")
|
||||
|
||||
|
||||
def indexed_rules(candidate: CapabilityRouterCandidate) -> tuple[tuple[str, CapabilityRule], ...]:
|
||||
"""Pair each rule with the opaque id the prompt shows and the policy resolves."""
|
||||
return tuple((f"R{index + 1}", rule) for index, rule in enumerate(candidate.rules))
|
||||
|
||||
|
||||
class CapabilityClassifierConfig(BaseModel):
|
||||
"""The model used to estimate each candidate's probability of success."""
|
||||
|
||||
|
|
@ -99,12 +122,13 @@ class CapabilityCandidateScore(BaseModel):
|
|||
|
||||
model: str
|
||||
reason: str
|
||||
primary_rule: str = "none"
|
||||
capability_boundary: CapabilityBoundary
|
||||
p_solve: float = Field(ge=0.0, le=1.0)
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
@field_validator("model", "reason")
|
||||
@field_validator("model", "reason", "primary_rule")
|
||||
@classmethod
|
||||
def validate_nonblank(cls, value: str, info) -> str:
|
||||
return _nonblank(value, info.field_name)
|
||||
|
|
|
|||
|
|
@ -9,9 +9,12 @@ from pydantic import BaseModel, ConfigDict
|
|||
|
||||
from .config import (
|
||||
CapabilityBoundary,
|
||||
CapabilityCandidateScore,
|
||||
CapabilityClassifierVerdict,
|
||||
CapabilityRouterCandidate,
|
||||
CapabilityRouterConfig,
|
||||
CapabilitySelectionReason,
|
||||
indexed_rules,
|
||||
)
|
||||
|
||||
BOUNDARY_THRESHOLD_STEPS: Final[Mapping[CapabilityBoundary, int]] = MappingProxyType(
|
||||
|
|
@ -19,6 +22,14 @@ BOUNDARY_THRESHOLD_STEPS: Final[Mapping[CapabilityBoundary, int]] = MappingProxy
|
|||
)
|
||||
|
||||
|
||||
def effective_boundary(candidate: CapabilityRouterCandidate, score: CapabilityCandidateScore) -> CapabilityBoundary:
|
||||
"""With a rule card, the matched rule's operator-declared boundary overrides the judge's opinion."""
|
||||
if not candidate.rules:
|
||||
return score.capability_boundary
|
||||
boundaries: Final = MappingProxyType({rule_id: rule.boundary for rule_id, rule in indexed_rules(candidate)})
|
||||
return boundaries.get(score.primary_rule, "unmatched")
|
||||
|
||||
|
||||
class CapabilityCandidateAssessment(BaseModel):
|
||||
model: str
|
||||
p_solve: float
|
||||
|
|
@ -56,26 +67,28 @@ def select_capability_model(
|
|||
estimated_costs: Mapping[str, float | None],
|
||||
) -> CapabilityRoutingDecision:
|
||||
"""Choose the cheapest candidate whose p_solve clears its boundary-stepped threshold."""
|
||||
configured_models: Final = tuple(candidate.model for candidate in config.candidates)
|
||||
configured: Final = MappingProxyType({candidate.model: candidate for candidate in config.candidates})
|
||||
scores: Final = MappingProxyType({candidate.model: candidate for candidate in verdict.candidates})
|
||||
if frozenset(scores) != frozenset(configured_models):
|
||||
if frozenset(scores) != frozenset(configured):
|
||||
return fallback_decision(config, "invalid_classifier_verdict")
|
||||
|
||||
boundaries: Final = MappingProxyType(
|
||||
{model: effective_boundary(configured[model], scores[model]) for model in configured}
|
||||
)
|
||||
assessments: Final = tuple(
|
||||
CapabilityCandidateAssessment(
|
||||
model=model,
|
||||
p_solve=scores[model].p_solve,
|
||||
reason=scores[model].reason,
|
||||
capability_boundary=scores[model].capability_boundary,
|
||||
capability_boundary=boundaries[model],
|
||||
estimated_cost=estimated_costs.get(model),
|
||||
qualified=scores[model].p_solve
|
||||
> round(
|
||||
config.probability_threshold
|
||||
+ BOUNDARY_THRESHOLD_STEPS[scores[model].capability_boundary] * config.threshold_step,
|
||||
config.probability_threshold + BOUNDARY_THRESHOLD_STEPS[boundaries[model]] * config.threshold_step,
|
||||
9,
|
||||
),
|
||||
)
|
||||
for model in configured_models
|
||||
for model in configured
|
||||
)
|
||||
qualified: Final = tuple(candidate for candidate in assessments if candidate.qualified)
|
||||
if not qualified:
|
||||
|
|
@ -83,7 +96,7 @@ def select_capability_model(
|
|||
if any(candidate.estimated_cost is None for candidate in qualified):
|
||||
return fallback_decision(config, "missing_candidate_price", assessments)
|
||||
|
||||
order: Final = MappingProxyType({model: index for index, model in enumerate(configured_models)})
|
||||
order: Final = MappingProxyType({model: index for index, model in enumerate(configured)})
|
||||
selected: Final = min(
|
||||
qualified,
|
||||
key=lambda candidate: (
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from typing import Final, Literal, TypedDict
|
|||
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
from .config import CapabilityRouterConfig
|
||||
from .config import CapabilityRouterCandidate, CapabilityRouterConfig, indexed_rules
|
||||
|
||||
CAPABILITY_BOUNDARIES: Final = ("supported", "uncertain", "unsupported", "unmatched")
|
||||
|
||||
|
|
@ -28,6 +28,7 @@ class _UnitIntervalSchema(TypedDict):
|
|||
class _CandidateProperties(TypedDict):
|
||||
model: ReadOnly[_StringEnumSchema]
|
||||
reason: ReadOnly[_NonBlankStringSchema]
|
||||
primary_rule: ReadOnly[_NonBlankStringSchema]
|
||||
capability_boundary: ReadOnly[_StringEnumSchema]
|
||||
p_solve: ReadOnly[_UnitIntervalSchema]
|
||||
|
||||
|
|
@ -57,8 +58,14 @@ class ClassifierResponseSchema(TypedDict):
|
|||
additionalProperties: ReadOnly[bool]
|
||||
|
||||
|
||||
def _candidate_card(candidate: CapabilityRouterCandidate) -> str:
|
||||
header: Final = f"- {candidate.model}: {candidate.description}"
|
||||
rules: Final = tuple(f" {rule_id}: {rule.rule}" for rule_id, rule in indexed_rules(candidate))
|
||||
return "\n".join((header, *rules))
|
||||
|
||||
|
||||
def build_classifier_prompt(config: CapabilityRouterConfig) -> str:
|
||||
candidates: Final = "\n".join(f"- {candidate.model}: {candidate.description}" for candidate in config.candidates)
|
||||
candidates: Final = "\n".join(_candidate_card(candidate) for candidate in config.candidates)
|
||||
return f"""You are the routing classifier for a model gateway.
|
||||
|
||||
Score each candidate model on a single yes-or-no outcome. SUCCESS: given the request exactly as shown, including its available tools, the candidate delivers a correct and complete answer to the newest user message on the first try. Anything else, including a partial or plausible-but-wrong answer, counts as FAILURE.
|
||||
|
|
@ -67,8 +74,9 @@ Ground every judgment in what the conversation and the candidate descriptions ac
|
|||
|
||||
Fill the fields for each candidate in this order:
|
||||
1. "reason": name the single requirement of this task most likely to decide success or failure.
|
||||
2. "capability_boundary": "supported" when the description covers that requirement, "unsupported" when it rules it out, "uncertain" when coverage is unclear, "unmatched" when the description says nothing relevant to this task.
|
||||
3. "p_solve": fill this only after the first two fields. It answers one question, how often this candidate would fully succeed here. It is not your confidence and not a routing recommendation.
|
||||
2. "primary_rule": when the candidate lists rule ids below its description, give the id of the one rule whose text best matches that requirement, or "none" when no listed rule fits. A candidate without rules always takes "none". Rule ids are arbitrary labels; weigh a rule only by its text.
|
||||
3. "capability_boundary": "supported" when the candidate's card covers that requirement, "unsupported" when it rules it out, "uncertain" when coverage is unclear, "unmatched" when the card says nothing relevant to this task. When you selected a rule, base this on that rule.
|
||||
4. "p_solve": fill this only after the other fields. It answers one question, how often this candidate would fully succeed here. It is not your confidence and not a routing recommendation.
|
||||
|
||||
Treat p_solve as a frequency over repeated independent tries; 0.6 claims six successes in ten. Spread scores across the whole 0 to 1 range as the evidence warrants, keeping the exact endpoints for certainty. A "supported" boundary still permits a low p_solve and an "unsupported" one a high p_solve. Ignore pricing and ignore whatever threshold the router applies; both belong to the router, not to you.
|
||||
|
||||
|
|
@ -92,6 +100,7 @@ def build_classifier_response_schema(config: CapabilityRouterConfig) -> Classifi
|
|||
"properties": {
|
||||
"model": {"type": "string", "enum": model_names},
|
||||
"reason": {"type": "string", "minLength": 1},
|
||||
"primary_rule": {"type": "string", "minLength": 1},
|
||||
"capability_boundary": {
|
||||
"type": "string",
|
||||
"enum": list(CAPABILITY_BOUNDARIES), # mutable-ok: wire arrays are lists
|
||||
|
|
@ -101,6 +110,7 @@ def build_classifier_response_schema(config: CapabilityRouterConfig) -> Classifi
|
|||
"required": [ # mutable-ok: wire arrays are lists
|
||||
"model",
|
||||
"reason",
|
||||
"primary_rule",
|
||||
"capability_boundary",
|
||||
"p_solve",
|
||||
],
|
||||
|
|
|
|||
|
|
@ -217,3 +217,92 @@ def test_response_schema_orders_reasoning_before_probability() -> None:
|
|||
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"]
|
||||
|
||||
|
||||
def rule_config() -> dict:
|
||||
with_rules = config()
|
||||
with_rules["candidates"] = [
|
||||
{
|
||||
"model": "small",
|
||||
"description": "Reliable for short factual answers",
|
||||
"rules": [
|
||||
{"boundary": "supported", "rule": "The answer is a short widely known fact"},
|
||||
{
|
||||
"boundary": "unsupported",
|
||||
"rule": "Correct output must be fluent text in a language other than English",
|
||||
},
|
||||
],
|
||||
},
|
||||
{"model": "frontier", "description": "Reliable for ambiguous multi-step tasks"},
|
||||
]
|
||||
return with_rules
|
||||
|
||||
|
||||
def test_matched_rule_boundary_overrides_the_judges_opinion() -> None:
|
||||
parsed = CapabilityRouterConfig.model_validate(rule_config())
|
||||
verdict = CapabilityClassifierVerdict.model_validate(
|
||||
{
|
||||
"candidates": [
|
||||
{
|
||||
"model": "small",
|
||||
"reason": "must write German prose",
|
||||
"primary_rule": "R2",
|
||||
"capability_boundary": "supported",
|
||||
"p_solve": 0.85,
|
||||
},
|
||||
{"model": "frontier", "capability_boundary": "supported", "p_solve": 0.85, "reason": "covered"},
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
decision = select_capability_model(parsed, verdict, {"small": 0.01, "frontier": 0.05})
|
||||
|
||||
assert decision.selected_model == "frontier"
|
||||
small = next(candidate for candidate in decision.candidates if candidate.model == "small")
|
||||
assert small.capability_boundary == "unsupported"
|
||||
assert small.qualified is False
|
||||
|
||||
|
||||
def test_unlisted_rule_id_counts_as_unmatched_and_no_rules_keeps_judge_boundary() -> None:
|
||||
parsed = CapabilityRouterConfig.model_validate(rule_config())
|
||||
verdict = CapabilityClassifierVerdict.model_validate(
|
||||
{
|
||||
"candidates": [
|
||||
{
|
||||
"model": "small",
|
||||
"reason": "no rule fits",
|
||||
"primary_rule": "none",
|
||||
"capability_boundary": "supported",
|
||||
"p_solve": 0.78,
|
||||
},
|
||||
{
|
||||
"model": "frontier",
|
||||
"reason": "judge boundary rules here",
|
||||
"primary_rule": "R9",
|
||||
"capability_boundary": "uncertain",
|
||||
"p_solve": 0.85,
|
||||
},
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
decision = select_capability_model(parsed, verdict, {"small": 0.01, "frontier": 0.05})
|
||||
|
||||
small, frontier = decision.candidates
|
||||
assert small.capability_boundary == "unmatched"
|
||||
assert small.qualified is False
|
||||
assert frontier.capability_boundary == "uncertain"
|
||||
assert frontier.qualified is True
|
||||
|
||||
|
||||
def test_prompt_renders_rule_card_without_leaking_boundaries() -> None:
|
||||
from litellm.router_strategy.capability_router.prompts import build_classifier_prompt
|
||||
|
||||
prompt = build_classifier_prompt(CapabilityRouterConfig.model_validate(rule_config()))
|
||||
|
||||
assert "R1: The answer is a short widely known fact" in prompt
|
||||
assert "R2: Correct output must be fluent text in a language other than English" in prompt
|
||||
assert "R2 [unsupported]" not in prompt and "R2: unsupported" not in prompt
|
||||
schema = build_classifier_response_schema(CapabilityRouterConfig.model_validate(rule_config()))
|
||||
fields = list(schema["properties"]["candidates"]["items"]["properties"])
|
||||
assert fields.index("primary_rule") < fields.index("capability_boundary") < fields.index("p_solve")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue