From a0baa8067361907a8745eb5a1915d84cdc93e338 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:28:22 +0000 Subject: [PATCH] chore(cookbook): add standalone capability routing pilot POC Co-Authored-By: bot_apk --- cookbook/capability_routing_pilot/README.md | 11 + .../images_pinned.txt | 9 + .../images_pinned_upstream.txt | 5 + .../Dockerfile.pvlib_1776_pvlib-python-1165 | 2 + .../Dockerfile.pvlib_1776_pvlib-python-1395 | 2 + .../Dockerfile.pydicom_1776_pydicom-1031 | 2 + .../capability_routing_pilot/pip_freeze.txt | 125 ++++ .../capability_routing_pilot/poc/__init__.py | 0 .../capability_routing_pilot/poc/analyze.py | 419 ++++++++++++ .../poc/guarded_model.py | 121 ++++ .../capability_routing_pilot/poc/label.py | 338 ++++++++++ .../capability_routing_pilot/poc/pilot.yaml | 11 + .../poc/pilot_models.json | 14 + .../capability_routing_pilot/poc/policy.py | 163 +++++ .../capability_routing_pilot/poc/run_pilot.py | 597 ++++++++++++++++++ cookbook/capability_routing_pilot/tasks.json | 44 ++ .../tests/test_poc.py | 463 ++++++++++++++ 17 files changed, 2326 insertions(+) create mode 100644 cookbook/capability_routing_pilot/README.md create mode 100644 cookbook/capability_routing_pilot/images_pinned.txt create mode 100644 cookbook/capability_routing_pilot/images_pinned_upstream.txt create mode 100644 cookbook/capability_routing_pilot/overlays/Dockerfile.pvlib_1776_pvlib-python-1165 create mode 100644 cookbook/capability_routing_pilot/overlays/Dockerfile.pvlib_1776_pvlib-python-1395 create mode 100644 cookbook/capability_routing_pilot/overlays/Dockerfile.pydicom_1776_pydicom-1031 create mode 100644 cookbook/capability_routing_pilot/pip_freeze.txt create mode 100644 cookbook/capability_routing_pilot/poc/__init__.py create mode 100644 cookbook/capability_routing_pilot/poc/analyze.py create mode 100644 cookbook/capability_routing_pilot/poc/guarded_model.py create mode 100644 cookbook/capability_routing_pilot/poc/label.py create mode 100644 cookbook/capability_routing_pilot/poc/pilot.yaml create mode 100644 cookbook/capability_routing_pilot/poc/pilot_models.json create mode 100644 cookbook/capability_routing_pilot/poc/policy.py create mode 100644 cookbook/capability_routing_pilot/poc/run_pilot.py create mode 100644 cookbook/capability_routing_pilot/tasks.json create mode 100644 cookbook/capability_routing_pilot/tests/test_poc.py diff --git a/cookbook/capability_routing_pilot/README.md b/cookbook/capability_routing_pilot/README.md new file mode 100644 index 00000000000..9ab3822cd2a --- /dev/null +++ b/cookbook/capability_routing_pilot/README.md @@ -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 diff --git a/cookbook/capability_routing_pilot/images_pinned.txt b/cookbook/capability_routing_pilot/images_pinned.txt new file mode 100644 index 00000000000..e3bd5c9f2a3 --- /dev/null +++ b/cookbook/capability_routing_pilot/images_pinned.txt @@ -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 diff --git a/cookbook/capability_routing_pilot/images_pinned_upstream.txt b/cookbook/capability_routing_pilot/images_pinned_upstream.txt new file mode 100644 index 00000000000..0e091eebd48 --- /dev/null +++ b/cookbook/capability_routing_pilot/images_pinned_upstream.txt @@ -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 diff --git a/cookbook/capability_routing_pilot/overlays/Dockerfile.pvlib_1776_pvlib-python-1165 b/cookbook/capability_routing_pilot/overlays/Dockerfile.pvlib_1776_pvlib-python-1165 new file mode 100644 index 00000000000..ffffc0a9464 --- /dev/null +++ b/cookbook/capability_routing_pilot/overlays/Dockerfile.pvlib_1776_pvlib-python-1165 @@ -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 diff --git a/cookbook/capability_routing_pilot/overlays/Dockerfile.pvlib_1776_pvlib-python-1395 b/cookbook/capability_routing_pilot/overlays/Dockerfile.pvlib_1776_pvlib-python-1395 new file mode 100644 index 00000000000..487416eed95 --- /dev/null +++ b/cookbook/capability_routing_pilot/overlays/Dockerfile.pvlib_1776_pvlib-python-1395 @@ -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 diff --git a/cookbook/capability_routing_pilot/overlays/Dockerfile.pydicom_1776_pydicom-1031 b/cookbook/capability_routing_pilot/overlays/Dockerfile.pydicom_1776_pydicom-1031 new file mode 100644 index 00000000000..3e5dfd7c060 --- /dev/null +++ b/cookbook/capability_routing_pilot/overlays/Dockerfile.pydicom_1776_pydicom-1031 @@ -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 diff --git a/cookbook/capability_routing_pilot/pip_freeze.txt b/cookbook/capability_routing_pilot/pip_freeze.txt new file mode 100644 index 00000000000..a6046c8b8f3 --- /dev/null +++ b/cookbook/capability_routing_pilot/pip_freeze.txt @@ -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 diff --git a/cookbook/capability_routing_pilot/poc/__init__.py b/cookbook/capability_routing_pilot/poc/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/cookbook/capability_routing_pilot/poc/analyze.py b/cookbook/capability_routing_pilot/poc/analyze.py new file mode 100644 index 00000000000..5b8584a7c1c --- /dev/null +++ b/cookbook/capability_routing_pilot/poc/analyze.py @@ -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() diff --git a/cookbook/capability_routing_pilot/poc/guarded_model.py b/cookbook/capability_routing_pilot/poc/guarded_model.py new file mode 100644 index 00000000000..af3d8838510 --- /dev/null +++ b/cookbook/capability_routing_pilot/poc/guarded_model.py @@ -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") diff --git a/cookbook/capability_routing_pilot/poc/label.py b/cookbook/capability_routing_pilot/poc/label.py new file mode 100644 index 00000000000..59f0de1629b --- /dev/null +++ b/cookbook/capability_routing_pilot/poc/label.py @@ -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\": , \"rationale\": }. 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() diff --git a/cookbook/capability_routing_pilot/poc/pilot.yaml b/cookbook/capability_routing_pilot/poc/pilot.yaml new file mode 100644 index 00000000000..52113666375 --- /dev/null +++ b/cookbook/capability_routing_pilot/poc/pilot.yaml @@ -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 diff --git a/cookbook/capability_routing_pilot/poc/pilot_models.json b/cookbook/capability_routing_pilot/poc/pilot_models.json new file mode 100644 index 00000000000..d2360230e63 --- /dev/null +++ b/cookbook/capability_routing_pilot/poc/pilot_models.json @@ -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} +} diff --git a/cookbook/capability_routing_pilot/poc/policy.py b/cookbook/capability_routing_pilot/poc/policy.py new file mode 100644 index 00000000000..d1185d8f091 --- /dev/null +++ b/cookbook/capability_routing_pilot/poc/policy.py @@ -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()] diff --git a/cookbook/capability_routing_pilot/poc/run_pilot.py b/cookbook/capability_routing_pilot/poc/run_pilot.py new file mode 100644 index 00000000000..2a8e3d6daaa --- /dev/null +++ b/cookbook/capability_routing_pilot/poc/run_pilot.py @@ -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() diff --git a/cookbook/capability_routing_pilot/tasks.json b/cookbook/capability_routing_pilot/tasks.json new file mode 100644 index 00000000000..bdb7971904e --- /dev/null +++ b/cookbook/capability_routing_pilot/tasks.json @@ -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" + } + ] +} \ No newline at end of file diff --git a/cookbook/capability_routing_pilot/tests/test_poc.py b/cookbook/capability_routing_pilot/tests/test_poc.py new file mode 100644 index 00000000000..78289c818f4 --- /dev/null +++ b/cookbook/capability_routing_pilot/tests/test_poc.py @@ -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