From ac141a6fe4b28118fda8f3aa020d84b32547feb9 Mon Sep 17 00:00:00 2001 From: Tin Date: Wed, 2 Sep 2026 01:19:25 -0700 Subject: [PATCH] feat(router): reject capability task family leakage --- .../router_strategy/capability_router/README.md | 4 ++-- .../router_strategy/capability_router/policy.py | 6 +----- .../capability_router/training.py | 17 ++++++++++++++++- .../capability_router/test_training.py | 12 ++++++++++++ 4 files changed, 31 insertions(+), 8 deletions(-) diff --git a/litellm/router_strategy/capability_router/README.md b/litellm/router_strategy/capability_router/README.md index e2207d0639d..7da2aa0885a 100644 --- a/litellm/router_strategy/capability_router/README.md +++ b/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: diff --git a/litellm/router_strategy/capability_router/policy.py b/litellm/router_strategy/capability_router/policy.py index 526fbbcc2cc..98d954da6b8 100644 --- a/litellm/router_strategy/capability_router/policy.py +++ b/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, ) diff --git a/litellm/router_strategy/capability_router/training.py b/litellm/router_strategy/capability_router/training.py index 7dd9dc5330e..61d2d261b64 100644 --- a/litellm/router_strategy/capability_router/training.py +++ b/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, 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 6f440d5055f..25e3512a234 100644 --- a/tests/test_litellm/router_strategy/capability_router/test_training.py +++ b/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"