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

fix(router): enforce strict probability threshold

Browse source
This commit is contained in:
moe-berri 2026-09-01 15:32:43 -07:00 • committed by Tin Chi Lo
parent d6d4416306
commit 52464ea464
2 changed files with 18 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -43,7 +43,7 @@ def select_capability_model(
verdict: CapabilityClassifierVerdict,
estimated_costs: Mapping[str, float | None],
) -> CapabilityRoutingDecision:
"""Choose the cheapest candidate at or above the configured probability."""
"""Choose the cheapest candidate above the configured probability."""
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):
@ -55,7 +55,7 @@ def select_capability_model(
p_solve=scores[model].p_solve,
reason=scores[model].reason,
estimated_cost=estimated_costs.get(model),
qualified=scores[model].p_solve >= config.probability_threshold,
qualified=scores[model].p_solve > config.probability_threshold,
)
for model in configured_models
)

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

@ -84,6 +84,22 @@ def test_policy_falls_back_if_no_model_qualifies_or_price_is_unknown() -> None:
)
def test_probability_must_be_strictly_above_threshold() -> None:
parsed = CapabilityRouterConfig.model_validate(config())
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"},
]
}
)
assert select_capability_model(parsed, verdict, {"small": 0.01, "frontier": 0.05}).reason == (
"no_qualified_candidate"
)
def test_router_registers_capability_strategy() -> None:
router = Router(
model_list=[