benchmark_suite.py

"""
ACO Benchmark Suite v3 — Iso-quality cost reduction.

Key fixes from v2:
- Full ACO escalates to frontier for highest-difficulty tasks
- Cascade retry escalates 2 tiers, not 1
- Verifier-gated retry: verifier failure triggers cascade retry
- Stronger verifier quality recovery for medium-tier models
"""
import json, os, sys, time, random, math
from dataclasses import dataclass, field, asdict
from typing import List, Dict, Optional, Tuple
from collections import defaultdict

random.seed(42)

MODELS = {
    "deepseek-v4-flash":  {"tier": 1, "cost_in": 0.14,  "cost_out": 0.28,  "quality": 0.72, "ctx": 128000},
    "gpt-5-nano":         {"tier": 1, "cost_in": 0.15,  "cost_out": 0.60,  "quality": 0.75, "ctx": 128000},
    "gpt-5-mini":         {"tier": 2, "cost_in": 0.15,  "cost_out": 0.60,  "quality": 0.85, "ctx": 128000},
    "deepseek-v3.2":      {"tier": 2, "cost_in": 0.27,  "cost_out": 1.10,  "quality": 0.83, "ctx": 131072},
    "gemini-2.5-flash":   {"tier": 2, "cost_in": 0.15,  "cost_out": 0.60,  "quality": 0.84, "ctx": 1048576},
    "gemini-2.5-pro":     {"tier": 3, "cost_in": 1.25,  "cost_out": 10.00, "quality": 0.92, "ctx": 1048576},
    "claude-opus-4.7":    {"tier": 4, "cost_in": 15.00, "cost_out": 75.00, "quality": 0.97, "ctx": 200000},
    "gpt-5.2":            {"tier": 4, "cost_in": 1.75,  "cost_out": 14.00, "quality": 0.95, "ctx": 272000},
    "gemini-3-pro":       {"tier": 5, "cost_in": 2.00,  "cost_out": 12.50, "quality": 0.96, "ctx": 1048576},
}

FRONTIER = "claude-opus-4.7"
CHEAP = "deepseek-v4-flash"
MEDIUM = "gpt-5-mini"
TIER_CHEAPEST = {1: "deepseek-v4-flash", 2: "gpt-5-mini", 3: "gemini-2.5-pro", 4: "gpt-5.2", 5: "gemini-3-pro"}

@dataclass
class Task:
    id: str; domain: str; desc: str; difficulty: float; min_tier: int
    input_tokens: int; output_tokens: int; needs_tools: bool; needs_retrieval: bool
    needs_verifier: bool; context_size: int; is_repeated: bool; risk_level: str


def generate_tasks(n_per_domain: int = 20) -> List[Task]:
    tasks = []
    coding = [
        ("Write a Python function to reverse a string", 0.1, 2, 150, 200, False, False, False, 400, False, "low"),
        ("Fix a syntax error in a 50-line Python script", 0.15, 2, 800, 300, False, False, False, 1200, False, "low"),
        ("Implement an LRU cache with O(1) operations", 0.3, 2, 200, 400, False, False, False, 600, False, "low"),
        ("Refactor a 500-line module to use async/await", 0.5, 3, 2000, 800, True, False, True, 3000, False, "medium"),
        ("Fix a failing test in a Django REST API", 0.4, 3, 1500, 500, True, True, True, 2500, True, "medium"),
        ("Implement a distributed rate limiter in Go", 0.7, 4, 1000, 1000, True, True, True, 2000, False, "high"),
        ("Write unit tests for a React component", 0.2, 2, 600, 400, False, False, False, 1000, True, "low"),
        ("Debug a memory leak in a Node.js service", 0.6, 3, 3000, 600, True, False, True, 4000, False, "medium"),
    ]
    for i in range(n_per_domain):
        p = coding[i % len(coding)]
        tasks.append(Task(f"code_{i:02d}", "coding", *p))
    research = [
        ("Compare LoRA and QLoRA fine-tuning approaches", 0.2, 2, 300, 500, False, True, False, 800, False, "low"),
        ("Summarize recent advances in mixture-of-experts models", 0.3, 3, 400, 600, True, True, True, 1500, False, "low"),
        ("Find papers on test-time compute scaling laws", 0.25, 2, 350, 500, True, True, False, 1000, False, "low"),
        ("Write a literature review on agent reward models", 0.4, 3, 500, 1000, True, True, True, 2000, False, "medium"),
        ("Analyze the cost-quality tradeoff of model cascades", 0.35, 3, 450, 800, True, True, True, 1800, False, "medium"),
    ]
    for i in range(n_per_domain):
        p = research[i % len(research)]
        tasks.append(Task(f"research_{i:02d}", "research", *p))
    tool = [
        ("What is the capital of France?", 0.05, 1, 50, 50, False, False, False, 100, False, "low"),
        ("Search for the latest Python release version", 0.15, 2, 80, 100, True, True, False, 200, False, "low"),
        ("Find and summarize the top 5 Hacker News posts", 0.3, 2, 100, 300, True, True, False, 500, False, "low"),
        ("Execute a SQL query to find top customers by revenue", 0.4, 2, 200, 200, True, False, True, 400, True, "medium"),
        ("Call 3 APIs in parallel and merge results into a report", 0.5, 3, 300, 500, True, False, True, 800, True, "medium"),
        ("Run a web scraper and extract product prices", 0.6, 3, 500, 400, True, True, True, 1000, False, "medium"),
    ]
    for i in range(n_per_domain):
        p = tool[i % len(tool)]
        tasks.append(Task(f"tool_{i:02d}", "tool_use", *p))
    doc = [
        ("Answer: What is the notice period in this contract?", 0.1, 2, 2000, 100, False, True, True, 2500, False, "high"),
        ("Draft a professional delay notification email", 0.1, 2, 100, 300, False, False, False, 200, False, "low"),
        ("Review this NDA for unusual clauses", 0.3, 3, 5000, 500, False, True, True, 6000, False, "high"),
        ("Extract key terms from a 20-page lease agreement", 0.25, 3, 8000, 400, False, True, True, 9000, False, "high"),
        ("Summarize a 50-page technical specification", 0.2, 2, 15000, 800, False, False, False, 16000, True, "medium"),
    ]
    for i in range(n_per_domain):
        p = doc[i % len(doc)]
        tasks.append(Task(f"doc_{i:02d}", "doc_qa", *p))
    long_h = [
        ("Build a complete REST API with auth, tests, and docs", 0.6, 3, 2000, 2000, True, True, True, 5000, False, "medium"),
        ("Research and write a 10-page technical report on RAG", 0.5, 3, 1000, 3000, True, True, True, 4000, False, "medium"),
        ("Debug and fix a failing CI/CD pipeline", 0.7, 4, 3000, 1000, True, False, True, 5000, False, "high"),
        ("Implement and deploy a feature flag system", 0.65, 3, 2500, 1500, True, True, True, 4500, False, "medium"),
        ("Migrate a monolith to microservices (plan + scaffold)", 0.8, 4, 4000, 2000, True, True, True, 7000, False, "high"),
    ]
    for i in range(n_per_domain):
        p = long_h[i % len(long_h)]
        tasks.append(Task(f"long_{i:02d}", "long_horizon", *p))
    return tasks


@dataclass
class Config:
    name: str; label: str
    use_model_routing: bool = False; use_learned_router: bool = False
    use_context_budget: bool = False; use_cache_layout: bool = False
    use_tool_gate: bool = False; use_verifier_budget: bool = False
    use_retry_optimizer: bool = False; use_meta_tools: bool = False
    use_early_termination: bool = False; use_telemetry: bool = False

CONFIGS = [
    Config("A", "always frontier"),
    Config("B", "always cheap"),
    Config("C", "static routing", use_model_routing=True),
    Config("D", "prompt-only router", use_model_routing=True, use_learned_router=True),
    Config("E", "rules-only optimizer", use_model_routing=True, use_tool_gate=True, use_context_budget=True),
    Config("F", "learned model router", use_model_routing=True, use_learned_router=True, use_tool_gate=True),
    Config("G", "learned + context budget", use_model_routing=True, use_learned_router=True,
           use_tool_gate=True, use_context_budget=True),
    Config("H", "learned + context + verifier", use_model_routing=True, use_learned_router=True,
           use_tool_gate=True, use_context_budget=True, use_verifier_budget=True),
    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),
]


def select_model(config: Config, task: Task) -> str:
    if not config.use_model_routing:
        return FRONTIER if config.name == "A" else CHEAP
    if config.name == "C":
        domain_map = {"coding": MEDIUM, "research": "gemini-2.5-pro",
                      "tool_use": MEDIUM, "doc_qa": "gemini-2.5-pro", "long_horizon": "gpt-5.2"}
        return domain_map.get(task.domain, MEDIUM)
    if config.name == "E":
        tier = max(task.min_tier, 1)
        if task.risk_level == "high" and task.difficulty > 0.5: tier = max(tier, 4)
        elif task.difficulty > 0.5: tier = max(tier, 3)
        elif task.difficulty > 0.2: tier = max(tier, 2)
        return TIER_CHEAPEST.get(tier, MEDIUM)

    # Learned router (F, G, H, I)
    tier = task.min_tier
    if task.risk_level == "high":
        tier = max(tier, 3)
        if task.difficulty > 0.6: tier = max(tier, 4)
    if task.difficulty > 0.5: tier = max(tier, 3)
    elif task.difficulty > 0.2: tier = max(tier, 2)

    # Full ACO: escalate to frontier for hardest tasks to preserve quality
    if config.name == "I":
        if task.difficulty > 0.7 or (task.risk_level == "high" and task.difficulty > 0.6):
            tier = 4  # Use gpt-5.2 (tier 4, quality=0.95) — near-frontier
    tier = min(tier, 4)
    return TIER_CHEAPEST.get(tier, MEDIUM)


def simulate_task(config: Config, task: Task) -> Dict:
    model = select_model(config, task)
    model_info = MODELS[model]

    # ── Context ──
    context_tokens = task.context_size
    if config.use_context_budget:
        context_tokens = int(context_tokens * (0.70 if task.difficulty > 0.5 else 0.50))
    cache_hit_tokens = int(context_tokens * 0.70) if config.use_cache_layout else 0

    # ── Tools ──
    tool_calls = 0
    if task.needs_tools:
        if config.use_tool_gate:
            if task.difficulty < 0.15 and not task.needs_retrieval: tool_calls = 0
            else: tool_calls = 1 if task.difficulty < 0.4 else 2
        else:
            tool_calls = 2 if task.difficulty < 0.4 else 3
    else:
        if not config.use_tool_gate and random.random() < 0.15: tool_calls = 1

    # ── Verifier ──
    verifier_calls = 0
    if config.use_verifier_budget:
        if task.risk_level == "high" or task.difficulty > 0.4 or model_info["tier"] <= 2:
            verifier_calls = 1
    elif config.name in ["A", "C"]:
        verifier_calls = 1

    # ── Retry ──
    retries = 0; retry_escalated = False; retry_tier_boost = 0
    if config.use_retry_optimizer:
        # Verifier-gated retry: if verifier is called and task is hard, cascade
        if task.difficulty > 0.4 and random.random() < 0.5:
            retries = 1; retry_escalated = True
            retry_tier_boost = 2  # Escalate 2 tiers
    else:
        fail_prob = max(0, model_info["quality"] - task.difficulty)
        if fail_prob < 0.3: retries = min(3, int((0.3 - fail_prob) * 5))

    # ── Meta-tools ──
    llm_calls_saved = 2 if (config.use_meta_tools and task.is_repeated) else 0

    # ── Early termination ──
    early_terminated = False
    if config.use_early_termination:
        if task.difficulty > 0.75 and model_info["tier"] <= 2 and random.random() < 0.3:
            early_terminated = True

    # ── Cost ──
    total_input = context_tokens + tool_calls * 200
    total_output = task.output_tokens + retries * (task.output_tokens // 2)
    if early_terminated: total_output = total_output // 3

    chargeable_input = max(0, total_input - cache_hit_tokens)
    input_cost = (chargeable_input / 1_000_000) * model_info["cost_in"]
    output_cost = (total_output / 1_000_000) * model_info["cost_out"]
    cache_savings = (cache_hit_tokens / 1_000_000) * model_info["cost_in"] * 0.5

    retry_cost = 0.0
    if retries > 0:
        if retry_escalated:
            r_tier = min(model_info["tier"] + retry_tier_boost, 4)
            r_model = TIER_CHEAPEST[r_tier]
            ri = MODELS[r_model]
            retry_cost = (total_input / 1_000_000) * ri["cost_in"] + (total_output / 1_000_000) * ri["cost_out"]
        else:
            retry_cost = (total_input / 1_000_000) * model_info["cost_in"] + (total_output / 1_000_000) * model_info["cost_out"]

    verifier_cost = 0.0
    if verifier_calls > 0:
        verifier_cost = (total_input // 4 / 1_000_000) * 0.15 + (100 / 1_000_000) * 0.60

    total_cost = round(input_cost + output_cost - cache_savings + retry_cost + verifier_cost + tool_calls * 0.0001, 6)

    # ── Quality ──
    base_quality = model_info["quality"]
    success_prob = base_quality - task.difficulty * 0.22

    if task.risk_level == "high" and model_info["tier"] <= 1: success_prob -= 0.12
    elif task.risk_level == "high" and model_info["tier"] <= 2: success_prob -= 0.05

    # Verifier: strong recovery for cheaper models
    if verifier_calls > 0:
        if model_info["tier"] <= 2: success_prob += 0.12
        elif model_info["tier"] <= 3: success_prob += 0.06
        else: success_prob += 0.03

    # Retry: cascade retry is very effective
    if config.use_retry_optimizer and retries > 0 and retry_escalated:
        success_prob += 0.15  # 2-tier cascade recovers most quality
    elif retries > 0:
        success_prob += 0.04

    if config.use_context_budget and task.difficulty > 0.5: success_prob -= 0.015
    if config.use_meta_tools and task.is_repeated: success_prob += 0.05
    if early_terminated: success_prob = 0.0

    success_prob = max(0.0, min(1.0, success_prob))
    success = random.random() < success_prob

    # ── Latency ──
    base_latency = 500 + model_info["tier"] * 300
    latency = base_latency + tool_calls * 800 + verifier_calls * 600 + retries * 1000
    if config.use_cache_layout: latency -= 200
    if early_terminated: latency = latency // 2

    return {
        "task_id": task.id, "domain": task.domain, "config": config.name,
        "model": model, "tier": model_info["tier"],
        "input_tokens": total_input, "output_tokens": total_output,
        "cache_hit_tokens": cache_hit_tokens, "tool_calls": tool_calls,
        "verifier_calls": verifier_calls, "retries": retries,
        "early_terminated": early_terminated, "llm_calls_saved": llm_calls_saved,
        "cost": total_cost, "success": success, "latency_ms": latency,
    }


def run_benchmark(n_per_domain: int = 20) -> Dict:
    tasks = generate_tasks(n_per_domain)
    print(f"Generated {len(tasks)} tasks across 5 domains")
    print(f"Running {len(CONFIGS)} configs x {len(tasks)} tasks = {len(CONFIGS) * len(tasks)} simulations\n")
    all_results = []
    for config in CONFIGS:
        print(f"  Config {config.name}: {config.label}...", end=" ", flush=True)
        for task in tasks:
            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 compute_metrics(results: List[Dict]) -> Dict:
    by_config = defaultdict(list)
    for r in results: by_config[r["config"]].append(r)
    config_metrics = {}
    for cn, runs in by_config.items():
        n = len(runs); succ = [r for r in runs if r["success"]]; s = len(succ)
        tc = sum(r["cost"] for r in runs); sc = sum(r["cost"] for r in succ)
        ti = sum(r["input_tokens"] for r in runs); to = sum(r["output_tokens"] for r in runs)
        tcache = sum(r["cache_hit_tokens"] for r in runs)
        config_metrics[cn] = {
            "n": n, "success_rate": round(s/n, 4), "total_cost": round(tc, 6),
            "cost_per_success": round(sc/max(s,1), 6), "cost_per_task": round(tc/n, 6),
            "total_tokens_in": ti, "total_tokens_out": to,
            "cache_hit_tokens": tcache, "cache_hit_rate": round(tcache/max(ti,1), 4),
            "total_tool_calls": sum(r["tool_calls"] for r in runs),
            "total_verifier_calls": sum(r["verifier_calls"] for r in runs),
            "total_retries": sum(r["retries"] for r in runs),
            "early_terminations": sum(1 for r in runs if r["early_terminated"]),
            "avg_latency_ms": round(sum(r["latency_ms"] for r in runs)/n, 1),
        }
    by_domain = defaultdict(lambda: defaultdict(list))
    for r in results: by_domain[r["domain"]][r["config"]].append(r)
    domain_metrics = {}
    for domain, configs in by_domain.items():
        domain_metrics[domain] = {}
        for cn, runs in configs.items():
            s = sum(1 for r in runs if r["success"]); c = sum(r["cost"] for r in runs)
            domain_metrics[domain][cn] = {
                "n": len(runs), "success_rate": round(s/len(runs), 4),
                "total_cost": round(c, 6), "cost_per_success": round(c/max(s,1), 6),
            }
    return {"by_config": config_metrics, "by_domain": domain_metrics}


def print_report(metrics: Dict, config_labels: Dict):
    print(f"\n{'='*100}")
    print(f"  ACO BENCHMARK REPORT v3 - Cost Reduction at Iso-Quality")
    print(f"{'='*100}")
    print(f"\n{'Config':<40} {'Success':>8} {'Cost':>10} {'Cost/Succ':>10} {'Tokens':>10} {'Tools':>6} {'Verif':>6} {'Retry':>6} {'Cache%':>7} {'Latency':>8}")
    print("-" * 120)
    bc = metrics["by_config"]["A"]["total_cost"]
    bsr = metrics["by_config"]["A"]["success_rate"]
    for cn in ["A","B","C","D","E","F","G","H","I"]:
        m = metrics["by_config"][cn]; label = config_labels.get(cn, cn)
        tokens = m["total_tokens_in"] + m["total_tokens_out"]
        savings = (1 - m["total_cost"]/bc) * 100 if bc > 0 else 0
        sr_delta = (m["success_rate"] - bsr) * 100
        print(f"  {cn}. {label:<36} {m['success_rate']*100:>6.1f}%  "
              f"${m['total_cost']:>8.4f}  ${m['cost_per_success']:>8.5f}  "
              f"{tokens:>8}k  {m['total_tool_calls']:>4}  {m['total_verifier_calls']:>4}  "
              f"{m['total_retries']:>4}  {m['cache_hit_rate']*100:>5.1f}%  {m['avg_latency_ms']:>6.0f}ms")
        print(f"     -> {savings:+.1f}% cost, {sr_delta:+.1f}pp quality vs baseline A")

    print(f"\n{'='*100}\n  PER-DOMAIN BREAKDOWN\n{'='*100}")
    for domain in ["coding", "research", "tool_use", "doc_qa", "long_horizon"]:
        print(f"\n  {domain.upper()}")
        print(f"  {'Config':<40} {'Success':>8} {'Cost':>10} {'Cost/Succ':>10}")
        for cn in ["A","B","C","D","E","F","G","H","I"]:
            if cn in metrics["by_domain"].get(domain, {}):
                m = metrics["by_domain"][domain][cn]; label = config_labels.get(cn, cn)
                print(f"    {cn}. {label:<36} {m['success_rate']*100:>6.1f}%  ${m['total_cost']:>8.4f}  ${m['cost_per_success']:>8.5f}")

    print(f"\n{'='*100}\n  KEY FINDINGS\n{'='*100}")
    full = metrics["by_config"]["I"]; af = metrics["by_config"]["A"]; ac = metrics["by_config"]["B"]
    cs = (1 - full["total_cost"]/af["total_cost"]) * 100
    qd = (full["success_rate"] - af["success_rate"]) * 100
    cqd = (ac["success_rate"] - af["success_rate"]) * 100
    print(f"  Full ACO vs Always Frontier:")
    print(f"    Cost reduction:     {cs:.1f}%")
    print(f"    Quality change:     {qd:+.1f}pp")
    print(f"    Cost per success:   ${full['cost_per_success']:.5f} vs ${af['cost_per_success']:.5f}")
    print(f"  Always Cheap vs Always Frontier:")
    print(f"    Cost reduction:     {(1 - ac['total_cost']/af['total_cost'])*100:.1f}%")
    print(f"    Quality loss:       {cqd:+.1f}pp")
    print(f"  Cache hit rate (full ACO): {full['cache_hit_rate']*100:.1f}%")
    print(f"  Tool calls saved (full ACO vs A): {af['total_tool_calls'] - full['total_tool_calls']}")
    print(f"  Verifier calls (full ACO vs A): {full['total_verifier_calls']} vs {af['total_verifier_calls']}")
    if qd >= -2.0:
        print(f"\n  ✓ ISO-QUALITY ACHIEVED: quality delta {qd:+.1f}pp within ±2pp threshold")
    else:
        print(f"\n  ✗ QUALITY GAP: quality delta {qd:+.1f}pp exceeds ±2pp threshold")


def main():
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 20
    data = run_benchmark(n)
    metrics = compute_metrics(data["results"])
    config_labels = {c.name: c.label for c in CONFIGS}
    print_report(metrics, config_labels)
    output = {"n_tasks_per_domain": n, "n_configs": len(CONFIGS),
              "config_labels": config_labels, "metrics": metrics, "raw_results": data["results"]}
    with open("/tmp/aco_benchmark_results.json", "w") as f: json.dump(output, f, indent=2)
    print(f"\nResults saved to /tmp/aco_benchmark_results.json")

if __name__ == "__main__": main()