feat(router): reject capability task family leakage

litellm/router_strategy/capability_router/README.md

@@ -14,12 +14,12 @@ Capability-card training uses outcomes from the real task harness. It does not r
 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}
+{"benchmark":"swe-bench-verified","task_id":"django__django-12345","task_family":"django/django","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
+Every benchmark task must have an explicit `train`, `validation`, or `test` split and outcomes for every configured candidate. A task cannot cross splits. Set `task_family` to a repository, domain, task generator, or another shared origin when related tasks could leak. The trainer rejects a family found across 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:
 

litellm/router_strategy/capability_router/policy.py

@@ -34,11 +34,7 @@ def effective_boundary(candidate: CapabilityRouterCandidate, score: CapabilityCa
 
 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
-        ),
+        (bucket.probability for bucket in candidate.probability_calibration if raw_probability <= bucket.upper_bound),
         raw_probability,
     )
 

litellm/router_strategy/capability_router/training.py

@@ -31,6 +31,7 @@ _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)
+    task_family: str | None = Field(default=None, min_length=1)
     split: Literal["train", "validation", "test"]
     model: str = Field(min_length=1)
     primary_rule: str = Field(default="none", min_length=1)
@@ -472,6 +473,19 @@ def train_capability_artifact(
     )
     if any(len(splits) != 1 for splits in task_splits):
         raise ValueError("a benchmark task must not cross splits")
+    family_splits: Final = tuple(
+        frozenset(record.split for record in family_records)
+        for _, grouped in groupby(
+            sorted(
+                (record for record in records if record.task_family is not None),
+                key=lambda record: (record.benchmark, record.task_family),
+            ),
+            key=lambda record: (record.benchmark, record.task_family),
+        )
+        for family_records in (tuple(grouped),)
+    )
+    if any(len(splits) != 1 for splits in family_splits):
+        raise ValueError("a benchmark task family 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")
@@ -523,6 +537,7 @@ def train_capability_artifact(
                     key=lambda record: (
                         record.benchmark,
                         record.task_id,
+                        record.task_family or "",
                         record.split,
                         record.model,
                         record.primary_rule,
@@ -533,7 +548,7 @@ def train_capability_artifact(
                 )
             ).encode()
         ).hexdigest(),
-        split_contract="split is explicit; task_id must not cross splits",
+        split_contract="split is explicit; task_id and task_family must not cross splits",
     )
     return CapabilityTrainingResult(
         artifact=artifact,

tests/test_litellm/router_strategy/capability_router/test_training.py

@@ -96,6 +96,18 @@ def test_training_rejects_a_task_that_crosses_splits() -> None:
         train_capability_artifact(crossed, training_config())
 
 
+def test_training_rejects_a_task_family_that_crosses_splits() -> None:
+    source = records()
+    crossed = (
+        source[0].model_copy(update={"task_family": "shared-repository"}),
+        source[1].model_copy(update={"task_id": "different-task", "task_family": "shared-repository", "split": "test"}),
+        *source[2:],
+    )
+
+    with pytest.raises(ValueError, match="task family 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"