ablation_study.py

"""
ACO Ablation Study + Cost-Quality Frontier Report.

10 ablations from spec:
 1. no model router
 2. no context budgeter
 3. no cache-aware layout
 4. no tool-use cost gate
 5. no verifier budgeter
 6. no retry optimizer
 7. no meta-tools
 8. no early termination
 9. no specialist models (force all routing through frontier-tier)
 10. no telemetry feedback

Each ablation removes one module from the full ACO and re-runs the benchmark.
Also runs all 9 baseline configs (A-I) for the frontier report.
"""
import json, sys, random, os
from dataclasses import dataclass, field, asdict
from collections import defaultdict
from typing import List, Dict

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) or ".")
import importlib.util

try:
    from benchmark_suite import (
        generate_tasks, simulate_task, Config, CONFIGS, MODELS,
        FRONTIER, CHEAP, MEDIUM, TIER_CHEAPEST,
        compute_metrics, run_benchmark
    )
except ImportError:
    import urllib.request
    url = "https://huggingface.co/narcolepticchicken/agent-cost-optimizer/resolve/main/benchmark_suite.py"
    path = "/tmp/benchmark_suite.py"
    if not os.path.exists(path):
        urllib.request.urlretrieve(url, path)
    sys.path.insert(0, "/tmp")
    from benchmark_suite import (
        generate_tasks, simulate_task, Config, CONFIGS, MODELS,
        FRONTIER, CHEAP, MEDIUM, TIER_CHEAPEST,
        compute_metrics, run_benchmark
    )

random.seed(42)

def make_ablation_configs() -> List[Config]:
    base = dict(
        use_model_routing=True, use_learned_router=True,
        use_context_budget=True, use_cache_layout=True,
        use_tool_gate=True, use_verifier_budget=True,
        use_retry_optimizer=True, use_meta_tools=True,
        use_early_termination=True, use_telemetry=True,
    )
    return [
        Config("abl1", "no model router", use_model_routing=False, use_learned_router=False,
               use_context_budget=True, use_cache_layout=True, use_tool_gate=True,
               use_verifier_budget=True, use_retry_optimizer=True, use_meta_tools=True,
               use_early_termination=True, use_telemetry=True),
        Config("abl2", "no context budgeter", **{**base, "use_context_budget": False}),
        Config("abl3", "no cache layout", **{**base, "use_cache_layout": False}),
        Config("abl4", "no tool gate", **{**base, "use_tool_gate": False}),
        Config("abl5", "no verifier budgeter", **{**base, "use_verifier_budget": False}),
        Config("abl6", "no retry optimizer", **{**base, "use_retry_optimizer": False}),
        Config("abl7", "no meta-tools", **{**base, "use_meta_tools": False}),
        Config("abl8", "no early termination", **{**base, "use_early_termination": False}),
        Config("abl9", "no specialist models",
               use_model_routing=True, use_learned_router=True,
               use_context_budget=True, use_cache_layout=True, use_tool_gate=True,
               use_verifier_budget=True, use_retry_optimizer=True, use_meta_tools=True,
               use_early_termination=True, use_telemetry=True),
        Config("abl10", "no telemetry feedback", **{**base, "use_telemetry": False}),
    ]

def run_ablation(n_per_domain: int = 20) -> Dict:
    """Run all baselines + 10 ablations."""
    tasks = generate_tasks(n_per_domain)
    full_aco = Config("I", "full ACO", use_model_routing=True, use_learned_router=True,
                      use_context_budget=True, use_cache_layout=True, use_tool_gate=True,
                      use_verifier_budget=True, use_retry_optimizer=True, use_meta_tools=True,
                      use_early_termination=True, use_telemetry=True)
    all_configs = list(CONFIGS) + make_ablation_configs()

    print(f"Generated {len(tasks)} tasks across 5 domains")
    print(f"Running {len(all_configs)} configs (9 baselines + 10 ablations) x {len(tasks)} tasks\n")

    all_results = []
    for config in all_configs:
        print(f"  {config.name}: {config.label}...", end=" ", flush=True)
        for task in tasks:
            if config.name == "abl9":
                result = simulate_task(config, task)
                if result["tier"] == 1:
                    result["model"] = "gpt-5-mini"; result["tier"] = 2
                    mi = MODELS["gpt-5-mini"]
                    result["cost"] = round(
                        (result["input_tokens"] / 1_000_000) * mi["cost_in"] +
                        (result["output_tokens"] / 1_000_000) * mi["cost_out"], 6)
                all_results.append(result)
            else:
                all_results.append(simulate_task(config, task))
        cr = [r for r in all_results if r["config"] == config.name]
        n = len(cr); s = sum(1 for r in cr if r["success"]); c = sum(r["cost"] for r in cr)
        print(f"{s}/{n} success, ${c:.4f} total")

    return {"tasks": [asdict(t) for t in tasks], "results": all_results}

def print_ablation_report(metrics: Dict, config_labels: Dict):
    print(f"\n{'='*100}")
    print(f"  ACO ABLATION REPORT - Which Modules Actually Save Money?")
    print(f"{'='*100}")

    full = metrics["by_config"]["I"]
    print(f"\n  Full ACO baseline: {full['success_rate']*100:.1f}% success, ${full['total_cost']:.4f} cost")

    print(f"\n{'Ablation':<40} {'Success':>8} {'Cost':>10} {'Cost Δ':>10} {'Quality Δ':>10} {'Verdict':>15}")
    print("-" * 100)

    verdicts = []
    for abl_name in [f"abl{i}" for i in range(1, 11)]:
        if abl_name not in metrics["by_config"]:
            continue
        m = metrics["by_config"][abl_name]
        label = config_labels.get(abl_name, abl_name)
        cost_delta = m["total_cost"] - full["total_cost"]
        cost_pct = (cost_delta / full["total_cost"]) * 100 if full["total_cost"] > 0 else 0
        quality_delta = (m["success_rate"] - full["success_rate"]) * 100

        if quality_delta < -3: verdict = "CRITICAL"
        elif cost_delta > 0 and quality_delta < -1: verdict = "HURTS QUALITY"
        elif cost_delta > 0.02: verdict = "SAVES MONEY"
        elif abs(cost_pct) < 2 and abs(quality_delta) < 1: verdict = "NOISE"
        elif cost_delta < 0 and quality_delta >= 0: verdict = "COST INCREASE"
        else: verdict = "MARGINAL"

        verdicts.append((abl_name, label, verdict, cost_pct, quality_delta))
        print(f"  {label:<38} {m['success_rate']*100:>6.1f}%  ${m['total_cost']:>8.4f}  "
              f"{cost_pct:>+8.1f}%  {quality_delta:>+8.1f}pp  {verdict:>15}")

    print(f"\n{'='*100}")
    print(f"  ABLATION SUMMARY")
    print(f"{'='*100}")

    critical = [v for v in verdicts if v[2] == "CRITICAL"]
    saves = [v for v in verdicts if v[2] == "SAVES MONEY"]
    noise = [v for v in verdicts if v[2] == "NOISE"]
    hurts = [v for v in verdicts if v[2] == "HURTS QUALITY"]
    marginal = [v for v in verdicts if v[2] == "MARGINAL"]

    print(f"\n  CRITICAL (removing causes >3pp quality loss):")
    for _, label, _, _, qd in critical: print(f"    - {label} ({qd:+.1f}pp)")
    if not critical: print("    (none)")

    print(f"\n  SAVES MONEY (removing increases cost):")
    for _, label, _, cp, _ in saves: print(f"    - {label} (+{cp:.1f}% cost without it)")
    if not saves: print("    (none)")

    print(f"\n  NOISE (removing has <2% cost, <1pp quality):")
    for _, label, _, _, _ in noise: print(f"    - {label}")
    if not noise: print("    (none)")

    print(f"\n  HURTS QUALITY (removing reduces quality but saves cost):")
    for _, label, _, _, qd in hurts: print(f"    - {label} ({qd:+.1f}pp)")
    if not hurts: print("    (none)")

    print(f"\n  MARGINAL (small effect either way):")
    for _, label, _, _, _ in marginal: print(f"    - {label}")
    if not marginal: print("    (none)")

def print_frontier_report(metrics: Dict, config_labels: Dict):
    print(f"\n{'='*100}")
    print(f"  COST-QUALITY FRONTIER")
    print(f"{'='*100}")

    # Use all configs that exist in metrics
    available = [k for k in ["A","B","C","D","E","F","G","H","I"] if k in metrics["by_config"]]
    configs = []
    for cn in available:
        m = metrics["by_config"][cn]
        configs.append((cn, config_labels.get(cn, cn), m["success_rate"], m["total_cost"]))

    configs.sort(key=lambda x: x[3])

    print(f"\n  {'Config':<40} {'Quality':>8} {'Cost':>10}  {'Visual':>35}")
    print("-" * 95)

    max_cost = max(c[3] for c in configs) if configs else 1
    for cn, label, sr, cost in configs:
        bar_len = int(sr * 30)
        bar = "█" * bar_len + "░" * (30 - bar_len)
        print(f"  {cn}. {label:<36} {sr*100:>6.1f}%  ${cost:>8.4f}  {bar}")

    pareto = []
    for i, (cn, label, sr, cost) in enumerate(configs):
        dominated = False
        for j, (cn2, label2, sr2, cost2) in enumerate(configs):
            if i != j and sr2 >= sr and cost2 <= cost and (sr2 > sr or cost2 < cost):
                dominated = True; break
        if not dominated: pareto.append((cn, label, sr, cost))

    print(f"\n  Pareto-optimal (no other config is both cheaper AND better):")
    for cn, label, sr, cost in pareto:
        print(f"    {cn}. {label}: {sr*100:.1f}% at ${cost:.4f}")

    frontier_q = max(c[2] for c in configs) if configs else 0
    print(f"\n  Iso-quality configs (within ±2pp of best quality {frontier_q*100:.1f}%):")
    for cn, label, sr, cost in configs:
        if sr >= frontier_q - 0.02:
            savings = (1 - cost / max_cost) * 100
            print(f"    {cn}. {label}: {sr*100:.1f}% at ${cost:.4f} ({savings:+.1f}% vs most expensive)")

def main():
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 20
    data = run_ablation(n)
    metrics = compute_metrics(data["results"])

    config_labels = {c.name: c.label for c in CONFIGS}
    for c in make_ablation_configs(): config_labels[c.name] = c.label

    print_ablation_report(metrics, config_labels)
    print_frontier_report(metrics, config_labels)

    output = {"n_tasks_per_domain": n, "metrics": metrics,
              "config_labels": config_labels, "raw_results": data["results"]}
    with open("/tmp/aco_ablation_results.json", "w") as f:
        json.dump(output, f, indent=2)
    print(f"\nResults saved to /tmp/aco_ablation_results.json")

if __name__ == "__main__": main()