From d619a4b2bb337e1d2c0ec387eadfb3018766ff38 Mon Sep 17 00:00:00 2001
From: Tin <tin@berri.ai>
Date: Thu, 3 Sep 2026 18:56:23 -0700
Subject: [PATCH] fix(router): make calibration pooling iterative

---
 .../capability_router/training.py             | 13 +++++++++--
 .../capability_router/test_training.py        | 22 +++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/litellm/router_strategy/capability_router/training.py b/litellm/router_strategy/capability_router/training.py
index 285e76da8ba..180fb26d25a 100644
--- a/litellm/router_strategy/capability_router/training.py
+++ b/litellm/router_strategy/capability_router/training.py
@@ -6,6 +6,7 @@ import json
 import math
 from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import dataclass
+from functools import reduce
 from itertools import groupby
 from pathlib import Path
 from types import MappingProxyType
@@ -140,7 +141,7 @@ class _TaskCandidate:
     boundary: CapabilityBoundary
 
 
-def _pool_adjacent_violators(blocks: tuple[_CalibrationBlock, ...]) -> tuple[_CalibrationBlock, ...]:
+def _merge_first_calibration_violation(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,
@@ -154,7 +155,15 @@ def _pool_adjacent_violators(blocks: tuple[_CalibrationBlock, ...]) -> tuple[_Ca
         successes=left.successes + right.successes,
         observations=left.observations + right.observations,
     )
-    return _pool_adjacent_violators((*blocks[:violation], merged, *blocks[violation + 2 :]))
+    return (*blocks[:violation], merged, *blocks[violation + 2 :])
+
+
+def _pool_adjacent_violators(blocks: tuple[_CalibrationBlock, ...]) -> tuple[_CalibrationBlock, ...]:
+    return reduce(
+        lambda pooled, _: _merge_first_calibration_violation(pooled),
+        range(len(blocks)),
+        blocks,
+    )
 
 
 def fit_probability_calibration(
diff --git a/tests/test_litellm/router_strategy/capability_router/test_training.py b/tests/test_litellm/router_strategy/capability_router/test_training.py
index 569a198a9d9..64186623b10 100644
--- a/tests/test_litellm/router_strategy/capability_router/test_training.py
+++ b/tests/test_litellm/router_strategy/capability_router/test_training.py
@@ -6,6 +6,7 @@ import pytest
 from litellm.router_strategy.capability_router.config import CapabilityRouterConfig
 from litellm.router_strategy.capability_router.training import (
     CapabilityTrainingRecord,
+    fit_probability_calibration,
     main,
     train_capability_artifact,
 )
@@ -79,6 +80,27 @@ def test_training_learns_rule_boundaries_and_improves_held_out_routing() -> None
     assert result.report.test_threshold_sweep
 
 
+def test_probability_calibration_pools_new_adjacent_violations() -> None:
+    calibration_records = tuple(
+        CapabilityTrainingRecord(
+            benchmark="agent-bench",
+            task_id=f"task-{index}",
+            split="train",
+            model="small",
+            raw_p_solve=probability,
+            success=success,
+            estimated_cost=1.0,
+        )
+        for index, (probability, success) in enumerate(((0.1, 1.0),) * 3 + ((0.2, 0.0),) * 3 + ((0.3, 0.0),) * 3)
+    )
+
+    calibration = fit_probability_calibration(calibration_records, max_bins=3)
+
+    assert len(calibration) == 1
+    assert calibration[0].upper_bound == 1.0
+    assert calibration[0].probability == pytest.approx(4.0 / 11.0)
+
+
 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"):