chore(cookbook): add standalone capability routing pilot POC

Co-Authored-By: bot_apk <apk@cognition.ai>
This commit is contained in:
Devin AI 2026-10-20 23:28:22 +00:00
parent f3e1307ed9
commit a0baa80673
17 changed files with 2326 additions and 0 deletions

View file

@ -0,0 +1,11 @@
# Capability-aware routing pilot (standalone research POC)
Five frozen SWE-bench dev tasks (`tasks.json`) solved by mini-swe-agent 2.4.6 with Claude Haiku 4.5, Sonnet 5 and Opus 5 through a LiteLLM gateway, graded only by the official swebench 5.0.2 harness. This is not production code and nothing under `litellm/` imports it
`poc/guarded_model.py` wraps the mini-swe-agent LiteLLM model and refuses a call before dispatch when spent plus a worst-case reservation would exceed the per-run budget. Actual cost comes from the gateway response cost header. `poc/run_pilot.py` runs the 15-run matrix with one setup-only retry for infrastructure failures and appends one record per attempt to `runs.jsonl`. `poc/label.py` writes offline GPT-5.4-mini labels and runtime Jev labels from the problem statement only. `poc/policy.py` builds the capability map and the conservative router, and `poc/analyze.py` writes `capability_map.json` and the report
The router picks a cheaper model only when, on the shared graded task IDs, a one-sided Hoeffding lower bound on the paired success difference is at least minus the tolerance and the paired mean cost is lower. With five tasks it always falls back to the Opus reference, by design
Three evaluator images failed gold validation with an unchanged grader (NumPy 2 in both pvlib images, pytest 8 in the pydicom image). `overlays/` holds the version pins applied locally under the original tags. Gold then resolved 5/5 and a no-op patch 0/5. `images_pinned.txt` lists the digests used and `images_pinned_upstream.txt` the upstream ones
Run tests with `python -m pytest -q tests` from this directory with mini-swe-agent, litellm, scipy and pyyaml installed

View file

@ -0,0 +1,9 @@
swebench/sweb.eval.x86_64.pydicom_1776_pydicom-1031:latest sha256:dd35fa4aecc830e588ae5966b3f3cfe91b3b34a6057c483847f82a44e4e77ac5
swebench/sweb.eval.x86_64.pvlib_1776_pvlib-python-1165:latest sha256:e4134171da28404a3d71ef75f1866b776def32c31b986aab528c6f5716b27423
swebench/sweb.eval.x86_64.pvlib_1776_pvlib-python-1395:latest sha256:2eb9ad59249fc77761ffab8158f8fe2b1f44111f4972f804ca09991e49ae1197
swebench/sweb.eval.x86_64.pylint-dev_1776_astroid-1614:latest sha256:7206ea531b6c8b02f7a8e1760d732e4651f95307781f12f55e4b88e982610159
swebench/sweb.eval.x86_64.sqlfluff_1776_sqlfluff-884:latest sha256:ca28dab1d331c1c7e4908784576a7e1f3b11387dbaa4a8db466ce42a0eda6bec
# overlays (local, built from upstream): pvlib-1395,pvlib-1165 +numpy==1.26.4; pydicom-1031 +pytest==7.4.4; originals tagged routing_poc/orig.*
routing_poc/orig.pydicom_1776_pydicom-1031:latest sha256:0b962e881760f4b8bf3cdbfc695624c90906198ccf1e2410cd639f6adc6e340b
routing_poc/orig.pvlib_1776_pvlib-python-1165:latest sha256:71e033964894077384c0c27a04ae0271e28eaad35602aa7570d40c683f1967d1
routing_poc/orig.pvlib_1776_pvlib-python-1395:latest sha256:c7b51f4978926d54a788c54ac4d96e6c4c4e2cd2ee62736b30e91c9fd0891f10

View file

@ -0,0 +1,5 @@
swebench/sweb.eval.x86_64.pydicom_1776_pydicom-1031:latest sha256:0b962e881760f4b8bf3cdbfc695624c90906198ccf1e2410cd639f6adc6e340b 5.49GB
swebench/sweb.eval.x86_64.pylint-dev_1776_astroid-1614:latest sha256:7206ea531b6c8b02f7a8e1760d732e4651f95307781f12f55e4b88e982610159 3.55GB
swebench/sweb.eval.x86_64.pvlib_1776_pvlib-python-1165:latest sha256:71e033964894077384c0c27a04ae0271e28eaad35602aa7570d40c683f1967d1 5.96GB
swebench/sweb.eval.x86_64.pvlib_1776_pvlib-python-1395:latest sha256:c7b51f4978926d54a788c54ac4d96e6c4c4e2cd2ee62736b30e91c9fd0891f10 5.89GB
swebench/sweb.eval.x86_64.sqlfluff_1776_sqlfluff-884:latest sha256:ca28dab1d331c1c7e4908784576a7e1f3b11387dbaa4a8db466ce42a0eda6bec 3.74GB

View file

@ -0,0 +1,2 @@
FROM routing_poc/orig.pvlib_1776_pvlib-python-1165:latest
RUN /opt/miniconda3/envs/testbed/bin/pip install --no-cache-dir "numpy==1.26.4" && /opt/miniconda3/envs/testbed/bin/pip check || true

View file

@ -0,0 +1,2 @@
FROM routing_poc/orig.pvlib_1776_pvlib-python-1395:latest
RUN /opt/miniconda3/envs/testbed/bin/pip install --no-cache-dir "numpy==1.26.4" && /opt/miniconda3/envs/testbed/bin/pip check || true

View file

@ -0,0 +1,2 @@
FROM routing_poc/orig.pydicom_1776_pydicom-1031:latest
RUN /opt/miniconda3/envs/testbed/bin/pip install --no-cache-dir "pytest==7.4.4" && /opt/miniconda3/envs/testbed/bin/pip check || true

View file

@ -0,0 +1,125 @@
aiohappyeyeballs==2.7.1
aiohttp==3.14.5
aiosignal==1.4.0
annotated-doc==0.0.5
annotated-types==0.8.0
anyio==4.15.1
attrs==26.1.0
beautifulsoup4==4.15.0
boto3==1.43.116
botocore==1.43.116
cbor2==6.1.5
certifi==2026.10.19
cffi==2.1.1
cfgv==3.5.0
chardet==7.7.0
charset-normalizer==3.5.2
click==8.5.0
cryptography==50.0.2
datasets==5.1.0
dill==0.4.1
distlib==0.4.4
distro==1.9.0
docker==7.2.0
fastcore==1.14.5
fastuuid==0.14.0
filelock==3.32.7
frozenlist==1.8.0
fsspec==2026.7.0
ghapi==1.1.1
gitdb==4.0.12
gitpython==3.2.1
grpclib==0.4.9
h11==0.16.0
h2==4.4.1
hf-xet==1.7.0
hpack==4.2.0
httpcore==1.0.9
httpcore2==2.15.1
httpx==0.28.1
httpx2==2.15.1
huggingface-hub==1.33.0
hyperframe==6.1.0
identify==2.6.20
idna==3.20
importlib-metadata==8.9.0
iniconfig==2.3.1
jinja2==3.1.6
jiter==0.17.0
jmespath==1.1.0
jsonschema==4.26.0
jsonschema-specifications==2025.9.1
linkify-it-py==2.2.0
litellm==1.106.1
markdown-it-py==4.2.0
markupsafe==3.0.4
mcp==2.3.0
mcp-types==2.3.0
mdit-py-plugins==0.6.1
mdurl==0.1.2
mini-swe-agent==2.4.6
modal==1.6.2
multidict==7.1.0
multiprocess==0.70.19
nodeenv==1.11.0
numpy==2.5.4
openai==2.54.0
opentelemetry-api==1.45.1
packaging==26.3
pandas==3.0.6
platformdirs==4.13.5
pluggy==1.7.0
pre-commit==4.7.0
prompt-toolkit==3.0.53
propcache==0.5.4
protobuf==7.36.2
pyarrow==26.0.0
pycparser==3.11
pydantic==2.14.1
pydantic-core==2.50.1
pydantic-settings==2.16.1
pygments==2.21.0
pyjwt==2.15.1
pytest==9.1.1
python-dateutil==2.9.0.post0
python-discovery==1.6.4
python-dotenv==1.2.4
python-multipart==0.0.32
pyyaml==6.0.3
referencing==0.37.0
regex==2026.9.29
requests==2.34.2
rich==15.0.0
rpds-py==2026.9.1
s3transfer==0.19.3
scipy==1.18.1
shellingham==1.5.4
six==1.17.0
smmap==5.0.3
sniffio==1.3.1
soupsieve==3.0.3
sse-starlette==3.5.0
starlette==1.9.0
swebench==5.0.2
synchronicity==0.12.6
tenacity==9.2.1
textual==8.2.8
tiktoken==0.14.0
tokenizers==0.23.3
toml==0.10.2
tqdm==4.70.1
truststore==0.10.5
typer==0.27.3
types-certifi==2021.10.8.3
types-toml==0.10.8.20260518
typing-extensions==4.16.0
typing-inspection==0.4.4
unidiff==1.0.1
urllib3==2.8.0
uvicorn==0.54.0
virtualenv==21.14.11
watchfiles==1.3.0
wcwidth==0.9.3
xxhash==4.0.2
yarl==1.25.1
zipp==4.1.1

View file

@ -0,0 +1,419 @@
import csv
import json
import math
from pathlib import Path
from poc.policy import (
PolicyConfig,
beta_cell,
build_capability_map,
choose_model,
load_jsonl,
paired_comparison,
parse_category,
task_outcomes,
)
ROOT = Path(__file__).resolve().parents[1]
MODEL_ORDER = ("haiku-4.5", "sonnet-5", "opus-5")
PREFLIGHT_PATH = "preflight_probes.json"
def _json(path: Path, fallback: object) -> object:
return json.loads(path.read_text()) if path.exists() else fallback
def _cost(value: object) -> float:
if isinstance(value, bool):
return 0.0
if isinstance(value, (int, float)):
return float(value) if math.isfinite(value) else 0.0
if isinstance(value, str):
try:
number = float(value)
except ValueError:
return 0.0
return number if math.isfinite(number) else 0.0
return 0.0
def _labels(path: Path) -> dict[str, dict[str, object]]:
raw = _json(path / "labels.json", [])
return {
row["task_id"]: row
for row in raw
if isinstance(row, dict) and isinstance(row.get("task_id"), str)
}
def _task_ids(path: Path) -> list[str]:
payload = _json(path / "tasks.json", {"tasks": []})
return [task["instance_id"] for task in payload["tasks"]]
def _preflight_cost(path: Path) -> float:
probes = _json(path / PREFLIGHT_PATH, [])
return sum(_cost(probe.get("cost")) for probe in probes if isinstance(probe, dict))
def _aux_costs(path: Path) -> list[dict[str, object]]:
cost_path = path / "costs_aux.jsonl"
return load_jsonl(str(cost_path)) if cost_path.exists() else []
def _model_outcome(outcomes: dict[tuple[str, str], dict[str, object]], task_id: str, model: str) -> dict[str, object]:
return outcomes.get((task_id, model), {"graded": False, "success": None, "cost_usd": 0.0, "attempts": 0})
def _last_exit_status(
records: list[dict[str, object]], task_id: str, model: str
) -> str:
attempts = [row for row in records if row["task_id"] == task_id and row["model"] == model]
return " | ".join(str(row["exit_status"]) for row in attempts if row.get("exit_status")) or "not_run"
def _oracle(outcomes: dict[tuple[str, str], dict[str, object]], task_id: str) -> str:
successful = [
model for model in MODEL_ORDER if _model_outcome(outcomes, task_id, model)["success"] is True
]
return min(successful, key=lambda model: (_model_outcome(outcomes, task_id, model)["cost_usd"], MODEL_ORDER.index(model))) if successful else "none"
def _fixed_model_summaries(
task_ids: list[str],
outcomes: dict[tuple[str, str], dict[str, object]],
records: list[dict[str, object]],
cfg: PolicyConfig,
) -> dict[str, dict[str, object]]:
return {
model: _fixed_model_summary(task_ids, model, outcomes, records, cfg)
for model in MODEL_ORDER
}
def _fixed_model_summary(
task_ids: list[str],
model: str,
outcomes: dict[tuple[str, str], dict[str, object]],
records: list[dict[str, object]],
cfg: PolicyConfig,
) -> dict[str, object]:
model_outcomes = [_model_outcome(outcomes, task_id, model) for task_id in task_ids]
solved = sum(outcome["success"] is True for outcome in model_outcomes)
total_cost = sum(float(outcome["cost_usd"]) for outcome in model_outcomes)
attempted_tasks = sum(outcome["attempts"] > 0 for outcome in model_outcomes)
cell = beta_cell(solved, len(task_ids), cfg.credible_mass, total_cost / attempted_tasks if attempted_tasks else None)
return {
"solved": solved,
"total_tasks": len(task_ids),
"raw_rate": solved / len(task_ids) if task_ids else None,
"posterior_mean": cell.posterior_mean,
"ci_low": cell.ci_low,
"ci_high": cell.ci_high,
"total_cost_usd": total_cost,
"attempted_tasks": attempted_tasks,
"mean_cost_per_attempted_task_usd": total_cost / attempted_tasks if attempted_tasks else None,
"cost_per_success_usd": total_cost / solved if solved else None,
"attempt_records": sum(row["model"] == model for row in records),
}
def _case_rows(
task_ids: list[str],
labels: dict[str, dict[str, object]],
records: list[dict[str, object]],
outcomes: dict[tuple[str, str], dict[str, object]],
capability_map: dict[str, dict[str, object]],
cfg: PolicyConfig,
) -> list[dict[str, object]]:
rows = []
for task_id in task_ids:
task_labels = labels.get(task_id, {})
offline_data = task_labels.get("offline", {})
jev_data = task_labels.get("jev", {})
offline_label = offline_data.get("label") if isinstance(offline_data, dict) else None
jev_label = jev_data.get("label") if isinstance(jev_data, dict) else None
decision = choose_model(jev_label, capability_map, cfg)
row: dict[str, object] = {
"task_id": task_id,
"offline_label": offline_label,
"jev_label": jev_label,
}
for model in MODEL_ORDER:
outcome = _model_outcome(outcomes, task_id, model)
row[f"{model}_success"] = outcome["success"]
row[f"{model}_cost_usd"] = outcome["cost_usd"]
row[f"{model}_exit_status"] = _last_exit_status(records, task_id, model)
row[f"{model}_attempts"] = outcome["attempts"]
row["jev_decision"] = decision.model
row["jev_decision_reason"] = decision.reason
row["retrospective_oracle"] = _oracle(outcomes, task_id)
rows.append(row)
return rows
def _leave_one_out(
task_ids: list[str],
labels: dict[str, dict[str, object]],
records: list[dict[str, object]],
categories: dict[str, str],
cfg: PolicyConfig,
) -> list[tuple[str, str, str]]:
return [
_leave_one_out_row(task_id, task_ids, labels, records, categories, cfg)
for task_id in task_ids
]
def _leave_one_out_row(
task_id: str,
task_ids: list[str],
labels: dict[str, dict[str, object]],
records: list[dict[str, object]],
categories: dict[str, str],
cfg: PolicyConfig,
) -> tuple[str, str, str]:
task_labels = labels.get(task_id, {})
jev_data = task_labels.get("jev", {})
jev_label = jev_data.get("label") if isinstance(jev_data, dict) else None
profile = build_capability_map(records, categories, list(MODEL_ORDER), cfg, exclude_task=task_id)
decision = choose_model(jev_label, profile, cfg)
return task_id, decision.model, decision.reason
def _render_report(
task_ids: list[str],
labels: dict[str, dict[str, object]],
records: list[dict[str, object]],
outcomes: dict[tuple[str, str], dict[str, object]],
case_rows: list[dict[str, object]],
fixed: dict[str, dict[str, object]],
loo_rows: list[tuple[str, str, str]],
aux: list[dict[str, object]],
preflight_cost: float,
total_runtime_s: float,
capability_map: dict[str, dict[str, object]],
cfg: PolicyConfig,
) -> str:
offline_count = sum(
isinstance(labels.get(task_id, {}).get("offline"), dict)
and labels[task_id]["offline"].get("label") is not None
for task_id in task_ids
)
agreement_rows = [
(labels[task_id]["offline"].get("label"), labels[task_id]["jev"].get("label"))
for task_id in task_ids
if isinstance(labels.get(task_id, {}).get("offline"), dict)
and isinstance(labels.get(task_id, {}).get("jev"), dict)
and labels[task_id]["offline"].get("label") is not None
and labels[task_id]["jev"].get("label") is not None
]
agreement = sum(offline == jev for offline, jev in agreement_rows)
routed_successes = sum(_model_outcome(outcomes, row["task_id"], str(row["jev_decision"]))["success"] is True for row in case_rows)
routed_cost = sum(float(_model_outcome(outcomes, row["task_id"], str(row["jev_decision"]))["cost_usd"]) for row in case_rows)
label_cost = sum(_cost(row.get("cost_usd")) for row in aux if row.get("kind") == "offline_label")
classifier_cost = sum(_cost(row.get("cost_usd")) for row in aux if row.get("kind") == "jev_classifier")
other_aux_cost = sum(
_cost(row.get("cost_usd"))
for row in aux
if row.get("kind") not in {"offline_label", "jev_classifier"}
)
solver_cost = sum(_cost(row.get("cost_usd")) for row in records)
infra_failures = sum(row.get("error_class") == "infra" for row in records)
solver_failures = sum(row.get("error_class") == "solver" for row in records)
report = [
"# Capability routing pilot analysis",
"",
"Five frozen tasks are too few to establish general routing superiority. No threshold was tuned, no validation set exists, and these results are not evidence of routing improvement",
"",
"## Per-task cases",
"",
"| Task | Offline | Jev | Haiku success / cost / exit / attempts | Sonnet success / cost / exit / attempts | Opus success / cost / exit / attempts | Jev decision | Decision reason | Retrospective oracle |",
"| --- | --- | --- | --- | --- | --- | --- | --- | --- |",
]
for row in case_rows:
cells = [
str(row["task_id"]),
str(row["offline_label"] or "missing"),
str(row["jev_label"] or "classifier failure"),
]
for model in MODEL_ORDER:
success = row[f"{model}_success"]
cells.append(
f"{success if success is not None else 'not graded'} / ${float(row[f'{model}_cost_usd']):.6f} / "
f"{row[f'{model}_exit_status']} / {row[f'{model}_attempts']}"
)
cells.extend([str(row["jev_decision"]), str(row["jev_decision_reason"]), str(row["retrospective_oracle"])])
report.append("| " + " | ".join(value.replace("|", "\\|") for value in cells) + " |")
report.extend(["", "## Fixed-model results", "", "| Model | Solved | Raw rate | Beta(1,1) mean | 90% CI | Total cost | Mean cost per attempted task | Cost per success |", "| --- | ---: | ---: | ---: | --- | ---: | ---: | ---: |"])
for model in MODEL_ORDER:
result = fixed[model]
ci = f"{result['ci_low']:.3f}–{result['ci_high']:.3f}"
mean_cost = result["mean_cost_per_attempted_task_usd"]
success_cost = result["cost_per_success_usd"]
report.append(
f"| {model} | {result['solved']}/{result['total_tasks']} | {result['raw_rate']:.3f} | "
f"{result['posterior_mean']:.3f} | {ci} | ${result['total_cost_usd']:.6f} | "
f"{f'${mean_cost:.6f}' if mean_cost is not None else 'n/a'} | "
f"{f'${success_cost:.6f}' if success_cost is not None else 'n/a'} |"
)
report.extend(
[
"",
"## Jev routing policy on this pilot",
"",
f"The router selected {routed_successes}/{len(task_ids)} outcomes for an observed cost of ${routed_cost:.6f}. "
f"It fell back to the reference model for {sum(row['jev_decision'] == 'opus-5' for row in case_rows)}/{len(task_ids)} tasks. "
f"For each candidate, T is the sorted intersection of reference and candidate task_success and task_cost keys. "
f"For each task in T, d is candidate success minus reference success, and the one-sided Hoeffding lower "
f"bound is mean(d) - sqrt(2*ln(1/alpha)/n), with alpha={cfg.alpha}. A candidate is eligible iff "
f"n >= {cfg.min_paired_tasks}, the lower bound is >= -{cfg.tolerance}, and its paired mean cost is lower "
"than the reference paired mean cost. The router picks the eligible candidate with the lowest paired "
"mean cost, breaking ties by model name, and otherwise falls back to the reference model",
"",
"Paired comparisons by category:",
"",
"| Category | Candidate | Paired n | Mean diff | Lower bound | Candidate paired mean cost | Reference paired mean cost |",
"| --- | --- | ---: | ---: | ---: | ---: | ---: |",
]
)
for category, cells in sorted(capability_map.items()):
reference = cells.get(cfg.reference)
candidates = [(model, cell) for model, cell in cells.items() if model != cfg.reference]
if not candidates:
report.append(f"| {category} | n/a | 0 | n/a | n/a | n/a | n/a |")
continue
if reference is None:
report.extend(f"| {category} | {model} | 0 | n/a | n/a | n/a | n/a |" for model, _ in candidates)
continue
for model, candidate in sorted(candidates, key=lambda item: item[0]):
comparison = paired_comparison(candidate, reference, cfg.alpha)
mean_diff = comparison["mean_diff"]
lower_bound = comparison["lower_bound"]
candidate_cost = comparison["cand_cost"]
reference_cost = comparison["ref_cost"]
report.append(
f"| {category} | {model} | {comparison['n']} | "
f"{f'{mean_diff:.3f}' if mean_diff is not None else 'n/a'} | "
f"{f'{lower_bound:.3f}' if lower_bound is not None else 'n/a'} | "
f"{f'${candidate_cost:.6f}' if candidate_cost is not None else 'n/a'} | "
f"{f'${reference_cost:.6f}' if reference_cost is not None else 'n/a'} |"
)
report.extend(
[
"",
"## RETROSPECTIVE UPPER BOUND, not achieved routing",
"",
"For each task, the oracle selects the cheapest model that succeeded on that task. This uses observed outcomes and is not an executable routing result",
"",
"## EXPLORATORY leave-one-task-out",
"",
"| Excluded task | Jev label decision | Reason |",
"| --- | --- | --- |",
]
)
report.extend(f"| {task_id} | {model} | {reason} |" for task_id, model, reason in loo_rows)
report.extend(
[
"",
"## Label agreement and resource use",
"",
f"Offline label coverage: {offline_count}/{len(task_ids)}. Offline/Jev agreement: {agreement}/{len(agreement_rows)} comparable labels",
f"Spend: solver ${solver_cost:.6f}, offline labeler ${label_cost:.6f}, Jev classifier ${classifier_cost:.6f}, "
f"preflight ${preflight_cost:.6f}, other auxiliary calls ${other_aux_cost:.6f}, total "
f"${solver_cost + label_cost + classifier_cost + preflight_cost + other_aux_cost:.6f}",
f"Total recorded runtime: {total_runtime_s:.2f}s. Infrastructure failure attempts: {infra_failures}. "
f"Solver failure attempts: {solver_failures}",
f"Capability map categories: {', '.join(sorted(capability_map)) or 'none'}",
]
)
return "\n".join(report) + "\n"
def analyze(data_dir: Path = ROOT) -> dict[str, object]:
records = load_jsonl(str(data_dir / "runs.jsonl")) if (data_dir / "runs.jsonl").exists() else []
labels = _labels(data_dir)
task_ids = _task_ids(data_dir)
model_payload = _json(data_dir / "poc" / "pilot_models.json", {"solvers": {}})
models = [model for model in MODEL_ORDER if model in model_payload["solvers"]]
cfg = PolicyConfig()
categories = {
task_id: parse_category(
str(labels.get(task_id, {}).get("offline", {}).get("label", "unknown"))
if isinstance(labels.get(task_id, {}).get("offline"), dict)
else "unknown"
)
for task_id in task_ids
}
capability_map = build_capability_map(records, categories, models, cfg)
outcomes = task_outcomes(records)
fixed = _fixed_model_summaries(task_ids, outcomes, records, cfg)
case_rows = _case_rows(task_ids, labels, records, outcomes, capability_map, cfg)
loo_rows = _leave_one_out(task_ids, labels, records, categories, cfg)
aux = _aux_costs(data_dir)
preflight_cost = _preflight_cost(data_dir)
label_runtime = sum(
float(value.get("latency_s", 0.0))
for task in labels.values()
for key in ("offline", "jev")
if isinstance((value := task.get(key)), dict)
and isinstance(value.get("latency_s"), (int, float))
)
preflight = _json(data_dir / PREFLIGHT_PATH, [])
preflight_runtime = sum(
float(row.get("latency", 0.0))
for row in preflight
if isinstance(row, dict) and isinstance(row.get("latency"), (int, float))
)
total_runtime_s = (
sum(_cost(row.get("solve_latency_s")) + _cost(row.get("eval_latency_s")) for row in records)
+ label_runtime
+ preflight_runtime
)
report = _render_report(
task_ids,
labels,
records,
outcomes,
case_rows,
fixed,
loo_rows,
aux,
preflight_cost,
total_runtime_s,
capability_map,
cfg,
)
(data_dir / "capability_map.json").write_text(json.dumps(capability_map, indent=2) + "\n")
(data_dir / "report.md").write_text(report)
if case_rows:
with (data_dir / "case_table.csv").open("w", newline="") as file:
writer = csv.DictWriter(file, fieldnames=list(case_rows[0]))
writer.writeheader()
writer.writerows(case_rows)
return {
"capability_map": capability_map,
"fixed_models": fixed,
"cases": case_rows,
"leave_one_out": loo_rows,
"report": report,
}
def main() -> None:
result = analyze()
print(
json.dumps(
{
"task_count": len(result["cases"]),
"model_summaries": result["fixed_models"],
"report": str(ROOT / "report.md"),
"case_table": str(ROOT / "case_table.csv"),
},
indent=2,
)
)
if __name__ == "__main__":
main()

View file

@ -0,0 +1,121 @@
import json
import os
import time
from pathlib import Path
import litellm
from minisweagent.exceptions import LimitsExceeded
from minisweagent.models.litellm_model import LitellmModel, LitellmModelConfig
from minisweagent.models.utils.actions_toolcall import BASH_TOOL
COST_HEADER = "llm_provider-x-litellm-response-cost"
class GuardedModelConfig(LitellmModelConfig):
run_budget_usd: float
worst_input_usd_per_token: float
output_usd_per_token: float
max_output_tokens: int
input_token_safety_factor: float = 1.25
fixed_overhead_tokens: int = 1024
call_log_path: str
def reserve_usd(input_tokens: int, cfg: GuardedModelConfig) -> float:
return (
input_tokens * cfg.input_token_safety_factor * cfg.worst_input_usd_per_token
+ cfg.max_output_tokens * cfg.output_usd_per_token
)
def estimate_input_tokens(counted: int, prev_counted: int | None, prev_actual: int | None,
cfg: GuardedModelConfig) -> int:
floor = counted + cfg.fixed_overhead_tokens
if prev_counted is None or prev_actual is None:
return floor
return max(floor, prev_actual + max(counted - prev_counted, 0))
def budget_allows(spent_usd: float, reservation_usd: float, budget_usd: float) -> bool:
return spent_usd + reservation_usd <= budget_usd
def gateway_cost(hidden_params: dict) -> float | None:
raw = (hidden_params.get("additional_headers") or {}).get(COST_HEADER)
return None if raw is None else float(raw)
class GuardedLitellmModel(LitellmModel):
def __init__(self, **kwargs):
super().__init__(config_class=GuardedModelConfig, **kwargs)
self.spent_usd = 0.0
self._pending_reservation = 0.0
self._pending_input_estimate = 0
self._pending_counted = 0
self._prev_counted: int | None = None
self._prev_actual: int | None = None
def query(self, messages: list[dict], **kwargs) -> dict:
os.environ["MSWEA_MODEL_RETRY_STOP_AFTER_ATTEMPT"] = "1"
prepared = self._prepare_messages_for_api(messages)
n_in = litellm.token_counter(
model=self.config.model_name.removeprefix("litellm_proxy/"), messages=prepared, tools=[BASH_TOOL]
)
n_est = estimate_input_tokens(n_in, self._prev_counted, self._prev_actual, self.config)
reservation = reserve_usd(n_est, self.config)
if not budget_allows(self.spent_usd, reservation, self.config.run_budget_usd):
self._log({"event": "budget_guard_refused", "spent_usd": self.spent_usd, "reservation_usd": reservation,
"counted_input_tokens": n_in, "est_input_tokens": n_est})
raise LimitsExceeded(
{"role": "exit", "content": "BudgetGuard", "extra": {"exit_status": "BudgetGuard", "submission": ""}}
)
self._pending_reservation = reservation
self._pending_input_estimate = n_est
self._pending_counted = n_in
return super().query(messages, **kwargs)
def _query(self, messages, **kwargs):
return super()._query(
messages,
**(
kwargs
| {
"max_tokens": self.config.max_output_tokens,
"timeout": 120.0,
"num_retries": 0,
}
),
)
def _calculate_cost(self, response) -> dict[str, float]:
cost = gateway_cost(response._hidden_params)
self.spent_usd += cost or 0.0
usage = response.usage.model_dump() if response.usage else {}
self._prev_counted = self._pending_counted
self._prev_actual = usage.get("prompt_tokens")
self._log({
"event": "call",
"call_id": response._hidden_params.get("litellm_call_id"),
"response_id": response.id,
"returned_model": response.model,
"gateway_model_id": (response._hidden_params.get("additional_headers") or {}).get(
"llm_provider-x-litellm-model-id"),
"cost_usd": cost,
"cost_source": "gateway_header" if cost is not None else "missing",
"reservation_usd": self._pending_reservation,
"est_input_tokens": self._pending_input_estimate,
"prompt_tokens": usage.get("prompt_tokens"),
"completion_tokens": usage.get("completion_tokens"),
"cache_read_input_tokens": usage.get("cache_read_input_tokens"),
"cache_creation_input_tokens": usage.get("cache_creation_input_tokens"),
"reasoning_tokens": (usage.get("completion_tokens_details") or {}).get("reasoning_tokens"),
"latency_ms": response._hidden_params.get("_response_ms"),
"finish_reason": response.choices[0].finish_reason,
"spent_usd_after": self.spent_usd,
"reservation_covered_actual": cost is not None and cost <= self._pending_reservation,
})
return {"cost": cost if cost is not None else 0.0}
def _log(self, row: dict) -> None:
with Path(self.config.call_log_path).open("a") as f:
f.write(json.dumps({"ts": time.time(), **row}) + "\n")

View file

@ -0,0 +1,338 @@
import argparse
import fcntl
import json
import math
import os
import re
import stat
import time
from datetime import UTC, datetime
from json import JSONDecodeError, JSONDecoder
from pathlib import Path
import httpx
from poc.policy import CATEGORIES, parse_category
ROOT = Path(__file__).resolve().parents[1]
BASE_URL = os.environ.get("LITELLM_GATEWAY_BASE_URL", "https://gateway.litellm-sandbox.ai").rstrip("/")
LABEL_MODEL = "openai/gpt-5.4-mini-2026-03-17"
JEV_MODEL = "typesafe/jev-latest"
REQUEST_TIMEOUT_S = 120.0
SYSTEM_PROMPT = (
"You label software engineering tasks by the main capability needed to solve them. Use only the problem "
"statement. Answer with JSON only: {\"label\": <one of repository_debugging, feature_implementation, "
"shell_build, algorithms_data, unknown>, \"rationale\": <one sentence>}. repository_debugging: diagnose and "
"fix incorrect behavior in an existing codebase. feature_implementation: add new functionality, options or "
"APIs. shell_build: build, packaging, environment, CI or shell tooling. algorithms_data: algorithmic, numerical "
"or data-processing logic is the core difficulty. unknown: cannot tell from the statement."
)
JEV_CRITERIA = {
"repository_debugging": "diagnose and fix incorrect behavior in an existing codebase",
"feature_implementation": "add new functionality, options or APIs",
"shell_build": "build, packaging, environment, CI or shell tooling",
"algorithms_data": "algorithmic, numerical or data-processing logic is the core difficulty",
"unknown": "cannot tell from the statement",
}
def _first_json_object(text: str) -> dict[str, object] | None:
decoder = JSONDecoder()
for index, character in enumerate(text):
if character != "{":
continue
try:
value, _ = decoder.raw_decode(text[index:])
except JSONDecodeError:
continue
if isinstance(value, dict):
return value
return None
def parse_label_response(text: str) -> tuple[str, str | None, bool]:
payload = _first_json_object(text)
if payload is None or not isinstance(payload.get("label"), str):
return "unknown", None, False
raw_label = payload["label"]
label = parse_category(raw_label)
if label == "unknown" and raw_label.strip().lower() != "unknown":
return "unknown", None, False
rationale = payload.get("rationale")
if rationale is not None and not isinstance(rationale, str):
return "unknown", None, False
return label, rationale, True
def _as_mapping(value: object) -> dict[str, object] | None:
return value if isinstance(value, dict) else None
def _finite_number(value: object) -> float | None:
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
return None
return float(value)
def parse_jev_response(
payload: object,
) -> tuple[str | None, float | None, dict[str, float]]:
response = _as_mapping(payload)
answers = None if response is None else _as_mapping(response.get("answers"))
answer = None if answers is None else _as_mapping(answers.get("capability"))
if answer is None and response is not None:
answer = _as_mapping(response.get("capability")) or _as_mapping(response.get("answer"))
if answer is None:
return None, None, {}
choice = answer.get("choice", answer.get("answer", answer.get("value")))
if not isinstance(choice, str):
return None, None, {}
label = parse_category(choice)
if label == "unknown" and choice.strip().lower() != "unknown":
return None, None, {}
confidence = _finite_number(answer.get("confidence"))
raw_probabilities = _as_mapping(answer.get("probabilities"))
probabilities = {
category: number
for key, raw_value in (raw_probabilities or {}).items()
if isinstance(key, str)
and (category := parse_category(key)) in CATEGORIES
and (category != "unknown" or key.strip().lower() == "unknown")
and (number := _finite_number(raw_value)) is not None
}
return label, confidence, probabilities
def _read_api_key() -> str:
path = ROOT / ".private" / "pilot_key"
if stat.S_IMODE(path.stat().st_mode) != 0o600:
raise PermissionError("pilot key must have mode 0600")
key = path.read_text().strip()
if not key:
raise ValueError("pilot key is empty")
return key
def _headers() -> dict[str, str]:
api_key = os.environ.get("LITELLM_PROXY_API_KEY")
if not api_key:
raise RuntimeError("LITELLM_PROXY_API_KEY is not set")
return {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
def _json_response(response: httpx.Response) -> dict[str, object]:
try:
payload = response.json()
except ValueError:
return {}
return payload if isinstance(payload, dict) else {}
def _response_metadata(response: httpx.Response, payload: dict[str, object], latency_s: float) -> dict[str, object]:
headers = response.headers
usage = payload.get("usage")
raw_cost = headers.get("x-litellm-response-cost") or headers.get("llm_provider-x-litellm-response-cost")
cost = _finite_number(raw_cost) if isinstance(raw_cost, (int, float)) else None
if isinstance(raw_cost, str):
try:
cost = _finite_number(float(raw_cost))
except ValueError:
cost = None
return {
"call_id": headers.get("x-litellm-call-id"),
"response_model": payload.get("model") or headers.get("x-litellm-response-model"),
"reported_version": payload.get("version")
or headers.get("x-litellm-model-version")
or headers.get("x-litellm-version"),
"cost_usd": cost,
"latency_s": latency_s,
"usage": usage if isinstance(usage, dict) else {},
}
def _append_auxiliary_call(kind: str, task_id: str, model: str, metadata: dict[str, object]) -> None:
path = ROOT / "costs_aux.jsonl"
path.parent.mkdir(parents=True, exist_ok=True)
row = {
"kind": kind,
"task_id": task_id,
"model": model,
"cost_usd": metadata["cost_usd"],
"call_id": metadata["call_id"],
"ts": datetime.now(UTC).isoformat(),
}
with path.open("a") as file:
fcntl.flock(file, fcntl.LOCK_EX)
file.write(json.dumps(row, separators=(",", ":")) + "\n")
file.flush()
fcntl.flock(file, fcntl.LOCK_UN)
def _sanitize_label_text(text: str) -> str:
return re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", text).strip()[:1200]
def _problem_statement_text(content: object) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
return "".join(
fragment.get("text", "")
for fragment in content
if isinstance(fragment, dict) and isinstance(fragment.get("text"), str)
)
return ""
def _offline_label(client: httpx.Client, task_id: str, problem_statement: str) -> dict[str, object]:
body = {
"model": LABEL_MODEL,
"max_completion_tokens": 400,
"messages": [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": f"Problem statement:\n<<<\n{problem_statement}\n>>>"},
],
}
started = time.perf_counter()
response: httpx.Response | None = None
payload: dict[str, object] = {}
error: str | None = None
try:
response = client.post(f"{BASE_URL}/v1/chat/completions", headers=_headers(), json=body)
payload = _json_response(response)
if response.is_error:
error = f"http_{response.status_code}"
except httpx.HTTPError as exception:
error = type(exception).__name__
latency_s = time.perf_counter() - started
metadata = (
_response_metadata(response, payload, latency_s)
if response is not None
else {"call_id": None, "response_model": None, "reported_version": None, "cost_usd": None,
"latency_s": latency_s, "usage": {}}
)
choices = payload.get("choices")
choice = choices[0] if isinstance(choices, list) and choices and isinstance(choices[0], dict) else {}
message = choice.get("message") if isinstance(choice, dict) else {}
raw_text = _problem_statement_text(message.get("content")) if isinstance(message, dict) else ""
label, rationale, parse_ok = parse_label_response(raw_text)
if error is not None:
label, rationale, parse_ok = "unknown", None, False
_append_auxiliary_call("offline_label", task_id, LABEL_MODEL, metadata)
return {
**metadata,
"label": label,
"rationale": rationale,
"parse_ok": parse_ok,
"raw_label_text": _sanitize_label_text(raw_text),
"error": error,
}
def _jev_label(client: httpx.Client, task_id: str, problem_statement: str) -> dict[str, object]:
body = {
"model": JEV_MODEL,
"state": problem_statement,
"questions": {
"capability": {
"type": "choice",
"instructions": "Which capability does solving this software task mainly require?",
"criteria": JEV_CRITERIA,
}
},
}
started = time.perf_counter()
response: httpx.Response | None = None
payload: dict[str, object] = {}
error: str | None = None
try:
response = client.post(f"{BASE_URL}/v1/systemone", headers=_headers(), json=body)
payload = _json_response(response)
if response.is_error:
error = f"http_{response.status_code}"
except httpx.HTTPError as exception:
error = type(exception).__name__
latency_s = time.perf_counter() - started
metadata = (
_response_metadata(response, payload, latency_s)
if response is not None
else {"call_id": None, "response_model": None, "reported_version": None, "cost_usd": None,
"latency_s": latency_s, "usage": {}}
)
label, confidence, probabilities = parse_jev_response(payload)
response_mapping = _as_mapping(payload)
answers = None if response_mapping is None else _as_mapping(response_mapping.get("answers"))
raw_answer = None if answers is None else _as_mapping(answers.get("capability"))
if raw_answer is None and response_mapping is not None:
raw_answer = _as_mapping(response_mapping.get("capability")) or _as_mapping(response_mapping.get("answer"))
raw_label_text = json.dumps(
{
key: value
for key, value in (raw_answer or {}).items()
if key in {"choice", "answer", "value", "confidence", "probabilities"}
},
ensure_ascii=False,
separators=(",", ":"),
)
if error is not None:
label, confidence, probabilities = None, None, {}
_append_auxiliary_call("jev_classifier", task_id, JEV_MODEL, metadata)
return {
**metadata,
"label": label,
"confidence": confidence,
"probabilities": probabilities,
"raw_label_text": _sanitize_label_text(raw_label_text),
"error": error,
}
def _load_tasks() -> tuple[dict[str, str], list[str]]:
frozen = json.loads((ROOT / "tasks.json").read_text())
task_ids = [task["instance_id"] for task in frozen["tasks"]]
statements = {
row["instance_id"]: row["problem_statement"]
for row in (json.loads(line) for line in (ROOT / "agent_dataset" / "train.jsonl").read_text().splitlines())
if row.get("instance_id") in task_ids and isinstance(row.get("problem_statement"), str)
}
missing = set(task_ids) - statements.keys()
if missing:
raise ValueError(f"missing sanitized training statements for {sorted(missing)}")
return statements, task_ids
def label_tasks() -> list[dict[str, object]]:
if (ROOT / "labels.json").exists():
raise FileExistsError("labels.json already exists; remove it only before intentionally relabeling")
os.environ["LITELLM_PROXY_API_KEY"] = _read_api_key()
statements, task_ids = _load_tasks()
with httpx.Client(transport=httpx.HTTPTransport(retries=0), timeout=REQUEST_TIMEOUT_S) as client:
rows = [
{
"task_id": task_id,
"offline": _offline_label(client, task_id, statements[task_id]),
"jev": _jev_label(client, task_id, statements[task_id]),
}
for task_id in task_ids
]
(ROOT / "labels.json").write_text(json.dumps(rows, indent=2) + "\n")
return rows
def main() -> None:
argparse.ArgumentParser(description="Label the frozen POC task statements").parse_args()
rows = label_tasks()
summary = [
{
"task_id": row["task_id"],
"offline": row["offline"]["label"],
"jev": row["jev"]["label"],
"cost_usd": (row["offline"]["cost_usd"] or 0.0) + (row["jev"]["cost_usd"] or 0.0),
}
for row in rows
]
print(json.dumps(summary, indent=2))
if __name__ == "__main__":
main()

View file

@ -0,0 +1,11 @@
agent:
step_limit: 50
cost_limit: 1.0
wall_time_limit_seconds: 540
environment:
run_args: ["--rm", "--label", "routing_poc=1"]
model:
model_class: poc.guarded_model.GuardedLitellmModel
run_budget_usd: 1.0
max_output_tokens: 8192
input_token_safety_factor: 1.25

View file

@ -0,0 +1,14 @@
{
"source": "gateway.litellm-sandbox.ai /v1/model/info, fetched 2026-10-20",
"reference": "opus-5",
"solvers": {
"haiku-4.5": {"gateway_model": "anthropic/claude-haiku-4-5-20251001", "deployment_id": "13ef3d99-7edf-4a6b-aab6-f088c8e25964",
"input": 1e-06, "output": 5e-06, "cache_read": 1e-07, "cache_write": 1.25e-06},
"sonnet-5": {"gateway_model": "anthropic/claude-sonnet-5", "deployment_id": "13ef3d99-7edf-4a6b-aab6-f088c8e25964",
"input": 2e-06, "output": 1e-05, "cache_read": 2e-07, "cache_write": 2.5e-06},
"opus-5": {"gateway_model": "anthropic/claude-opus-5", "deployment_id": "13ef3d99-7edf-4a6b-aab6-f088c8e25964",
"input": 5e-06, "output": 2.5e-05, "cache_read": 5e-07, "cache_write": 6.25e-06}
},
"labeler": {"gateway_model": "openai/gpt-5.4-mini-2026-03-17", "input": 7.5e-07, "output": 4.5e-06},
"classifier": {"gateway_model": "typesafe/jev-latest", "reported_version": "jev-1.13.0", "input": 4.2e-08, "output": 0.0}
}

View file

@ -0,0 +1,163 @@
import json
import math
import re
from collections import defaultdict
from dataclasses import dataclass
from scipy.stats import beta
CATEGORIES = ("repository_debugging", "feature_implementation", "shell_build", "algorithms_data", "unknown")
@dataclass(frozen=True, slots=True)
class PolicyConfig:
reference: str = "opus-5"
min_paired_tasks: int = 10
tolerance: float = 0.05
alpha: float = 0.05
credible_mass: float = 0.90
@dataclass(frozen=True, slots=True)
class Cell:
n: int
s: int
raw_rate: float | None
posterior_mean: float
ci_low: float
ci_high: float
mean_cost_per_task_usd: float | None
@dataclass(frozen=True, slots=True)
class Decision:
model: str
reason: str
def parse_category(raw: object) -> str:
if not isinstance(raw, str):
return "unknown"
norm = re.sub(r"[^a-z]+", "_", raw.strip().lower()).strip("_")
aliases = {"shell_build_work": "shell_build", "algorithms_data_processing": "algorithms_data"}
norm = aliases.get(norm, norm)
return norm if norm in CATEGORIES else "unknown"
def beta_cell(s: int, n: int, mass: float, mean_cost: float | None) -> Cell:
a, b = s + 1, n - s + 1
tail = (1 - mass) / 2
return Cell(n=n, s=s, raw_rate=(s / n) if n else None, posterior_mean=a / (a + b),
ci_low=float(beta.ppf(tail, a, b)), ci_high=float(beta.ppf(1 - tail, a, b)),
mean_cost_per_task_usd=mean_cost)
def task_outcomes(records: list[dict]) -> dict[tuple[str, str], dict]:
"""One outcome per (task, model). Infra-failed attempts add cost but never success or n."""
grouped: dict[tuple[str, str], list[dict]] = defaultdict(list)
for r in records:
grouped[(r["task_id"], r["model"])].append(r)
out = {}
for key, attempts in grouped.items():
graded = [a for a in attempts if a["error_class"] != "infra"]
if len(graded) > 1:
raise ValueError(f"duplicate graded attempts for {key}: single-attempt pilot")
out[key] = {
"graded": bool(graded),
"success": graded[0]["success"] if graded else None,
"cost_usd": sum(a["cost_usd"] for a in attempts),
"attempts": len(attempts),
}
return out
def build_capability_map(records: list[dict], task_category: dict[str, str], models: list[str],
cfg: PolicyConfig, exclude_task: str | None = None) -> dict:
outcomes = {k: v for k, v in task_outcomes(records).items() if k[0] != exclude_task}
cmap: dict = {}
for cat in CATEGORIES:
tasks = sorted({t for (t, _m) in outcomes if task_category.get(t, "unknown") == cat})
if not tasks:
continue
cmap[cat] = {}
for m in models:
graded = [outcomes[(t, m)] for t in tasks if (t, m) in outcomes and outcomes[(t, m)]["graded"]]
attempted = [outcomes[(t, m)] for t in tasks if (t, m) in outcomes]
mean_cost = sum(o["cost_usd"] for o in attempted) / len(attempted) if attempted else None
cell = beta_cell(sum(o["success"] for o in graded), len(graded), cfg.credible_mass, mean_cost)
cmap[cat][m] = {
f: getattr(cell, f) for f in Cell.__slots__} | {"graded_tasks": sorted(
t for t in tasks if (t, m) in outcomes and outcomes[(t, m)]["graded"])} | {
"task_success": {
t: outcomes[(t, m)]["success"]
for t in tasks
if (t, m) in outcomes and outcomes[(t, m)]["graded"]
},
"task_cost": {
t: outcomes[(t, m)]["cost_usd"]
for t in tasks
if (t, m) in outcomes
},
}
return cmap
def paired_lower_bound(diffs: list[int], alpha: float) -> float:
"""One-sided Hoeffding lower bound for paired differences in [-1, 1]."""
if not diffs:
raise ValueError("paired differences cannot be empty")
if not 0 < alpha < 1:
raise ValueError("alpha must be between 0 and 1")
return sum(diffs) / len(diffs) - math.sqrt(2 * math.log(1 / alpha) / len(diffs))
def paired_comparison(cand_cell: dict, ref_cell: dict, alpha: float) -> dict[str, int | float | None]:
cand_success = cand_cell["task_success"]
ref_success = ref_cell["task_success"]
cand_cost = cand_cell["task_cost"]
ref_cost = ref_cell["task_cost"]
paired_tasks = sorted(set(cand_success) & set(ref_success) & set(cand_cost) & set(ref_cost))
if not paired_tasks:
return {"n": 0, "mean_diff": None, "lower_bound": None, "cand_cost": None, "ref_cost": None}
diffs = [int(cand_success[task]) - int(ref_success[task]) for task in paired_tasks]
return {
"n": len(paired_tasks),
"mean_diff": sum(diffs) / len(diffs),
"lower_bound": paired_lower_bound(diffs, alpha),
"cand_cost": sum(cand_cost[task] for task in paired_tasks) / len(paired_tasks),
"ref_cost": sum(ref_cost[task] for task in paired_tasks) / len(paired_tasks),
}
def choose_model(category: str | None, cmap: dict, cfg: PolicyConfig) -> Decision:
cat = parse_category(category)
if category is None:
return Decision(cfg.reference, "fallback: classifier failure")
if cat == "unknown" or cat not in cmap or cfg.reference not in cmap[cat]:
return Decision(cfg.reference, f"fallback: no profile for category {cat}")
cells = cmap[cat]
ref = cells[cfg.reference]
eligible = []
for m, c in cells.items():
if m == cfg.reference:
continue
comparison = paired_comparison(c, ref, cfg.alpha)
if comparison["n"] < cfg.min_paired_tasks:
continue
if (
comparison["lower_bound"] is not None
and comparison["lower_bound"] >= -cfg.tolerance
and comparison["cand_cost"] is not None
and comparison["ref_cost"] is not None
and comparison["cand_cost"] < comparison["ref_cost"]
):
eligible.append((comparison["cand_cost"], m))
if not eligible:
return Decision(cfg.reference, f"fallback: no cheaper model with >= {cfg.min_paired_tasks} paired tasks "
f"and paired lower bound >= -{cfg.tolerance} (alpha={cfg.alpha})")
return Decision(min(eligible)[1], "cheapest model with paired lower bound within tolerance of reference")
def load_jsonl(path: str) -> list[dict]:
with open(path) as f:
return [json.loads(line) for line in f if line.strip()]

View file

@ -0,0 +1,597 @@
import argparse
import fcntl
import json
import os
import shutil
import stat
import subprocess
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import UTC, datetime
from pathlib import Path
import yaml
ROOT = Path(__file__).resolve().parents[1]
RUNS_PATH = ROOT / "runs.jsonl"
AUX_COSTS_PATH = ROOT / "costs_aux.jsonl"
RUNS_ROOT = ROOT / "runs_raw"
GLOBAL_BUDGET_USD = 15.0
RESERVATION_USD = 1.0
MIN_FREE_BYTES = 20 * 1024**3
MAX_LOG_BYTES = 200 * 1024
INFRA_EXIT_STATUSES = {
"APIConnectionError",
"InternalServerError",
"ServiceUnavailableError",
"RateLimitError",
"BadGatewayError",
"Timeout",
"ImportError",
"ModuleNotFoundError",
"DockerException",
"ImageNotFound",
"ContainerError",
"EnvironmentError",
"ConnectionError",
}
SOLVER_EXIT_STATUSES = {"Submitted", "LimitsExceeded", "BudgetGuard", "TimeExceeded", "FormatError"}
DATASET_METADATA = {
"agent": "princeton-nlp/SWE-bench@e48e2bd1e9fecd5bbd641e9414ac59da9f2e69f6 dev (sanitized)",
"eval": "SWE-bench/SWE-bench@c6fe717fd7a4c3ac1daa4055a4fd082c6a1d28a2 dev",
}
TOKEN_FIELDS = {
"prompt": "prompt_tokens",
"completion": "completion_tokens",
"cache_read": "cache_read_input_tokens",
"cache_creation": "cache_creation_input_tokens",
"reasoning": "reasoning_tokens",
}
def append_jsonl(path: Path, row: dict[str, object]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("a") as file:
fcntl.flock(file, fcntl.LOCK_EX)
file.write(json.dumps(row, separators=(",", ":")) + "\n")
file.flush()
fcntl.flock(file, fcntl.LOCK_UN)
def read_jsonl_locked(path: Path) -> list[dict[str, object]]:
if not path.exists():
return []
with path.open() as file:
fcntl.flock(file, fcntl.LOCK_SH)
rows = [json.loads(line) for line in file if line.strip()]
fcntl.flock(file, fcntl.LOCK_UN)
return [row for row in rows if isinstance(row, dict)]
def committed_spend_usd(runs_path: Path = RUNS_PATH, aux_path: Path = AUX_COSTS_PATH) -> float:
run_cost = sum(_cost(row.get("cost_usd")) for row in read_jsonl_locked(runs_path))
aux_cost = sum(_cost(row.get("cost_usd")) for row in read_jsonl_locked(aux_path))
return run_cost + aux_cost
def _cost(value: object) -> float:
return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else 0.0
def _utility_env(env: dict[str, str]) -> dict[str, str]:
return {
name: value
for name, value in env.items()
if name not in {"LITELLM_PROXY_API_KEY", "LITELLM_GATEWAY_ADMIN_KEY"}
}
def global_budget_allows(committed_spend: float, in_flight: int, budget: float = GLOBAL_BUDGET_USD) -> bool:
return committed_spend + RESERVATION_USD * (in_flight + 1) <= budget
def sum_call_costs(calls: list[dict[str, object]]) -> tuple[float, bool]:
present_costs = [call.get("cost_usd") for call in calls if call.get("event") == "call"]
return sum(_cost(cost) for cost in present_costs), any(
not isinstance(cost, (int, float)) or isinstance(cost, bool) for cost in present_costs
)
def sum_call_tokens(calls: list[dict[str, object]]) -> dict[str, int]:
totals = {name: 0 for name in TOKEN_FIELDS}
for call in calls:
if call.get("event") != "call":
continue
for name, source in TOKEN_FIELDS.items():
value = call.get(source)
totals[name] += int(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else 0
return totals
def classify_error(
exit_status: str | None,
rc: int,
submission: str,
eval_status: str | None,
trajectory_found: bool,
exception_text: str = "",
) -> str:
if eval_status == "resolved":
return "none"
if eval_status == "error":
return "infra"
if rc == 124 or exit_status in SOLVER_EXIT_STATUSES:
return "solver"
if not trajectory_found and rc != 124:
return "infra"
if exit_status in INFRA_EXIT_STATUSES:
return "infra"
if exit_status == "RuntimeError" and any(
marker in exception_text.lower() for marker in ("environment", "docker", "startup command")
):
return "infra"
return "solver" if submission or eval_status == "unresolved" or trajectory_found else "infra"
def _timestamp() -> str:
return datetime.now(UTC).strftime("%Y%m%dT%H%M%S%fZ")
def _utc_now() -> str:
return datetime.now(UTC).isoformat()
def _read_pilot_key() -> str:
path = ROOT / ".private" / "pilot_key"
if stat.S_IMODE(path.stat().st_mode) != 0o600:
raise PermissionError("pilot key must have mode 0600")
key = path.read_text().strip()
if not key:
raise ValueError("pilot key is empty")
return key
def _read_json(path: Path) -> object:
return json.loads(path.read_text())
def _labels_by_task() -> dict[str, dict[str, object]]:
path = ROOT / "labels.json"
if not path.exists():
return {}
entries = _read_json(path)
return {row["task_id"]: row for row in entries if isinstance(row, dict) and isinstance(row.get("task_id"), str)}
def _model_config(model: str, run_id: str, model_data: dict[str, object], base_url: str) -> dict[str, object]:
return {
"model": {
"model_name": f"litellm_proxy/{model_data['gateway_model']}",
"model_kwargs": {"api_base": base_url},
"worst_input_usd_per_token": max(
float(model_data["input"]),
float(model_data.get("cache_write", model_data["input"])),
),
"output_usd_per_token": float(model_data["output"]),
"call_log_path": str(RUNS_ROOT / run_id / "calls.jsonl"),
},
"environment": {
"run_args": ["--rm", "--label", "routing_poc=1", "--label", f"routing_poc_run={run_id}"],
},
}
def _bounded_log(path: Path) -> None:
size = path.stat().st_size
if size <= MAX_LOG_BYTES:
return
with path.open("rb") as file:
file.seek(-MAX_LOG_BYTES, os.SEEK_END)
tail = file.read()
path.write_bytes(tail)
def _calls(path: Path) -> list[dict[str, object]]:
if not path.exists():
return []
return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
def _trajectory(path: Path) -> tuple[dict[str, object], bool]:
if not path.exists():
return {}, False
payload = _read_json(path)
info = payload.get("info") if isinstance(payload, dict) else None
return (info if isinstance(info, dict) else {}), True
def _tail_exception_text(info: dict[str, object]) -> str:
return str(info.get("exception_str", ""))
def _report_status(report_dir: Path, task_id: str) -> str | None:
for path in report_dir.rglob("*.json"):
try:
report = _read_json(path)
except (OSError, json.JSONDecodeError):
continue
if not isinstance(report, dict):
continue
if task_id in report.get("resolved_ids", []):
return "resolved"
if task_id in report.get("unresolved_ids", []):
return "unresolved"
if task_id in report.get("error_ids", []):
return "error"
return None
def _grade(task_id: str, model: str, submission: str, run_id: str, run_dir: Path, env: dict[str, str]) -> tuple[str, float]:
if not submission.strip():
return "skipped_empty_patch", 0.0
eval_dir = run_dir / "eval"
eval_dir.mkdir(parents=True, exist_ok=True)
predictions_path = run_dir / "predictions.json"
predictions_path.write_text(
json.dumps(
[
{
"instance_id": task_id,
"model_name_or_path": model,
"model_patch": submission,
}
],
indent=2,
)
+ "\n"
)
command = [
str(ROOT / ".venv" / "bin" / "python"),
"-m",
"swebench.harness.run_evaluation",
"-d",
"SWE-bench/SWE-bench",
"-s",
"dev",
"-i",
task_id,
"-p",
str(predictions_path),
"-id",
run_id,
"--max_workers",
"1",
"-t",
"900",
"--report_dir",
str(eval_dir),
]
started = time.monotonic()
with (eval_dir / "eval.log").open("wb") as log:
try:
subprocess.run(
command,
cwd=eval_dir,
env=_utility_env(env),
stdout=log,
stderr=subprocess.STDOUT,
timeout=930,
)
except subprocess.TimeoutExpired:
pass
_bounded_log(eval_dir / "eval.log")
status = _report_status(eval_dir, task_id)
if status is None:
return "error", time.monotonic() - started
return status, time.monotonic() - started
def _cleanup_containers(run_id: str, env: dict[str, str]) -> int:
result = subprocess.run(
["docker", "ps", "-aq", "--filter", f"label=routing_poc_run={run_id}"],
cwd=ROOT,
env=_utility_env(env),
capture_output=True,
text=True,
check=False,
)
container_ids = result.stdout.split() if result.returncode == 0 else []
if not container_ids:
return 0
removed = subprocess.run(
["docker", "rm", "-f", *container_ids],
cwd=ROOT,
env=_utility_env(env),
capture_output=True,
)
return len(container_ids) if removed.returncode == 0 else 0
def _dataset_metadata(task_id: str) -> tuple[str | None, dict[str, str]]:
rows = [json.loads(line) for line in (ROOT / "agent_dataset" / "train.jsonl").read_text().splitlines()]
task = next((row for row in rows if row.get("instance_id") == task_id), {})
image_lines = (ROOT / "images_pinned.txt").read_text().splitlines()
image_pins = {
parts[0]: parts[1]
for line in image_lines
if len(parts := line.split()) >= 2
}
return task.get("base_commit"), image_pins
def _attempt(
task_id: str,
model: str,
attempt: int,
task: dict[str, object],
model_data: dict[str, object],
category_offline: object,
category_jev: object,
env: dict[str, str],
base_url: str,
) -> dict[str, object]:
started_at = _utc_now()
started = time.monotonic()
run_id = f"{task_id}__{model}__a{attempt}__{_timestamp()}"
run_dir = RUNS_ROOT / run_id
run_dir.mkdir(parents=True, exist_ok=False)
config = _model_config(model, run_id, model_data, base_url)
model_config_path = run_dir / "model.yaml"
model_config_path.write_text(yaml.safe_dump(config, sort_keys=False))
solve_dir = run_dir / "solve"
solve_log = run_dir / "solve.log"
command = [
"timeout",
"600",
str(ROOT / ".venv" / "bin" / "mini-extra"),
"swebench",
"--subset",
str(ROOT / "agent_dataset"),
"--split",
"train",
"--filter",
f"^{task_id}$",
"-o",
str(solve_dir),
"-w",
"1",
"-c",
"swebench.yaml",
"-c",
"poc/pilot.yaml",
"-c",
str(model_config_path),
]
with solve_log.open("wb") as log:
try:
completed = subprocess.run(command, cwd=ROOT, env=env, stdout=log, stderr=subprocess.STDOUT, timeout=615)
rc = completed.returncode
except subprocess.TimeoutExpired:
rc = 124
_bounded_log(solve_log)
solve_latency_s = time.monotonic() - started
containers_removed = _cleanup_containers(run_id, env)
trajectory_path = solve_dir / task_id / f"{task_id}.traj.json"
info, trajectory_found = _trajectory(trajectory_path)
exit_status = info.get("exit_status")
submission = info.get("submission") if isinstance(info.get("submission"), str) else ""
calls_path = run_dir / "calls.jsonl"
calls = _calls(calls_path)
cost_usd, cost_incomplete = sum_call_costs(calls)
tokens = sum_call_tokens(calls)
eval_status, eval_latency_s = _grade(task_id, model, submission, run_id, run_dir, env)
if eval_status == "skipped_empty_patch":
success = False
else:
success = eval_status == "resolved"
error_class = classify_error(
str(exit_status) if exit_status is not None else None,
rc,
submission,
eval_status,
trajectory_found,
_tail_exception_text(info),
)
base_commit, image_pins = _dataset_metadata(task_id)
call_rows = [call for call in calls if call.get("event") == "call"]
model_stats = info.get("model_stats")
trajectory_calls = model_stats.get("api_calls") if isinstance(model_stats, dict) else None
result = {
"task_id": task_id,
"run_id": run_id,
"attempt": attempt,
"model": model,
"gateway_model": model_data["gateway_model"],
"returned_models": sorted({str(call["returned_model"]) for call in call_rows if call.get("returned_model")}),
"deployment_ids": sorted({str(call["gateway_model_id"]) for call in call_rows if call.get("gateway_model_id")}),
"category_offline": category_offline,
"category_jev": category_jev,
"versions": {
"mini_swe_agent": _package_version("mini-swe-agent"),
"swebench": _package_version("swebench"),
"litellm": _package_version("litellm"),
"python": os.sys.version.split()[0],
},
"dataset": DATASET_METADATA,
"image": task.get("image"),
"image_digest": image_pins.get(str(task.get("image"))),
"base_commit": base_commit,
"exit_status": exit_status,
"traj_api_calls": trajectory_calls,
"success": success,
"eval_status": eval_status,
"error_class": error_class,
"rc": rc,
"cost_usd": cost_usd,
"cost_incomplete": cost_incomplete,
"n_calls": len(call_rows),
"tokens": tokens,
"solve_latency_s": solve_latency_s,
"eval_latency_s": eval_latency_s,
"call_ids": [call.get("call_id") for call in call_rows if call.get("call_id")],
"budget_guard_refused": any(call.get("event") == "budget_guard_refused" for call in calls),
"containers_removed": containers_removed,
"started_at": started_at,
"finished_at": _utc_now(),
}
append_jsonl(RUNS_PATH, result)
return result
def _package_version(name: str) -> str:
from importlib.metadata import version
return version(name)
class DispatchGuard:
def __init__(self, stop_path: Path, budget: float = GLOBAL_BUDGET_USD, min_free_bytes: int = MIN_FREE_BYTES):
self._lock = threading.Lock()
self._in_flight = 0
self._stopped = False
self._stop_path = stop_path
self._budget = budget
self._min_free_bytes = min_free_bytes
def acquire(self, task_id: str) -> bool:
with self._lock:
if self._stopped:
return False
free_bytes = shutil.disk_usage("/").free
if free_bytes < self._min_free_bytes:
reason = f"disk_free_below_{self._min_free_bytes}"
elif not global_budget_allows(committed_spend_usd(), self._in_flight, self._budget):
reason = "global_spend_reservation_exceeds_budget"
else:
self._in_flight += 1
return True
self._stopped = True
append_jsonl(
self._stop_path,
{"task_id": task_id, "reason": reason, "free_bytes": free_bytes, "ts": _utc_now()},
)
return False
def release(self) -> None:
with self._lock:
self._in_flight -= 1
def _run_task(
task_id: str,
task: dict[str, object],
models: list[str],
solvers: dict[str, dict[str, object]],
labels: dict[str, dict[str, object]],
guard: DispatchGuard,
base_url: str,
base_env: dict[str, str],
) -> list[dict[str, object]]:
offline = labels.get(task_id, {}).get("offline", {})
jev = labels.get(task_id, {}).get("jev", {})
category_offline = offline.get("label") if isinstance(offline, dict) else None
category_jev = jev.get("label") if isinstance(jev, dict) else None
results = []
for model in models:
attempt = 1
while attempt <= 2:
if not guard.acquire(task_id):
return results
try:
row = _attempt(
task_id,
model,
attempt,
task,
solvers[model],
category_offline,
category_jev,
base_env,
base_url,
)
results.append(row)
finally:
guard.release()
if attempt == 1 and row["error_class"] == "infra":
attempt += 1
continue
break
return results
def run_pilot(workers: int = 3, only: str | None = None, models: list[str] | None = None) -> list[dict[str, object]]:
if workers < 1:
raise ValueError("workers must be positive")
tasks_payload = _read_json(ROOT / "tasks.json")
model_payload = _read_json(ROOT / "poc" / "pilot_models.json")
if not isinstance(tasks_payload, dict) or not isinstance(model_payload, dict):
raise ValueError("invalid task or model manifest")
tasks = tasks_payload["tasks"]
solvers = model_payload["solvers"]
model_order = ["haiku-4.5", "sonnet-5", "opus-5"]
if models is not None and (not models or any(model not in model_order for model in models)):
raise ValueError("models must be selected from the pinned solver list")
selected_models = [model for model in model_order if models is None or model in models]
if not selected_models or any(model not in solvers for model in selected_models):
raise ValueError("models must be selected from the pinned solver list")
selected_tasks = [task for task in tasks if only is None or task["instance_id"] == only]
if only is not None and not selected_tasks:
raise ValueError(f"unknown task ID: {only}")
labels = _labels_by_task()
base_env = os.environ.copy()
base_env.pop("LITELLM_GATEWAY_ADMIN_KEY", None)
base_env["LITELLM_PROXY_API_KEY"] = _read_pilot_key()
base_env["PYTHONPATH"] = str(ROOT)
base_env["MSWEA_MODEL_RETRY_STOP_AFTER_ATTEMPT"] = "1"
stop_path = ROOT / "dispatch_stops.jsonl"
guard = DispatchGuard(stop_path)
base_url = os.environ.get("LITELLM_GATEWAY_BASE_URL", "https://gateway.litellm-sandbox.ai").rstrip("/")
results = []
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = {
executor.submit(
_run_task,
task["instance_id"],
task,
selected_models,
solvers,
labels,
guard,
base_url,
base_env,
): task["instance_id"]
for task in selected_tasks
}
for future in as_completed(futures):
results.extend(future.result())
return results
def main() -> None:
parser = argparse.ArgumentParser(description="Run the frozen capability-routing pilot matrix")
parser.add_argument("--workers", type=int, default=3)
parser.add_argument("--only")
parser.add_argument("--models", nargs="+")
arguments = parser.parse_args()
results = run_pilot(arguments.workers, arguments.only, arguments.models)
print(
json.dumps(
[
{
"task_id": row["task_id"],
"model": row["model"],
"attempt": row["attempt"],
"error_class": row["error_class"],
"success": row["success"],
"cost_usd": row["cost_usd"],
}
for row in results
],
indent=2,
)
)
if __name__ == "__main__":
main()

View file

@ -0,0 +1,44 @@
{
"frozen_at": "2026-10-20T23:10Z",
"dataset": "princeton-nlp/SWE-bench",
"split": "dev",
"seed": "litellm-capability-routing-pilot-v1",
"rule": "sort dev ids by sha256(seed+id); take first 5 with public swebench eval image",
"tasks": [
{
"instance_id": "pvlib__pvlib-python-1395",
"image": "swebench/sweb.eval.x86_64.pvlib_1776_pvlib-python-1395:latest",
"exists": true,
"size": 1692788107,
"digest": "sha256:c7b51f4978926d54a788c54ac4d96e6c4c4e2cd2ee62736b30e91c9fd0891f10"
},
{
"instance_id": "sqlfluff__sqlfluff-884",
"image": "swebench/sweb.eval.x86_64.sqlfluff_1776_sqlfluff-884:latest",
"exists": true,
"size": 1038961212,
"digest": "sha256:ca28dab1d331c1c7e4908784576a7e1f3b11387dbaa4a8db466ce42a0eda6bec"
},
{
"instance_id": "pylint-dev__astroid-1614",
"image": "swebench/sweb.eval.x86_64.pylint-dev_1776_astroid-1614:latest",
"exists": true,
"size": 977627664,
"digest": "sha256:7206ea531b6c8b02f7a8e1760d732e4651f95307781f12f55e4b88e982610159"
},
{
"instance_id": "pvlib__pvlib-python-1165",
"image": "swebench/sweb.eval.x86_64.pvlib_1776_pvlib-python-1165:latest",
"exists": true,
"size": 1720675327,
"digest": "sha256:71e033964894077384c0c27a04ae0271e28eaad35602aa7570d40c683f1967d1"
},
{
"instance_id": "pydicom__pydicom-1031",
"image": "swebench/sweb.eval.x86_64.pydicom_1776_pydicom-1031:latest",
"exists": true,
"size": 1623110388,
"digest": "sha256:0b962e881760f4b8bf3cdbfc695624c90906198ccf1e2410cd639f6adc6e340b"
}
]
}

View file

@ -0,0 +1,463 @@
import csv
import json
import math
import os
from pathlib import Path
import pytest
from minisweagent.exceptions import LimitsExceeded
from poc.analyze import analyze
from poc.guarded_model import (
GuardedLitellmModel,
GuardedModelConfig,
budget_allows,
estimate_input_tokens,
gateway_cost,
reserve_usd,
)
from poc.label import parse_jev_response, parse_label_response
from poc.policy import (
PolicyConfig,
build_capability_map,
choose_model,
paired_comparison,
paired_lower_bound,
parse_category,
task_outcomes,
)
from poc.run_pilot import (
_utility_env,
classify_error,
global_budget_allows,
sum_call_costs,
sum_call_tokens,
)
def _record(
task_id: str,
model: str,
*,
error_class: str = "solver",
success: bool = False,
cost_usd: float = 0.0,
exit_status: str = "Submitted",
) -> dict[str, object]:
return {
"task_id": task_id,
"model": model,
"error_class": error_class,
"success": success,
"cost_usd": cost_usd,
"exit_status": exit_status,
"solve_latency_s": 1.0,
"eval_latency_s": 2.0,
}
def test_category_parsing_handles_aliases_garbage_and_non_strings() -> None:
assert parse_category("Repository Debugging") == "repository_debugging"
assert parse_category("Shell/Build Work") == "shell_build"
assert parse_category("Algorithms & Data Processing") == "algorithms_data"
assert parse_category("garbage") == "unknown"
assert parse_category(None) == "unknown"
assert parse_category(4) == "unknown"
def test_label_response_parser_accepts_fences_and_rejects_garbage() -> None:
assert parse_label_response(
'Answer: ```json\n{"label":"Feature Implementation","rationale":"Adds an API."}\n```'
) == ("feature_implementation", "Adds an API.", True)
assert parse_label_response('{"label":"nonsense","rationale":"x"}') == ("unknown", None, False)
assert parse_label_response("no json") == ("unknown", None, False)
assert parse_label_response('{"label":"unknown","rationale":"unclear"}') == ("unknown", "unclear", True)
def test_jev_response_parser_reads_structured_answer_and_fails_closed() -> None:
response = {
"answers": {
"capability": {
"choice": "shell_build",
"confidence": 0.82,
"probabilities": {
"shell_build": 0.82,
"repository_debugging": 0.18,
"not_a_category": 0.7,
},
}
}
}
assert parse_jev_response(response) == (
"shell_build",
0.82,
{"shell_build": 0.82, "repository_debugging": 0.18},
)
assert parse_jev_response({"answers": {"capability": {"choice": "unknown"}}}) == (
"unknown",
None,
{},
)
assert parse_jev_response({"answers": {"capability": {"confidence": 0.5}}}) == (None, None, {})
assert parse_jev_response({"answers": {"capability": {"choice": "unrecognized"}}}) == (None, None, {})
assert parse_jev_response("malformed") == (None, None, {})
def test_task_outcomes_cluster_retries_and_keep_all_attempt_costs() -> None:
records = [
_record("task-a", "haiku-4.5", error_class="infra", cost_usd=0.1),
_record("task-a", "haiku-4.5", error_class="infra", cost_usd=0.2),
_record("task-a", "haiku-4.5", success=True, cost_usd=0.4),
_record("task-b", "haiku-4.5", error_class="infra", cost_usd=0.5),
]
outcomes = task_outcomes(records)
task_a_outcome = outcomes[("task-a", "haiku-4.5")]
assert {
"graded": task_a_outcome["graded"],
"success": task_a_outcome["success"],
"attempts": task_a_outcome["attempts"],
} == {
"graded": True,
"success": True,
"attempts": 3,
}
assert task_a_outcome["cost_usd"] == pytest.approx(0.7)
assert outcomes[("task-b", "haiku-4.5")] == {
"graded": False,
"success": None,
"cost_usd": 0.5,
"attempts": 1,
}
capability_map = build_capability_map(
records,
{"task-a": "repository_debugging", "task-b": "repository_debugging"},
["haiku-4.5", "opus-5"],
PolicyConfig(),
)
assert capability_map["repository_debugging"]["haiku-4.5"]["n"] == 1
assert capability_map["repository_debugging"]["haiku-4.5"]["s"] == 1
assert capability_map["repository_debugging"]["haiku-4.5"]["posterior_mean"] == pytest.approx(2 / 3)
assert capability_map["repository_debugging"]["haiku-4.5"]["mean_cost_per_task_usd"] == pytest.approx(0.6)
assert capability_map["repository_debugging"]["haiku-4.5"]["graded_tasks"] == ["task-a"]
assert capability_map["repository_debugging"]["haiku-4.5"]["task_success"] == {"task-a": True}
task_cost = capability_map["repository_debugging"]["haiku-4.5"]["task_cost"]
assert task_cost["task-a"] == pytest.approx(0.7)
assert task_cost["task-b"] == pytest.approx(0.5)
def test_task_outcomes_reject_multiple_graded_attempts() -> None:
records = [
_record("task-a", "haiku-4.5", success=False),
_record("task-a", "haiku-4.5", success=True),
]
with pytest.raises(ValueError, match=r"duplicate graded attempts for \('task-a', 'haiku-4.5'\): single-attempt pilot"):
task_outcomes(records)
def test_paired_lower_bound_matches_one_sided_hoeffding_formula() -> None:
expected = -math.sqrt(2 * math.log(20) / 1000)
assert PolicyConfig().alpha == 0.05
assert paired_lower_bound([0] * 1000, 0.05) == pytest.approx(expected, abs=1e-12)
with pytest.raises(ValueError, match="cannot be empty"):
paired_lower_bound([], 0.05)
def test_model_selection_falls_back_for_failures_unknown_and_sparse_evidence() -> None:
cfg = PolicyConfig(min_paired_tasks=10)
assert choose_model(None, {}, cfg).reason == "fallback: classifier failure"
assert choose_model("unknown", {}, cfg).model == "opus-5"
assert choose_model("repository_debugging", {}, cfg).reason == (
"fallback: no profile for category repository_debugging"
)
sparse_map = build_capability_map(
[
_record(f"task-{index}", model, success=index < successes, cost_usd=cost)
for index in range(9)
for model, successes, cost in (("opus-5", 9, 0.1), ("haiku-4.5", 9, 0.01))
],
{f"task-{index}": "repository_debugging" for index in range(9)},
["haiku-4.5", "opus-5"],
cfg,
)
assert choose_model("repository_debugging", sparse_map, cfg).model == "opus-5"
def test_model_selection_requires_paired_evidence() -> None:
cfg = PolicyConfig(min_paired_tasks=10, tolerance=0.05)
disjoint_records = [
_record(f"reference-{index}", "opus-5", success=False, cost_usd=0.1)
for index in range(12)
] + [
_record(f"candidate-{index}", "haiku-4.5", success=True, cost_usd=0.01)
for index in range(12)
]
disjoint_categories = {
**{f"reference-{index}": "repository_debugging" for index in range(12)},
**{f"candidate-{index}": "repository_debugging" for index in range(12)},
}
disjoint_map = build_capability_map(
disjoint_records,
disjoint_categories,
["haiku-4.5", "opus-5"],
cfg,
)
disjoint_decision = choose_model("repository_debugging", disjoint_map, cfg)
assert disjoint_decision.model == "opus-5"
assert disjoint_decision.reason == (
"fallback: no cheaper model with >= 10 paired tasks and paired lower bound >= -0.05 (alpha=0.05)"
)
both_success_records = [
_record(f"task-{index}", model, success=True, cost_usd=cost)
for index in range(12)
for model, cost in (("opus-5", 0.1), ("haiku-4.5", 0.01))
]
both_success_categories = {f"task-{index}": "repository_debugging" for index in range(12)}
both_success_map = build_capability_map(
both_success_records,
both_success_categories,
["haiku-4.5", "opus-5"],
cfg,
)
comparison = paired_comparison(
both_success_map["repository_debugging"]["haiku-4.5"],
both_success_map["repository_debugging"]["opus-5"],
cfg.alpha,
)
assert comparison["n"] == 12
assert comparison["mean_diff"] == 0
assert comparison["lower_bound"] == pytest.approx(-math.sqrt(2 * math.log(20) / 12))
assert choose_model("repository_debugging", both_success_map, cfg).model == "opus-5"
def test_model_selection_uses_paired_cost_and_chooses_cheapest_eligible() -> None:
cfg = PolicyConfig(min_paired_tasks=10, tolerance=0.1)
paired_records = [
_record(f"task-{index}", model, success=True, cost_usd=cost)
for index in range(1000)
for model, cost in (("opus-5", 0.1), ("haiku-4.5", 0.05))
]
paired_categories = {f"task-{index}": "repository_debugging" for index in range(1000)}
paired_map = build_capability_map(
paired_records,
paired_categories,
["haiku-4.5", "opus-5"],
cfg,
)
decision = choose_model("repository_debugging", paired_map, cfg)
assert decision.model == "haiku-4.5"
assert decision.reason == "cheapest model with paired lower bound within tolerance of reference"
multiple_candidate_records = [
_record(f"task-{index}", model, success=True, cost_usd=cost)
for index in range(1000)
for model, cost in (("opus-5", 0.1), ("haiku-4.5", 0.05), ("sonnet-5", 0.03))
]
multiple_candidate_map = build_capability_map(
multiple_candidate_records,
paired_categories,
["haiku-4.5", "sonnet-5", "opus-5"],
cfg,
)
assert choose_model("repository_debugging", multiple_candidate_map, cfg).model == "sonnet-5"
paired_cost_records = [
_record(f"paired-{index}", "opus-5", success=True, cost_usd=0.1)
for index in range(1000)
] + [
_record(f"paired-{index}", "haiku-4.5", success=True, cost_usd=0.11)
for index in range(1000)
] + [
_record(f"candidate-only-{index}", "haiku-4.5", success=True, cost_usd=0.0001)
for index in range(1000)
]
paired_cost_categories = {
**{f"paired-{index}": "repository_debugging" for index in range(1000)},
**{f"candidate-only-{index}": "repository_debugging" for index in range(1000)},
}
paired_cost_map = build_capability_map(
paired_cost_records,
paired_cost_categories,
["haiku-4.5", "opus-5"],
cfg,
)
assert paired_cost_map["repository_debugging"]["haiku-4.5"]["mean_cost_per_task_usd"] < (
paired_cost_map["repository_debugging"]["opus-5"]["mean_cost_per_task_usd"]
)
assert choose_model("repository_debugging", paired_cost_map, cfg).model == "opus-5"
def test_reservations_and_spend_guards_include_boundary(tmp_path: Path) -> None:
config = GuardedModelConfig(
model_name="litellm_proxy/test-model",
run_budget_usd=1.0,
worst_input_usd_per_token=0.000002,
output_usd_per_token=0.00001,
max_output_tokens=100,
call_log_path="/unused/calls.jsonl",
)
assert reserve_usd(100, config) == 0.00125
assert budget_allows(0.99875, 0.00125, 1.0)
assert not budget_allows(0.998751, 0.00125, 1.0)
assert global_budget_allows(13.0, 1, 15.0)
assert not global_budget_allows(13.0001, 1, 15.0)
assert _utility_env(
{
"LITELLM_PROXY_API_KEY": "proxy-test-value",
"LITELLM_GATEWAY_ADMIN_KEY": "admin-test-value",
"PATH": "/bin",
}
) == {"PATH": "/bin"}
guarded_model = GuardedLitellmModel(
model_name="litellm_proxy/anthropic/claude-haiku-4-5-20251001",
model_kwargs={"api_base": "https://gateway.litellm-sandbox.ai"},
run_budget_usd=0.0001,
worst_input_usd_per_token=1.25e-6,
output_usd_per_token=5e-6,
max_output_tokens=64,
call_log_path=str(tmp_path / "guard_refusal.jsonl"),
)
with pytest.raises(LimitsExceeded):
guarded_model.query([{"role": "user", "content": "Reply OK."}])
assert os.environ["MSWEA_MODEL_RETRY_STOP_AFTER_ATTEMPT"] == "1"
refusal = json.loads((tmp_path / "guard_refusal.jsonl").read_text())
assert refusal["event"] == "budget_guard_refused"
def test_gateway_and_run_cost_accounting_handles_missing_costs() -> None:
from poc.guarded_model import COST_HEADER
assert gateway_cost({"additional_headers": {COST_HEADER: "0.0042"}}) == 0.0042
assert gateway_cost({}) is None
calls = [
{
"event": "call",
"cost_usd": 0.1,
"prompt_tokens": 10,
"completion_tokens": 3,
"cache_read_input_tokens": 2,
"cache_creation_input_tokens": 1,
"reasoning_tokens": 1,
},
{"event": "call", "cost_usd": None, "prompt_tokens": 5},
{"event": "budget_guard_refused", "cost_usd": 50},
]
assert sum_call_costs(calls) == (0.1, True)
assert sum_call_tokens(calls) == {
"prompt": 15,
"completion": 3,
"cache_read": 2,
"cache_creation": 1,
"reasoning": 1,
}
def test_error_classification_distinguishes_infra_and_solver_failures() -> None:
assert classify_error("Submitted", 0, "patch", "resolved", True) == "none"
assert classify_error("Submitted", 0, "patch", "unresolved", True) == "solver"
assert classify_error("BudgetGuard", 1, "", "skipped_empty_patch", True) == "solver"
assert classify_error("RuntimeError", 1, "", "skipped_empty_patch", True, "Docker startup command failed") == "infra"
assert classify_error(None, 1, "", "skipped_empty_patch", False) == "infra"
assert classify_error("Submitted", 124, "patch", "unresolved", True) == "solver"
assert classify_error("Submitted", 1, "patch", "error", True) == "infra"
assert classify_error("ModuleNotFoundError", 1, "", "skipped_empty_patch", True) == "infra"
def _write_json(path: Path, payload: object) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(payload) + "\n")
def _write_jsonl(path: Path, rows: list[dict[str, object]]) -> None:
path.write_text("".join(json.dumps(row) + "\n" for row in rows))
def test_analysis_reports_fixed_model_costs_oracle_and_case_table(tmp_path: Path) -> None:
task_ids = ["task-a", "task-b"]
_write_json(tmp_path / "tasks.json", {"tasks": [{"instance_id": task_id} for task_id in task_ids]})
_write_json(
tmp_path / "poc" / "pilot_models.json",
{
"solvers": {
"haiku-4.5": {},
"sonnet-5": {},
"opus-5": {},
}
},
)
_write_json(
tmp_path / "labels.json",
[
{
"task_id": "task-a",
"offline": {"label": "repository_debugging", "latency_s": 0.1},
"jev": {"label": "repository_debugging", "latency_s": 0.2},
},
{
"task_id": "task-b",
"offline": {"label": "feature_implementation", "latency_s": 0.1},
"jev": {"label": "shell_build", "latency_s": 0.2},
},
],
)
_write_jsonl(
tmp_path / "runs.jsonl",
[
_record("task-a", "haiku-4.5", success=True, cost_usd=0.1),
_record("task-a", "haiku-4.5", error_class="infra", cost_usd=0.05),
_record("task-a", "sonnet-5", success=False, cost_usd=0.2),
_record("task-a", "opus-5", success=True, cost_usd=0.5),
_record("task-b", "haiku-4.5", error_class="infra", cost_usd=0.03),
_record("task-b", "sonnet-5", success=True, cost_usd=0.25),
_record("task-b", "opus-5", success=False, cost_usd=0.4),
],
)
_write_jsonl(
tmp_path / "costs_aux.jsonl",
[
{"kind": "offline_label", "cost_usd": 0.01},
{"kind": "jev_classifier", "cost_usd": 0.02},
{"kind": "smoke", "cost_usd": 0.003},
],
)
_write_json(
tmp_path / "preflight_probes.json",
[{"cost": "0.0006", "latency": 0.05}],
)
result = analyze(tmp_path)
haiku = result["fixed_models"]["haiku-4.5"]
assert haiku["solved"] == 1
assert haiku["total_tasks"] == 2
assert haiku["total_cost_usd"] == pytest.approx(0.18)
assert haiku["mean_cost_per_attempted_task_usd"] == pytest.approx(0.09)
assert haiku["cost_per_success_usd"] == pytest.approx(0.18)
assert [row["retrospective_oracle"] for row in result["cases"]] == ["haiku-4.5", "sonnet-5"]
assert all(row["jev_decision"] == "opus-5" for row in result["cases"])
assert "RETROSPECTIVE UPPER BOUND, not achieved routing" in result["report"]
assert "EXPLORATORY leave-one-task-out" in result["report"]
assert "task_success and task_cost keys" in result["report"]
assert "Paired comparisons by category:" in result["report"]
assert "| repository_debugging | haiku-4.5 | 1 | 0.000 | -2.448 | $0.150000 | $0.500000 |" in result["report"]
assert "No threshold was tuned" in result["report"]
assert "Spend: solver" in result["report"]
assert "total $1.563600" in result["report"]
assert "Offline label coverage: 2/2. Offline/Jev agreement: 1/2 comparable labels" in result["report"]
assert "Total recorded runtime: 21.65s" in result["report"]
assert (tmp_path / "capability_map.json").exists()
assert (tmp_path / "report.md").exists()
with (tmp_path / "case_table.csv").open(newline="") as file:
rows = list(csv.DictReader(file))
assert rows[0]["retrospective_oracle"] == "haiku-4.5"
assert rows[1]["sonnet-5_success"] == "True"
def test_input_estimate_covers_tool_overhead_and_tracks_observed_prompt_growth():
config = GuardedModelConfig(
model_name="litellm_proxy/test-model", run_budget_usd=1.0, worst_input_usd_per_token=0.000002,
output_usd_per_token=0.00001, max_output_tokens=64, call_log_path="unused.jsonl",
)
assert estimate_input_tokens(59, None, None, config) == 59 + config.fixed_overhead_tokens
assert estimate_input_tokens(9000, 5000, 7000, config) == 11000
assert estimate_input_tokens(5100, 5000, 5200, config) == 6124
assert estimate_input_tokens(4000, 5000, 7000, config) == 7000