feat(router): calibrate matched capability rules

litellm/router_strategy/capability_router/README.md

@@ -7,6 +7,7 @@ The capability router asks an LLM to forecast each configured candidate's probab
 Capability-card training uses outcomes from the real task harness. It does not replace the LLM classifier. The classifier still reads the task and cards at runtime; the trained artifact improves the evidence and policy around its forecast by:
 
 - assigning each card rule a boundary from observed end-to-end outcomes and a 95% Wilson interval;
+- attaching a smoothed observed success probability to each matched rule;
 - fitting monotonic probability calibration for each candidate;
 - tuning the global probability threshold and boundary step on validation tasks; and
 - reporting raw and calibrated Brier score, log loss, calibration error, solve rate, cost, and quality-cost utility on untouched test tasks
@@ -54,12 +55,20 @@ Report the complete solve-rate versus cost curve rather than one threshold. Tune
 
 Recommended executable public suites are Terminal-Bench 2.1, SWE-bench Verified or Pro, tau2-bench, AppWorld, BFCL, and ToolSandbox. RouterBench is useful as a cheap classifier and calibration smoke test, but it is not evidence of agentic end-to-end performance
 
+### SWE-bench Verified coding run
+
+The coding run used all 500 SWE-bench Verified issues and public per-instance outcomes for Claude 4.5 Haiku, Sonnet, and Opus under the same mini-SWE-agent v2.0.0 harness and reasoning setting. Repository families were held apart: 291 issues across six repositories trained the cards, 97 issues across two repositories selected the operating point, and 112 issues across four different repositories formed the test split. The classifier saw only the repository name and the first 2,000 characters of the issue, matching its runtime context limit
+
+Rule-conditioned calibration improved held-out Brier score from 0.3047 to 0.1974, log loss from 0.9078 to 0.5829, and expected calibration error from 0.2362 to 0.0184. At the quality-first `0.95` objective weight, the validation-selected route matched the original cards at 80.36% test solve rate, while mean recorded cost changed from 0.7839 to 0.8263. A lower-cost point on the learned curve reached 77.68% solve rate at 0.6134 cost
+
+The repository holdout exposed that model-wide calibration alone removed too much task discrimination. Attaching a smoothed outcome probability to the matched capability rule raised the learned `0.9`-weight route from 67.86% to 70.54% solve rate and from 0.7107 to 0.7204 utility. It still did not beat the original cards on test routing utility, so this result supports the calibration and evaluation machinery rather than a coding preset quality claim
+
 ### tau2-bench pipeline run
 
 The pipeline was exercised against public tau2-bench trajectories for `claude-sonnet-4-5` and `claude-opus-4-5`, with four recorded end-to-end attempts aggregated per task. Entire domains were held apart: 50 airline tasks trained the artifact, 113 retail tasks selected the operating point, and 110 telecom tasks were evaluated once. A local `mlx-community/Qwen3-4B-Instruct-2507-4bit` model produced the capability forecasts
 
 At a `0.7` quality weight, the validation-selected configuration reached 85.00% test solve rate at 0.4678 mean recorded cost. The original cards reached 92.27% at 0.7057 cost. This is 33.7% lower cost with a 7.27-point solve-rate loss, and improves the configured normalized utility from 0.6650 to 0.8950. The learned test curve also contains a 91.36% solve-rate point at 0.6743 cost, 4.4% below the original cost with a 0.91-point solve-rate loss
 
-Calibration generalized across the held-out domain. Brier score improved from 0.2207 to 0.0568, log loss from 0.7657 to 0.3994, and expected calibration error from 0.2984 to 0.1224
+Calibration generalized across the held-out domain. Brier score improved from 0.2207 to 0.0687, log loss from 0.7657 to 0.4326, and expected calibration error from 0.2984 to 0.1627
 
 This run validates artifact training, domain-disjoint evaluation, and the quality-cost tradeoff. It does not show a strict raw solve-rate improvement over the original cards, and it is not a direct comparison with a published Switchyard result. Larger cross-benchmark training data and another untouched test family are required before treating these cards as a general preset

litellm/router_strategy/capability_router/config.py

@@ -29,6 +29,7 @@ class CapabilityRule(BaseModel):
 
     boundary: CapabilityRuleBoundary
     rule: str
+    observed_success_probability: float | None = Field(default=None, ge=0.0, le=1.0)
 
     model_config = ConfigDict(extra="forbid")
 

litellm/router_strategy/capability_router/policy.py

@@ -32,7 +32,19 @@ def effective_boundary(candidate: CapabilityRouterCandidate, score: CapabilityCa
     return boundaries.get(score.primary_rule, "unmatched")
 
 
-def calibrated_probability(candidate: CapabilityRouterCandidate, raw_probability: float) -> float:
+def calibrated_probability(
+    candidate: CapabilityRouterCandidate, raw_probability: float, primary_rule: str = "none"
+) -> float:
+    rule_probability: Final = next(
+        (
+            rule.observed_success_probability
+            for rule_id, rule in indexed_rules(candidate)
+            if rule_id == primary_rule and rule.observed_success_probability is not None
+        ),
+        None,
+    )
+    if rule_probability is not None:
+        return rule_probability
     return next(
         (bucket.probability for bucket in candidate.probability_calibration if raw_probability <= bucket.upper_bound),
         raw_probability,
@@ -90,12 +102,12 @@ def select_capability_model(
         CapabilityCandidateAssessment(
             model=model,
             raw_p_solve=scores[model].p_solve,
-            p_solve=calibrated_probability(configured[model], scores[model].p_solve),
+            p_solve=calibrated_probability(configured[model], scores[model].p_solve, scores[model].primary_rule),
             reason=scores[model].reason,
             primary_rule=scores[model].primary_rule,
             capability_boundary=boundaries[model],
             estimated_cost=estimated_costs.get(model),
-            qualified=calibrated_probability(configured[model], scores[model].p_solve)
+            qualified=calibrated_probability(configured[model], scores[model].p_solve, scores[model].primary_rule)
             > round(
                 config.probability_threshold + BOUNDARY_THRESHOLD_STEPS[boundaries[model]] * config.threshold_step,
                 9,

litellm/router_strategy/capability_router/training.py

@@ -250,7 +250,16 @@ def _train_candidate(
         for _, rule, rule_records in learned
     )
     rules: Final = tuple(
-        CapabilityRule(boundary=boundary[0], rule=rule.rule) for (_, rule, _), boundary in zip(learned, boundaries)
+        CapabilityRule(
+            boundary=boundary[0],
+            rule=rule.rule,
+            observed_success_probability=(
+                (sum(record.success for record in rule_records) + 1.0) / (len(rule_records) + 2.0)
+                if rule_records
+                else None
+            ),
+        )
+        for (_, rule, rule_records), boundary in zip(learned, boundaries)
     )
     statistics: Final = tuple(
         CapabilityRuleStatistic(
@@ -317,7 +326,9 @@ def _aggregate_candidate(
     primary_rule: Final = min(rule_counts, key=lambda item: (-item[1], item[0]))[0]
     return _TaskCandidate(
         model=candidate.model,
-        probability=calibrated_probability(candidate, raw_probability) if calibrated else raw_probability,
+        probability=(
+            calibrated_probability(candidate, raw_probability, primary_rule) if calibrated else raw_probability
+        ),
         success=sum(record.success for record in records) / len(records),
         cost=sum(record.estimated_cost for record in records) / len(records),
         boundary=_effective_boundary(candidate, primary_rule),
@@ -424,7 +435,9 @@ def _probability_metrics(
     candidates: Final = MappingProxyType({candidate.model: candidate for candidate in config.candidates})
     rows: Final = tuple(record for record in records if record.split == "test" and record.model in candidates)
     predictions: Final = tuple(
-        calibrated_probability(candidates[record.model], record.raw_p_solve) if calibrated else record.raw_p_solve
+        calibrated_probability(candidates[record.model], record.raw_p_solve, record.primary_rule)
+        if calibrated
+        else record.raw_p_solve
         for record in rows
     )
     outcomes: Final = tuple(record.success for record in rows)

tests/test_litellm/router_strategy/capability_router/test_capability_router.py

@@ -216,6 +216,43 @@ def test_policy_qualifies_on_calibrated_probability_and_keeps_raw_forecast() ->
     assert decision.selected_model == "frontier"
 
 
+def test_policy_prefers_matched_rule_probability_over_model_calibration() -> None:
+    configured = config()
+    configured["candidates"][0]["rules"] = [
+        {
+            "boundary": "supported",
+            "rule": "The task has a bounded verifier",
+            "observed_success_probability": 0.85,
+        }
+    ]
+    configured["candidates"][0]["probability_calibration"] = [{"upper_bound": 1.0, "probability": 0.4}]
+    parsed = CapabilityRouterConfig.model_validate(configured)
+    verdict = CapabilityClassifierVerdict.model_validate(
+        {
+            "candidates": [
+                {
+                    "model": "small",
+                    "primary_rule": "R1",
+                    "capability_boundary": "supported",
+                    "p_solve": 0.2,
+                    "reason": "bounded verifier",
+                },
+                {
+                    "model": "frontier",
+                    "capability_boundary": "supported",
+                    "p_solve": 0.95,
+                    "reason": "covered",
+                },
+            ]
+        }
+    )
+
+    decision = select_capability_model(parsed, verdict, {"small": 0.01, "frontier": 0.05})
+
+    assert decision.candidates[0].p_solve == 0.85
+    assert decision.selected_model == "small"
+
+
 def test_boundary_buckets_step_the_effective_threshold() -> None:
     parsed = CapabilityRouterConfig.model_validate(config())
     verdict = CapabilityClassifierVerdict.model_validate(

tests/test_litellm/router_strategy/capability_router/test_training.py

@@ -65,6 +65,7 @@ def test_training_learns_rule_boundaries_and_improves_held_out_routing() -> None
     small = result.artifact.config.candidates[0]
 
     assert [rule.boundary for rule in small.rules] == ["supported", "unsupported"]
+    assert [rule.observed_success_probability for rule in small.rules] == pytest.approx([21.0 / 22.0, 1.0 / 22.0])
     assert small.probability_calibration[-1].upper_bound == 1.0
     assert result.artifact.records == len(records())
     assert len(result.artifact.records_sha256) == 64