feat(router): train capability cards from outcomes

litellm/router_strategy/capability_router/README.md

@@ -0,0 +1,55 @@
+# Capability router
+
+The capability router asks an LLM to forecast each configured candidate's probability of completing the current task, then selects the cheapest candidate whose forecast clears the configured reliability threshold. This is separate from the complexity router and its local or trained heuristic classifiers
+
+## Train capability cards
+
+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;
+- 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
+
+Create one JSONL row for every `(task, candidate, run)`:
+
+```json
+{"benchmark":"terminal-bench-2.1","task_id":"build-linux-kernel-qemu","split":"train","model":"efficient","primary_rule":"R3","raw_p_solve":0.72,"success":1.0,"estimated_cost":0.18}
+```
+
+`raw_p_solve` and `primary_rule` come from the capability classifier's routing decision. `success` must come from the benchmark's end-to-end verifier, not an LLM estimate. Repeated runs may use the same task and model; the trainer averages them when measuring routing quality
+
+Every benchmark task must have an explicit `train`, `validation`, or `test` split and outcomes for every configured candidate. A task cannot cross splits. For a general preset, assign entire benchmark families to a split and deduplicate related tasks before training. A random row split over near-duplicate tasks overstates generalization
+
+Run:
+
+```shell
+python -m litellm.router_strategy.capability_router.training outcomes.jsonl \
+  --config capability-router.json \
+  --artifact-output trained-capability-router.json \
+  --quality-weight 0.7
+```
+
+The artifact wraps a ready-to-use `CapabilityRouterConfig`. Its candidate rules carry learned boundaries and `probability_calibration` bins. The report printed to stdout compares the trained route with the original untrained cards on the same test tasks
+
+### Seed-card design
+
+Write rules around observable failure mechanisms, not broad labels such as "easy", "hard", "coding", or "reasoning". A useful rule tells the classifier what property decides success: whether the procedure is explicit, the required state is inspectable, a validator covers the output, policy conditions conflict, an action is irreversible, or complete search has no boundary. Match the hardest requirement rather than an easy side task
+
+Keep rule text shared across candidates when possible, then let each candidate's learned boundary express its coverage. Do not put prices, routing instructions, thresholds, or claimed percentages in a card. Those belong to deterministic policy and the learned calibration artifact
+
+Retain an existing boundary when a rule has no training observations. With observations, the trainer changes it to supported only when the lower confidence bound clears the reliability target, unsupported only when the upper bound misses it, and uncertain otherwise
+
+## Benchmark protocol
+
+Use the same candidate model revisions, agent, tools, task budget, provider settings, and number of attempts for every arm. At minimum compare:
+
+1. always efficient;
+2. always capable;
+3. the original Switchyard-style qualitative cards and raw `p_solve`;
+4. the trained cards and calibrated `p_solve`; and
+5. an oracle computed from the recorded candidate outcomes
+
+Report the complete solve-rate versus cost curve rather than one threshold. Tune cards, calibration, and thresholds on training and validation tasks only. Run the final configuration once on the test split and keep that result unchanged
+
+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

litellm/router_strategy/capability_router/capability_router.py

@@ -403,6 +403,15 @@ class CapabilityRouter(CustomLogger):
             candidate_probabilities={  # mutable-ok: safe_dumps stringifies non-dict mappings in the spend log
                 candidate.model: candidate.p_solve for candidate in decision.candidates
             },
+            raw_candidate_probabilities={  # mutable-ok: safe_dumps stringifies non-dict mappings in the spend log
+                candidate.model: candidate.raw_p_solve for candidate in decision.candidates
+            },
+            candidate_rules={  # mutable-ok: safe_dumps stringifies non-dict mappings in the spend log
+                candidate.model: candidate.primary_rule for candidate in decision.candidates
+            },
+            candidate_boundaries={  # mutable-ok: safe_dumps stringifies non-dict mappings in the spend log
+                candidate.model: candidate.capability_boundary for candidate in decision.candidates
+            },
             candidate_costs={  # mutable-ok: safe_dumps stringifies non-dict mappings in the spend log
                 candidate.model: candidate.estimated_cost
                 for candidate in decision.candidates

litellm/router_strategy/capability_router/config.py

@@ -3,7 +3,7 @@
 import math
 from typing import Final, Literal, TypeAlias
 
-from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator, model_validator
 from typing_extensions import Self
 
 
@@ -38,12 +38,29 @@ class CapabilityRule(BaseModel):
         return _nonblank(value, "capability rule")
 
 
+class CapabilityCalibrationBin(BaseModel):
+    """One monotonic post-hoc calibration bucket learned from end-to-end outcomes."""
+
+    upper_bound: float = Field(ge=0.0, le=1.0)
+    probability: float = Field(ge=0.0, le=1.0)
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    @field_validator("upper_bound", "probability", mode="before")
+    @classmethod
+    def validate_finite_value(cls, value: object) -> object:
+        if isinstance(value, (int, float)) and not math.isfinite(value):
+            raise ValueError("calibration values must be finite")
+        return value
+
+
 class CapabilityRouterCandidate(BaseModel):
     """A model group and the operator's description of when it succeeds."""
 
     model: str
     description: str
     rules: tuple[CapabilityRule, ...] = ()
+    probability_calibration: tuple[CapabilityCalibrationBin, ...] = ()
 
     model_config = ConfigDict(extra="forbid")
 
@@ -57,6 +74,18 @@ class CapabilityRouterCandidate(BaseModel):
     def validate_description(cls, value: str) -> str:
         return _nonblank(value, "candidate description")
 
+    @model_validator(mode="after")
+    def validate_probability_calibration(self) -> Self:
+        upper_bounds: Final = tuple(bucket.upper_bound for bucket in self.probability_calibration)
+        probabilities: Final = tuple(bucket.probability for bucket in self.probability_calibration)
+        if any(right <= left for left, right in zip(upper_bounds, upper_bounds[1:])):
+            raise ValueError("calibration upper bounds must be strictly increasing")
+        if any(right < left for left, right in zip(probabilities, probabilities[1:])):
+            raise ValueError("calibration probabilities must be nondecreasing")
+        if upper_bounds and upper_bounds[-1] != 1.0:
+            raise ValueError("the final calibration upper bound must be 1")
+        return self
+
 
 def indexed_rules(candidate: CapabilityRouterCandidate) -> tuple[tuple[str, CapabilityRule], ...]:
     """Pair each rule with the opaque id the prompt shows and the policy resolves."""
@@ -130,8 +159,8 @@ class CapabilityCandidateScore(BaseModel):
 
     @field_validator("model", "reason", "primary_rule")
     @classmethod
-    def validate_nonblank(cls, value: str, info) -> str:
-        return _nonblank(value, info.field_name)
+    def validate_nonblank(cls, value: str, info: ValidationInfo) -> str:
+        return _nonblank(value, info.field_name or "classifier field")
 
     @field_validator("p_solve", mode="before")
     @classmethod

litellm/router_strategy/capability_router/policy.py

@@ -26,14 +26,29 @@ def effective_boundary(candidate: CapabilityRouterCandidate, score: CapabilityCa
     """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)})
+    boundaries: Final[Mapping[str, CapabilityBoundary]] = MappingProxyType(
+        {rule_id: rule.boundary for rule_id, rule in indexed_rules(candidate)}
+    )
     return boundaries.get(score.primary_rule, "unmatched")
 
 
+def calibrated_probability(candidate: CapabilityRouterCandidate, raw_probability: float) -> float:
+    return next(
+        (
+            bucket.probability
+            for bucket in candidate.probability_calibration
+            if raw_probability <= bucket.upper_bound
+        ),
+        raw_probability,
+    )
+
+
 class CapabilityCandidateAssessment(BaseModel):
     model: str
+    raw_p_solve: float
     p_solve: float
     reason: str
+    primary_rule: str
     capability_boundary: CapabilityBoundary
     estimated_cost: float | None
     qualified: bool
@@ -72,17 +87,19 @@ def select_capability_model(
     if frozenset(scores) != frozenset(configured):
         return fallback_decision(config, "invalid_classifier_verdict")
 
-    boundaries: Final = MappingProxyType(
+    boundaries: Final[Mapping[str, CapabilityBoundary]] = MappingProxyType(
         {model: effective_boundary(configured[model], scores[model]) for model in configured}
     )
     assessments: Final = tuple(
         CapabilityCandidateAssessment(
             model=model,
-            p_solve=scores[model].p_solve,
+            raw_p_solve=scores[model].p_solve,
+            p_solve=calibrated_probability(configured[model], scores[model].p_solve),
             reason=scores[model].reason,
+            primary_rule=scores[model].primary_rule,
             capability_boundary=boundaries[model],
             estimated_cost=estimated_costs.get(model),
-            qualified=scores[model].p_solve
+            qualified=calibrated_probability(configured[model], scores[model].p_solve)
             > round(
                 config.probability_threshold + BOUNDARY_THRESHOLD_STEPS[boundaries[model]] * config.threshold_step,
                 9,

litellm/router_strategy/capability_router/training.py

@@ -0,0 +1,576 @@
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+from collections import Counter
+from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import dataclass
+from itertools import groupby
+from pathlib import Path
+from typing import Final, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .config import (
+    CapabilityBoundary,
+    CapabilityCalibrationBin,
+    CapabilityRouterCandidate,
+    CapabilityRouterConfig,
+    CapabilityRule,
+    CapabilityRuleBoundary,
+    indexed_rules,
+)
+from .policy import BOUNDARY_THRESHOLD_STEPS, calibrated_probability
+
+_THRESHOLDS: Final = (0.5, 0.6, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95)
+_THRESHOLD_STEPS: Final = (0.0, 0.05, 0.1, 0.15)
+
+
+class CapabilityTrainingRecord(BaseModel):
+    benchmark: str = Field(min_length=1)
+    task_id: str = Field(min_length=1)
+    split: Literal["train", "validation", "test"]
+    model: str = Field(min_length=1)
+    primary_rule: str = Field(default="none", min_length=1)
+    raw_p_solve: float = Field(ge=0.0, le=1.0)
+    success: float = Field(ge=0.0, le=1.0)
+    estimated_cost: float = Field(ge=0.0)
+
+    model_config = ConfigDict(extra="forbid", frozen=True, allow_inf_nan=False)
+
+
+class CapabilityRuleStatistic(BaseModel):
+    model: str
+    rule_id: str
+    observations: int
+    success_rate: float
+    interval_low: float
+    interval_high: float
+    learned_boundary: CapabilityBoundary
+
+    model_config = ConfigDict(frozen=True)
+
+
+class CapabilityProbabilityMetrics(BaseModel):
+    observations: int
+    brier: float
+    log_loss: float
+    ece: float
+
+    model_config = ConfigDict(frozen=True)
+
+
+class CapabilityRouteMetrics(BaseModel):
+    tasks: int
+    success_rate: float
+    mean_cost: float
+    normalized_cost: float
+    quality_cost_utility: float
+
+    model_config = ConfigDict(frozen=True)
+
+
+class CapabilityThresholdPoint(BaseModel):
+    probability_threshold: float
+    threshold_step: float
+    metrics: CapabilityRouteMetrics
+
+    model_config = ConfigDict(frozen=True)
+
+
+class CapabilityTrainingReport(BaseModel):
+    objective: str
+    validation: CapabilityRouteMetrics
+    test: CapabilityRouteMetrics
+    test_untrained: CapabilityRouteMetrics
+    test_always_candidates: Mapping[str, CapabilityRouteMetrics]
+    test_oracle: CapabilityRouteMetrics
+    test_threshold_sweep: tuple[CapabilityThresholdPoint, ...]
+    test_probability_raw: CapabilityProbabilityMetrics
+    test_probability_calibrated: CapabilityProbabilityMetrics
+
+    model_config = ConfigDict(frozen=True)
+
+
+class CapabilityTrainingArtifact(BaseModel):
+    config: CapabilityRouterConfig
+    rule_statistics: tuple[CapabilityRuleStatistic, ...]
+    datasets: tuple[str, ...]
+    records: int
+    split_counts: Mapping[str, int]
+    records_sha256: str
+    split_contract: str
+
+    model_config = ConfigDict(frozen=True)
+
+
+class _Arguments(BaseModel):
+    records: Path
+    config: Path
+    artifact_output: Path
+    quality_weight: float
+
+
+@dataclass(frozen=True, slots=True)
+class CapabilityTrainingResult:
+    artifact: CapabilityTrainingArtifact
+    report: CapabilityTrainingReport
+
+
+@dataclass(frozen=True, slots=True)
+class _CalibrationBlock:
+    upper_bound: float
+    successes: float
+    observations: int
+
+    @property
+    def probability(self) -> float:
+        return (self.successes + 1.0) / (self.observations + 2.0)
+
+
+@dataclass(frozen=True, slots=True)
+class _TaskCandidate:
+    model: str
+    probability: float
+    success: float
+    cost: float
+    boundary: CapabilityBoundary
+
+
+def _pool_adjacent_violators(blocks: tuple[_CalibrationBlock, ...]) -> tuple[_CalibrationBlock, ...]:
+    violation: Final = next(
+        (index for index in range(len(blocks) - 1) if blocks[index].probability > blocks[index + 1].probability),
+        None,
+    )
+    if violation is None:
+        return blocks
+    left: Final = blocks[violation]
+    right: Final = blocks[violation + 1]
+    merged: Final = _CalibrationBlock(
+        upper_bound=right.upper_bound,
+        successes=left.successes + right.successes,
+        observations=left.observations + right.observations,
+    )
+    return _pool_adjacent_violators((*blocks[:violation], merged, *blocks[violation + 2 :]))
+
+
+def fit_probability_calibration(
+    records: Sequence[CapabilityTrainingRecord], max_bins: int = 10
+) -> tuple[CapabilityCalibrationBin, ...]:
+    if max_bins < 1:
+        raise ValueError("max_bins must be at least 1")
+    if not records:
+        return ()
+    ordered: Final = tuple(sorted(records, key=lambda record: record.raw_p_solve))
+    bin_count: Final = min(max_bins, max(1, int(math.sqrt(len(ordered)))))
+    cutpoints: Final = tuple(
+        sorted(
+            frozenset(
+                ordered[min(len(ordered) - 1, math.ceil(len(ordered) * index / bin_count) - 1)].raw_p_solve
+                for index in range(1, bin_count + 1)
+            )
+        )
+    )
+    chunks: Final = tuple(
+        tuple(
+            record
+            for record in ordered
+            if (index == 0 or record.raw_p_solve > cutpoints[index - 1]) and record.raw_p_solve <= cutpoint
+        )
+        for index, cutpoint in enumerate(cutpoints)
+    )
+    blocks: Final = _pool_adjacent_violators(
+        tuple(
+            _CalibrationBlock(
+                upper_bound=max(record.raw_p_solve for record in chunk),
+                successes=sum(record.success for record in chunk),
+                observations=len(chunk),
+            )
+            for chunk in chunks
+        )
+    )
+    return tuple(
+        CapabilityCalibrationBin(
+            upper_bound=1.0 if index == len(blocks) - 1 else block.upper_bound,
+            probability=block.probability,
+        )
+        for index, block in enumerate(blocks)
+    )
+
+
+def _wilson_interval(successes: float, observations: int) -> tuple[float, float]:
+    if observations == 0:
+        return 0.0, 1.0
+    z: Final = 1.959963984540054
+    mean: Final = successes / observations
+    denominator: Final = 1.0 + z**2 / observations
+    center: Final = (mean + z**2 / (2.0 * observations)) / denominator
+    margin: Final = z * math.sqrt(mean * (1.0 - mean) / observations + z**2 / (4.0 * observations**2)) / denominator
+    return max(0.0, center - margin), min(1.0, center + margin)
+
+
+def _learned_boundary(
+    current: CapabilityRuleBoundary, successes: float, observations: int, target_success: float
+) -> tuple[CapabilityRuleBoundary, float, float]:
+    low, high = _wilson_interval(successes, observations)
+    if observations == 0:
+        return current, low, high
+    if low >= target_success:
+        return "supported", low, high
+    if high < target_success:
+        return "unsupported", low, high
+    return "uncertain", low, high
+
+
+def _train_candidate(
+    candidate: CapabilityRouterCandidate,
+    records: Sequence[CapabilityTrainingRecord],
+    target_success: float,
+) -> tuple[CapabilityRouterCandidate, tuple[CapabilityRuleStatistic, ...]]:
+    candidate_records: Final = tuple(record for record in records if record.model == candidate.model)
+    calibration: Final = fit_probability_calibration(candidate_records)
+    learned: Final = tuple(
+        (
+            rule_id,
+            rule,
+            tuple(record for record in candidate_records if record.primary_rule == rule_id),
+        )
+        for rule_id, rule in indexed_rules(candidate)
+    )
+    boundaries: Final = tuple(
+        _learned_boundary(
+            rule.boundary,
+            sum(record.success for record in rule_records),
+            len(rule_records),
+            target_success,
+        )
+        for _, rule, rule_records in learned
+    )
+    rules: Final = tuple(
+        CapabilityRule(boundary=boundary[0], rule=rule.rule)
+        for (_, rule, _), boundary in zip(learned, boundaries)
+    )
+    statistics: Final = tuple(
+        CapabilityRuleStatistic(
+            model=candidate.model,
+            rule_id=rule_id,
+            observations=len(rule_records),
+            success_rate=(
+                sum(record.success for record in rule_records) / len(rule_records) if rule_records else 0.0
+            ),
+            interval_low=boundary[1],
+            interval_high=boundary[2],
+            learned_boundary=boundary[0],
+        )
+        for (rule_id, _, rule_records), boundary in zip(learned, boundaries)
+    )
+    return candidate.model_copy(update={"rules": rules, "probability_calibration": calibration}), statistics
+
+
+def _effective_boundary(candidate: CapabilityRouterCandidate, primary_rule: str) -> CapabilityBoundary:
+    if not candidate.rules:
+        return "unmatched"
+    boundaries: Final[dict[str, CapabilityBoundary]] = {
+        rule_id: rule.boundary for rule_id, rule in indexed_rules(candidate)
+    }
+    return boundaries.get(primary_rule, "unmatched")
+
+
+def _task_candidates(
+    records: Sequence[CapabilityTrainingRecord],
+    split: Literal["validation", "test"],
+    config: CapabilityRouterConfig,
+    calibrated: bool,
+) -> tuple[tuple[_TaskCandidate, ...], ...]:
+    candidates: Final = {candidate.model: candidate for candidate in config.candidates}
+    selected: Final = sorted(
+        (record for record in records if record.split == split and record.model in candidates),
+        key=lambda record: (record.benchmark, record.task_id, record.model),
+    )
+    grouped_tasks: Final = tuple(
+        tuple(group)
+        for _, group in groupby(selected, key=lambda record: (record.benchmark, record.task_id))
+    )
+    return tuple(
+        tuple(
+            _aggregate_candidate(tuple(model_records), candidates[model], calibrated)
+            for model, model_records in groupby(task_records, key=lambda record: record.model)
+        )
+        for task_records in grouped_tasks
+        if frozenset(record.model for record in task_records) == frozenset(candidates)
+    )
+
+
+def _aggregate_candidate(
+    records: tuple[CapabilityTrainingRecord, ...],
+    candidate: CapabilityRouterCandidate,
+    calibrated: bool,
+) -> _TaskCandidate:
+    raw_probability: Final = sum(record.raw_p_solve for record in records) / len(records)
+    primary_rule: Final = min(Counter(record.primary_rule for record in records).items(), key=lambda item: (-item[1], item[0]))[0]
+    return _TaskCandidate(
+        model=candidate.model,
+        probability=calibrated_probability(candidate, raw_probability) 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),
+    )
+
+
+def _route_metrics(
+    records: Sequence[CapabilityTrainingRecord],
+    split: Literal["validation", "test"],
+    config: CapabilityRouterConfig,
+    quality_weight: float,
+    calibrated: bool,
+) -> CapabilityRouteMetrics:
+    tasks: Final = _task_candidates(records, split, config, calibrated)
+    if not tasks:
+        raise ValueError(f"{split} has no tasks with outcomes for every configured candidate")
+    order: Final = {candidate.model: index for index, candidate in enumerate(config.candidates)}
+    selected: Final = tuple(_select_task_candidate(task, config, order) for task in tasks)
+    return _summarize_routes(tasks, selected, quality_weight)
+
+
+def _summarize_routes(
+    tasks: tuple[tuple[_TaskCandidate, ...], ...],
+    selected: tuple[_TaskCandidate, ...],
+    quality_weight: float,
+) -> CapabilityRouteMetrics:
+    normalized_costs: Final = tuple(_normalized_task_cost(choice, task) for choice, task in zip(selected, tasks))
+    success_rate: Final = sum(choice.success for choice in selected) / len(selected)
+    normalized_cost: Final = sum(normalized_costs) / len(normalized_costs)
+    return CapabilityRouteMetrics(
+        tasks=len(tasks),
+        success_rate=success_rate,
+        mean_cost=sum(choice.cost for choice in selected) / len(selected),
+        normalized_cost=normalized_cost,
+        quality_cost_utility=quality_weight * success_rate + (1.0 - quality_weight) * (1.0 - normalized_cost),
+    )
+
+
+def _always_candidate_metrics(
+    records: Sequence[CapabilityTrainingRecord],
+    config: CapabilityRouterConfig,
+    quality_weight: float,
+) -> Mapping[str, CapabilityRouteMetrics]:
+    tasks: Final = _task_candidates(records, "test", config, True)
+    return {
+        model: _summarize_routes(
+            tasks,
+            tuple(next(candidate for candidate in task if candidate.model == model) for task in tasks),
+            quality_weight,
+        )
+        for model in (candidate.model for candidate in config.candidates)
+    }
+
+
+def _oracle_metrics(
+    records: Sequence[CapabilityTrainingRecord],
+    config: CapabilityRouterConfig,
+    quality_weight: float,
+) -> CapabilityRouteMetrics:
+    tasks: Final = _task_candidates(records, "test", config, True)
+    selected: Final = tuple(
+        max(
+            task,
+            key=lambda candidate: (
+                quality_weight * candidate.success
+                + (1.0 - quality_weight) * (1.0 - _normalized_task_cost(candidate, task)),
+                -candidate.cost,
+            ),
+        )
+        for task in tasks
+    )
+    return _summarize_routes(tasks, selected, quality_weight)
+
+
+def _select_task_candidate(
+    task: tuple[_TaskCandidate, ...], config: CapabilityRouterConfig, order: dict[str, int]
+) -> _TaskCandidate:
+    qualified: Final = tuple(
+        candidate
+        for candidate in task
+        if candidate.probability
+        > round(
+            config.probability_threshold
+            + BOUNDARY_THRESHOLD_STEPS[candidate.boundary] * config.threshold_step,
+            9,
+        )
+    )
+    fallback: Final = next(candidate for candidate in task if candidate.model == config.fallback_model)
+    return min(qualified, key=lambda candidate: (candidate.cost, order[candidate.model])) if qualified else fallback
+
+
+def _normalized_task_cost(selected: _TaskCandidate, task: tuple[_TaskCandidate, ...]) -> float:
+    low: Final = min(candidate.cost for candidate in task)
+    high: Final = max(candidate.cost for candidate in task)
+    return 0.0 if high == low else (selected.cost - low) / (high - low)
+
+
+def _probability_metrics(
+    records: Sequence[CapabilityTrainingRecord],
+    config: CapabilityRouterConfig,
+    calibrated: bool,
+) -> CapabilityProbabilityMetrics:
+    candidates: Final = {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
+        for record in rows
+    )
+    outcomes: Final = tuple(record.success for record in rows)
+    return CapabilityProbabilityMetrics(
+        observations=len(rows),
+        brier=sum((prediction - outcome) ** 2 for prediction, outcome in zip(predictions, outcomes)) / len(rows),
+        log_loss=-sum(
+            outcome * math.log(max(1e-9, prediction))
+            + (1.0 - outcome) * math.log(max(1e-9, 1.0 - prediction))
+            for prediction, outcome in zip(predictions, outcomes)
+        )
+        / len(rows),
+        ece=sum(
+            len(bucket)
+            / len(rows)
+            * abs(sum(item[0] for item in bucket) / len(bucket) - sum(item[1] for item in bucket) / len(bucket))
+            for index in range(10)
+            for bucket in (
+                tuple(
+                    (prediction, outcome)
+                    for prediction, outcome in zip(predictions, outcomes)
+                    if min(9, int(prediction * 10)) == index
+                ),
+            )
+            if bucket
+        ),
+    )
+
+
+def train_capability_artifact(
+    records: Sequence[CapabilityTrainingRecord],
+    config: CapabilityRouterConfig,
+    quality_weight: float = 0.7,
+) -> CapabilityTrainingResult:
+    if not 0.0 <= quality_weight <= 1.0:
+        raise ValueError("quality_weight must be between 0 and 1")
+    required_splits: Final = frozenset(record.split for record in records)
+    if required_splits != frozenset(("train", "validation", "test")):
+        raise ValueError("records must explicitly contain train, validation, and test splits")
+    task_splits: Final = tuple(
+        frozenset(record.split for record in task_records)
+        for _, grouped in groupby(
+            sorted(records, key=lambda record: (record.benchmark, record.task_id)),
+            key=lambda record: (record.benchmark, record.task_id),
+        )
+        for task_records in (tuple(grouped),)
+    )
+    if any(len(splits) != 1 for splits in task_splits):
+        raise ValueError("a benchmark task must not cross splits")
+    configured_models: Final = frozenset(candidate.model for candidate in config.candidates)
+    if frozenset(record.model for record in records) != configured_models:
+        raise ValueError("record models must exactly match configured candidates")
+    training: Final = tuple(record for record in records if record.split == "train")
+    trained_rows: Final = tuple(
+        _train_candidate(candidate, training, config.probability_threshold) for candidate in config.candidates
+    )
+    trained_candidates: Final = tuple(candidate for candidate, _ in trained_rows)
+    rule_statistics: Final = tuple(statistic for _, statistics in trained_rows for statistic in statistics)
+    calibrated_config: Final = config.model_copy(update={"candidates": trained_candidates})
+    candidates: Final = tuple(
+        calibrated_config.model_copy(update={"probability_threshold": threshold, "threshold_step": step})
+        for threshold in _THRESHOLDS
+        for step in _THRESHOLD_STEPS
+        if threshold + 2.0 * step <= 1.0
+    )
+    scored: Final = tuple(
+        (_route_metrics(records, "validation", candidate, quality_weight, True), candidate)
+        for candidate in candidates
+    )
+    validation, trained_config = max(
+        scored,
+        key=lambda item: (
+            item[0].quality_cost_utility,
+            -item[0].mean_cost,
+            item[0].success_rate,
+        ),
+    )
+    datasets: Final = tuple(sorted(frozenset(record.benchmark for record in records)))
+    artifact: Final = CapabilityTrainingArtifact(
+        config=trained_config,
+        rule_statistics=rule_statistics,
+        datasets=datasets,
+        records=len(records),
+        split_counts={split: sum(record.split == split for record in records) for split in sorted(required_splits)},
+        records_sha256=hashlib.sha256(
+            "\n".join(
+                record.model_dump_json()
+                for record in sorted(
+                    records,
+                    key=lambda record: (
+                        record.benchmark,
+                        record.task_id,
+                        record.split,
+                        record.model,
+                        record.primary_rule,
+                        record.raw_p_solve,
+                        record.success,
+                        record.estimated_cost,
+                    ),
+                )
+            ).encode()
+        ).hexdigest(),
+        split_contract="split is explicit; task_id must not cross splits",
+    )
+    return CapabilityTrainingResult(
+        artifact=artifact,
+        report=CapabilityTrainingReport(
+            objective=(
+                f"{quality_weight:g} * observed success + {1.0 - quality_weight:g} * normalized cost score"
+            ),
+            validation=validation,
+            test=_route_metrics(records, "test", trained_config, quality_weight, True),
+            test_untrained=_route_metrics(records, "test", config, quality_weight, False),
+            test_always_candidates=_always_candidate_metrics(records, trained_config, quality_weight),
+            test_oracle=_oracle_metrics(records, trained_config, quality_weight),
+            test_threshold_sweep=tuple(
+                CapabilityThresholdPoint(
+                    probability_threshold=candidate.probability_threshold,
+                    threshold_step=candidate.threshold_step,
+                    metrics=_route_metrics(records, "test", candidate, quality_weight, True),
+                )
+                for candidate in candidates
+            ),
+            test_probability_raw=_probability_metrics(records, config, False),
+            test_probability_calibrated=_probability_metrics(records, trained_config, True),
+        ),
+    )
+
+
+def _read_records(path: Path) -> tuple[CapabilityTrainingRecord, ...]:
+    with path.open() as record_file:
+        return tuple(CapabilityTrainingRecord.model_validate_json(line) for line in record_file if line.strip())
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser: Final = argparse.ArgumentParser(description="Train and benchmark capability-router cards")
+    parser.add_argument("records", type=Path)
+    parser.add_argument("--config", type=Path, required=True)
+    parser.add_argument("--artifact-output", type=Path, required=True)
+    parser.add_argument("--quality-weight", type=float, default=0.7)
+    return parser
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    args: Final = _Arguments.model_validate(vars(_parser().parse_args(argv)))
+    records: Final = _read_records(args.records)
+    config: Final = CapabilityRouterConfig.model_validate(json.loads(args.config.read_text()))
+    result: Final = train_capability_artifact(records, config, args.quality_weight)
+    args.artifact_output.write_text(result.artifact.model_dump_json(indent=2) + "\n")
+    print(result.report.model_dump_json(indent=2))  # noqa: T201  # CLI emits benchmark JSON
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

litellm/types/utils.py

@@ -2936,6 +2936,9 @@ class StandardLoggingRoutingDecision(TypedDict, total=False):
     tier_litellm_params: Mapping[str, object]  # writable-ok: Pydantic warns on ReadOnly TypedDict fields
     probability_threshold: float  # writable-ok: Pydantic warns on ReadOnly TypedDict fields
     candidate_probabilities: Mapping[str, float]  # writable-ok: Pydantic warns on ReadOnly TypedDict fields
+    raw_candidate_probabilities: Mapping[str, float]  # writable-ok: Pydantic warns on ReadOnly TypedDict fields
+    candidate_rules: Mapping[str, str]  # writable-ok: Pydantic warns on ReadOnly TypedDict fields
+    candidate_boundaries: Mapping[str, str]  # writable-ok: Pydantic warns on ReadOnly TypedDict fields
     candidate_costs: Mapping[str, float]  # writable-ok: Pydantic warns on ReadOnly TypedDict fields
     qualified_models: Sequence[str]  # writable-ok: Pydantic warns on ReadOnly TypedDict fields
     fallback_reason: str | None  # writable-ok: Pydantic warns on ReadOnly TypedDict fields
@@ -2970,6 +2973,9 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset(
         "tier_litellm_params",
         "probability_threshold",
         "candidate_probabilities",
+        "raw_candidate_probabilities",
+        "candidate_rules",
+        "candidate_boundaries",
         "candidate_costs",
         "qualified_models",
         "fallback_reason",

tests/test_litellm/router_strategy/capability_router/test_capability_router.py

@@ -41,6 +41,14 @@ def test_config_requires_unique_candidates_and_candidate_fallback() -> None:
     with pytest.raises(ValidationError, match="one of the candidate"):
         CapabilityRouterConfig.model_validate(missing_fallback)
 
+    invalid_calibration = config()
+    invalid_calibration["candidates"][0]["probability_calibration"] = [
+        {"upper_bound": 0.5, "probability": 0.8},
+        {"upper_bound": 1.0, "probability": 0.4},
+    ]
+    with pytest.raises(ValidationError, match="nondecreasing"):
+        CapabilityRouterConfig.model_validate(invalid_calibration)
+
 
 def test_policy_selects_cheapest_model_above_global_threshold() -> None:
     parsed = CapabilityRouterConfig.model_validate(config())
@@ -178,9 +186,36 @@ async def test_same_user_turn_reuses_cached_decision() -> None:
     assert first.model == second.model == "small"
     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
+    assert first.routing_decision["raw_candidate_probabilities"] == {"small": 0.9, "frontier": 0.95}
+    assert first.routing_decision["candidate_rules"] == {"small": "none", "frontier": "none"}
     strategy._new_decision.assert_awaited_once()
 
 
+def test_policy_qualifies_on_calibrated_probability_and_keeps_raw_forecast() -> None:
+    configured = config()
+    configured["candidates"][0]["probability_calibration"] = [
+        {"upper_bound": 0.8, "probability": 0.4},
+        {"upper_bound": 1.0, "probability": 0.9},
+    ]
+    parsed = CapabilityRouterConfig.model_validate(configured)
+    verdict = CapabilityClassifierVerdict.model_validate(
+        {
+            "candidates": [
+                {"model": "small", "capability_boundary": "supported", "p_solve": 0.78, "reason": "raw optimism"},
+                {"model": "frontier", "capability_boundary": "supported", "p_solve": 0.95, "reason": "covered"},
+            ]
+        }
+    )
+
+    decision = select_capability_model(parsed, verdict, {"small": 0.01, "frontier": 0.05})
+
+    small = decision.candidates[0]
+    assert small.raw_p_solve == 0.78
+    assert small.p_solve == 0.4
+    assert small.qualified is False
+    assert decision.selected_model == "frontier"
+
+
 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

@@ -0,0 +1,120 @@
+import json
+from pathlib import Path
+
+import pytest
+
+from litellm.router_strategy.capability_router.config import CapabilityRouterConfig
+from litellm.router_strategy.capability_router.training import (
+    CapabilityTrainingRecord,
+    main,
+    train_capability_artifact,
+)
+
+
+def training_config() -> CapabilityRouterConfig:
+    return CapabilityRouterConfig.model_validate(
+        {
+            "candidates": [
+                {
+                    "model": "small",
+                    "description": "Efficient model with bounded-task and open-ended-task rules",
+                    "rules": [
+                        {"boundary": "uncertain", "rule": "The task has an executable correctness check"},
+                        {"boundary": "uncertain", "rule": "The task requires resolving hidden behavior"},
+                    ],
+                },
+                {
+                    "model": "strong",
+                    "description": "Capable fallback model",
+                    "rules": [
+                        {"boundary": "supported", "rule": "The task has an executable correctness check"},
+                        {"boundary": "supported", "rule": "The task requires resolving hidden behavior"},
+                    ],
+                },
+            ],
+            "classifier": {"model": "judge"},
+            "probability_threshold": 0.5,
+            "threshold_step": 0.0,
+            "fallback_model": "strong",
+        }
+    )
+
+
+def records() -> tuple[CapabilityTrainingRecord, ...]:
+    split_sizes = {"train": 20, "validation": 4, "test": 4}
+    return tuple(
+        CapabilityTrainingRecord(
+            benchmark="agent-bench",
+            task_id=f"{split}-{difficulty}-{index}",
+            split=split,
+            model=model,
+            primary_rule="R1" if difficulty == "bounded" else "R2",
+            raw_p_solve=(0.9 if difficulty == "bounded" else 0.6) if model == "small" else 0.9,
+            success=1.0 if model == "strong" or difficulty == "bounded" else 0.0,
+            estimated_cost=1.0 if model == "small" else 10.0,
+        )
+        for split, count in split_sizes.items()
+        for difficulty in ("bounded", "hidden")
+        for index in range(count)
+        for model in ("small", "strong")
+    )
+
+
+def test_training_learns_rule_boundaries_and_improves_held_out_routing() -> None:
+    result = train_capability_artifact(records(), training_config())
+    small = result.artifact.config.candidates[0]
+
+    assert [rule.boundary for rule in small.rules] == ["supported", "unsupported"]
+    assert small.probability_calibration[-1].upper_bound == 1.0
+    assert result.artifact.records == len(records())
+    assert len(result.artifact.records_sha256) == 64
+    assert result.report.test.success_rate == 1.0
+    assert result.report.test_untrained.success_rate == 0.5
+    assert result.report.test.quality_cost_utility > result.report.test_untrained.quality_cost_utility
+    assert result.report.test_probability_calibrated.brier < result.report.test_probability_raw.brier
+    assert result.report.test_always_candidates["small"].success_rate == 0.5
+    assert result.report.test_always_candidates["strong"].success_rate == 1.0
+    assert result.report.test_oracle.quality_cost_utility >= result.report.test.quality_cost_utility
+    assert result.report.test_threshold_sweep
+
+
+def test_training_requires_explicit_splits_and_exact_candidate_models() -> None:
+    only_training = tuple(record for record in records() if record.split == "train")
+    with pytest.raises(ValueError, match="train, validation, and test"):
+        train_capability_artifact(only_training, training_config())
+
+    without_strong = tuple(record for record in records() if record.model == "small")
+    with pytest.raises(ValueError, match="exactly match"):
+        train_capability_artifact(without_strong, training_config())
+
+
+def test_training_rejects_a_task_that_crosses_splits() -> None:
+    source = records()
+    crossed = (*source, source[0].model_copy(update={"split": "test"}))
+
+    with pytest.raises(ValueError, match="must not cross splits"):
+        train_capability_artifact(crossed, training_config())
+
+
+def test_training_cli_writes_a_ready_to_use_artifact(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None:
+    records_path = tmp_path / "outcomes.jsonl"
+    config_path = tmp_path / "config.json"
+    artifact_path = tmp_path / "artifact.json"
+    records_path.write_text("".join(record.model_dump_json() + "\n" for record in records()))
+    config_path.write_text(training_config().model_dump_json())
+
+    exit_code = main(
+        (
+            str(records_path),
+            "--config",
+            str(config_path),
+            "--artifact-output",
+            str(artifact_path),
+        )
+    )
+
+    artifact = json.loads(artifact_path.read_text())
+    report = json.loads(capsys.readouterr().out)
+    assert exit_code == 0
+    assert artifact["config"]["candidates"][0]["probability_calibration"]
+    assert report["test"]["success_rate"] == 1.0

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -24582,6 +24582,16 @@ export interface components {
              */
             status: "cancelled";
         };
+        /**
+         * CapabilityCalibrationBin
+         * @description One monotonic post-hoc calibration bucket learned from end-to-end outcomes.
+         */
+        CapabilityCalibrationBin: {
+            /** Probability */
+            probability: number;
+            /** Upper Bound */
+            upper_bound: number;
+        };
         /**
          * CapabilityClassifierConfig
          * @description The model used to estimate each candidate's probability of success.
@@ -24614,6 +24624,11 @@ export interface components {
             description: string;
             /** Model */
             model: string;
+            /**
+             * Probability Calibration
+             * @default []
+             */
+            probability_calibration: components["schemas"]["CapabilityCalibrationBin"][];
             /**
              * Rules
              * @default []
@@ -35849,6 +35864,10 @@ export interface components {
         StandardLoggingRoutingDecision: {
             /** Cached */
             cached?: boolean;
+            /** Candidate Boundaries */
+            candidate_boundaries?: {
+                [key: string]: string;
+            };
             /** Candidate Costs */
             candidate_costs?: {
                 [key: string]: number;
@@ -35857,6 +35876,10 @@ export interface components {
             candidate_probabilities?: {
                 [key: string]: number;
             };
+            /** Candidate Rules */
+            candidate_rules?: {
+                [key: string]: string;
+            };
             /**
              * Cause
              * @enum {string}
@@ -35884,6 +35907,10 @@ export interface components {
             probability_threshold?: number;
             /** Qualified Models */
             qualified_models?: string[];
+            /** Raw Candidate Probabilities */
+            raw_candidate_probabilities?: {
+                [key: string]: number;
+            };
             /** Reasoning Override Min Score */
             reasoning_override_min_score?: number;
             /** Request Type */