benchmark_live.py

"""Live benchmark: runs real agent tasks through the ACO proxy and measures savings.

Sends coding, QA, and tool-use tasks to the proxy, recording cost/latency/success.
Compares: baseline (always frontier) vs ACO proxy.

Usage:
    uv run --with openai,httpx benchmark_live.py --proxy http://localhost:8080/v1
    uv run --with openai,httpx benchmark_live.py --provider openai  # direct, no proxy

Or via hf_jobs (tests the proxy itself):
    hf_jobs run --script benchmark_live.py --deps openai,httpx --hardware cpu-basic --timeout 2h
"""
import json, time, sys, os, argparse
from typing import List, Dict, Any
from dataclasses import dataclass, field

# ── Test tasks: diverse mix of easy, medium, hard ──

TASKS = [
    # Quick answer (should route to tier 1, no tools needed)
    {"id": "qa_01", "desc": "quick answer — trivia",
     "messages": [{"role": "user", "content": "What is the capital of France?"}],
     "expect_success": True, "expected_tier": 1},

    {"id": "qa_02", "desc": "quick answer — definition",
     "messages": [{"role": "user", "content": "Briefly explain what a Python decorator is."}],
     "expect_success": True, "expected_tier": 1},

    # Simple coding (should route to tier 2)
    {"id": "code_01", "desc": "simple coding — fizzbuzz",
     "messages": [{"role": "user", "content": "Write a Python function fizzbuzz(n) that prints numbers 1 to n, but prints 'Fizz' for multiples of 3, 'Buzz' for multiples of 5, and 'FizzBuzz' for multiples of both."}],
     "expect_success": True, "expected_tier": 2},

    {"id": "code_02", "desc": "simple coding — reverse string",
     "messages": [{"role": "user", "content": "Write a Python function to reverse a string without using [::-1]."}],
     "expect_success": True, "expected_tier": 2},

    # Tool-use task (needs search)
    {"id": "tool_01", "desc": "tool use — search",
     "messages": [{"role": "user", "content": "What is the current version of Python and when was it released? Search for the answer."}],
     "expect_success": True, "expected_tier": 2, "needs_tools": True},

    # Medium coding (should route to tier 2-3)
    {"id": "code_03", "desc": "medium coding — LRU cache",
     "messages": [{"role": "user", "content": "Implement an LRU (Least Recently Used) cache in Python with get(key) and put(key, value) methods. Both should be O(1)."}],
     "expect_success": True, "expected_tier": 3},

    # Research-style (should route to tier 3+)
    {"id": "research_01", "desc": "research — transformer architectures",
     "messages": [{"role": "user", "content": "Compare LoRA and QLoRA for fine-tuning large language models. What are the key differences in memory usage and performance?"}],
     "expect_success": True, "expected_tier": 3},

    # Document drafting
    {"id": "doc_01", "desc": "document drafting — email",
     "messages": [{"role": "user", "content": "Write a professional email to a client explaining that their project will be delayed by 2 weeks due to an unexpected dependency issue. Be diplomatic."}],
     "expect_success": True, "expected_tier": 2},

    # Long context (compression should kick in)
    {"id": "long_01", "desc": "long context — log analysis",
     "messages": [
         {"role": "system", "content": "You are a log analyzer. Analyze the provided server log and identify the root cause of the outage."},
         {"role": "user", "content": "Here is a server log. Find the root cause of the outage:\n\n" +
          "Traceback (most recent call last):\n  File \"app.py\", line 42, in handle_request\n" * 30 +
          "\nERROR: Connection to database 'prod-db' timed out after 30s\n" * 5 +
          "\n[the rest of the log is normal operations]"}
     ],
     "expect_success": True, "expected_tier": 2},

    # Edge case: ambiguous
    {"id": "edge_01", "desc": "ambiguous query",
     "messages": [{"role": "user", "content": "Help me with this thing."}],
     "expect_success": True, "expected_tier": 2},
]

# ── Runner ──

@dataclass
class BenchResult:
    task_id: str
    desc: str
    success: bool
    model: str
    tier: int
    cost: float
    input_tokens: int
    output_tokens: int
    latency_ms: float
    tool_gated: bool
    model_routed: bool
    compression_ratio: float
    error: str = ""


def run_benchmark(api_base: str, model: str, api_key: str = None,
                  max_tasks: int = None, timeout: int = 60) -> List[BenchResult]:
    """Run benchmark tasks through the given endpoint."""
    import openai
    client = openai.OpenAI(base_url=api_base, api_key=api_key or "no-key-needed")

    tasks = TASKS[:max_tasks] if max_tasks else TASKS
    results = []
    total_cost = 0.0

    for i, task in enumerate(tasks):
        print(f"\n[{i+1}/{len(tasks)}] {task['id']}: {task['desc']}")
        t_start = time.time()

        try:
            kwargs = {"model": model, "messages": task["messages"],
                      "max_tokens": 500, "temperature": 0.0,
                      "timeout": timeout}
            if task.get("needs_tools"):
                kwargs["tools"] = [{"type": "function", "function": {
                    "name": "web_search", "description": "Search the web for current information",
                    "parameters": {"type": "object", "properties": {
                        "query": {"type": "string"}}}}}}]

            response = client.chat.completions.create(**kwargs)
            latency = (time.time() - t_start) * 1000

            usage = getattr(response, "usage", None) if hasattr(response, "usage") else None
            input_tokens = getattr(usage, "prompt_tokens", 0)
            output_tokens = getattr(usage, "completion_tokens", 0)

            # ACO telemetry from response (if going through proxy)
            aco_cost = getattr(usage, "aco_cost_usd", None) if usage else None
            aco_model = getattr(usage, "aco_model", None) if usage else None
            aco_tier = getattr(usage, "aco_tier", None) if usage else None
            aco_cache = getattr(usage, "aco_cache_hit_tokens", 0) if usage else 0
            aco_compress = getattr(usage, "aco_compression_ratio", 1.0) if usage else 1.0
            aco_tool_gated = getattr(usage, "aco_tool_gated", False) if usage else False

            used_model = aco_model or model
            used_tier = aco_tier or 0

            # Check response
            content = response.choices[0].message.content if response.choices else ""
            success = bool(content and len(content) > 10)

            result = BenchResult(
                task_id=task["id"],
                desc=task["desc"],
                success=success,
                model=used_model,
                tier=used_tier,
                cost=aco_cost or 0.0,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                latency_ms=round(latency, 1),
                tool_gated=aco_tool_gated,
                model_routed=used_model != model,
                compression_ratio=round(aco_compress, 2) if aco_compress else 1.0,
            )
            results.append(result)
            total_cost += result.cost
            print(f"  → model={used_model} tier={used_tier} cost=${result.cost:.5f} "
                  f"lat={latency:.0f}ms {'✓' if success else '✗'}")

        except Exception as e:
            latency = (time.time() - t_start) * 1000
            result = BenchResult(
                task_id=task["id"], desc=task["desc"],
                success=False, model=model, tier=0, cost=0.0,
                input_tokens=0, output_tokens=0, latency_ms=round(latency, 1),
                tool_gated=False, model_routed=False, compression_ratio=1.0,
                error=str(e)[:200],
            )
            results.append(result)
            print(f"  → ERROR: {str(e)[:100]}")

    return results


def print_report(results: List[BenchResult], title: str):
    """Print a benchmark report."""
    n = len(results)
    success_n = sum(1 for r in results if r.success)
    total_cost = sum(r.cost for r in results)
    avg_lat = sum(r.latency_ms for r in results) / max(n, 1)
    tier_counts = {}
    for r in results:
        tier_counts[r.tier] = tier_counts.get(r.tier, 0) + 1

    print(f"\n{'='*60}")
    print(f"  {title}")
    print(f"{'='*60}")
    print(f"  Tasks: {n} | Success: {success_n} ({success_n/n*100:.0f}%)")
    print(f"  Total cost: ${total_cost:.6f}")
    print(f"  Avg latency: {avg_lat:.0f}ms")
    print(f"  Tier distribution: {dict(sorted(tier_counts.items()))}")
    print(f"\n  Per-task:")
    for r in results:
        status = "✓" if r.success else "✗"
        print(f"    {r.task_id:<12} model={r.model:<22} tier={r.tier} "
              f"cost=${r.cost:.5f} lat={r.latency_ms:.0f}ms "
              f"gated={'Y' if r.tool_gated else '-'} {status}")
    return {
        "n": n, "success_rate": success_n / max(n, 1),
        "total_cost": total_cost, "avg_latency": avg_lat,
        "tier_distribution": dict(tier_counts),
    }


def main():
    parser = argparse.ArgumentParser(description="ACO Live Benchmark")
    parser.add_argument("--proxy", default=None, help="Proxy URL (e.g. http://localhost:8080/v1)")
    parser.add_argument("--provider", default="openai", help="Direct provider (skip proxy)")
    parser.add_argument("--model", default="gpt-5-mini", help="Model for direct calls")
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--tasks", type=int, default=None)
    args = parser.parse_args()

    if args.proxy:
        # Route through proxy
        print(f"Benchmarking through ACO proxy: {args.proxy}")
        results = run_benchmark(args.proxy, "gpt-5-mini", args.api_key, args.tasks)
        print_report(results, "ACO PROXY")
    else:
        # Direct provider call (baseline)
        base_urls = {"openai": "https://api.openai.com/v1",
                     "deepseek": "https://api.deepseek.com/v1"}
        base = base_urls.get(args.provider, "https://api.openai.com/v1")
        print(f"Benchmarking DIRECT to {args.provider}: {base}")
        results = run_benchmark(base, args.model, args.api_key or os.environ.get("OPENAI_API_KEY"), args.tasks)
        print_report(results, f"DIRECT ({args.provider})")

    # Save results
    with open("/tmp/aco_benchmark.json", "w") as f:
        json.dump([{
            "task_id": r.task_id, "desc": r.desc,
            "success": r.success, "model": r.model, "tier": r.tier,
            "cost": r.cost, "input_tokens": r.input_tokens,
            "output_tokens": r.output_tokens, "latency_ms": r.latency_ms,
            "tool_gated": r.tool_gated, "model_routed": r.model_routed,
            "compression_ratio": r.compression_ratio, "error": r.error,
        } for r in results], f, indent=2)
    print("\nResults saved to /tmp/aco_benchmark.json")


if __name__ == "__main__":
    main()