Upload aco/benchmark.py

aco/benchmark.py

@@ -0,0 +1,482 @@
+"""ACO Benchmark Suite
+
+Evaluates the Agent Cost Optimizer across 5 task categories:
+A. Coding agent tasks
+B. Research agent tasks  
+C. Tool-use tasks
+D. Document/contract/QA tasks
+E. Long-horizon agent tasks
+
+For each, measures: task success, LLM calls, tool calls, retries, cost, latency,
+cache hit rate, context tokens, false-DONE rate.
+
+Baselines compared:
+A. always frontier model
+B. always cheap model
+C. static model routing
+D. prompt-only router
+
+Primary metric: cost reduction at equal or better task success.
+Secondary: cost per completed artifact.
+"""
+
+import json, time, hashlib
+from typing import Dict, List, Tuple, Optional, Any, Callable
+from dataclasses import dataclass, field
+from collections import defaultdict
+
+
+# =============================================================================
+# Core Metrics & Types
+# =============================================================================
+
+@dataclass
+class StepMetrics:
+    step_num: int
+    model_used: str
+    provider: str
+    input_tokens: int
+    output_tokens: int
+    reasoning_tokens: int = 0
+    cache_hit: bool = False
+    context_size: int = 0
+    tool_calls: int = 0
+    tool_failures: int = 0
+    verifier_calls: int = 0
+    retries: int = 0
+    latency_ms: float = 0.0
+    cost_usd: float = 0.0
+
+
+@dataclass
+class TaskResult:
+    task_id: str
+    task_type: str  # "coding", "research", "tool_use", "qa", "long_horizon"
+    success: bool
+    steps: List[StepMetrics] = field(default_factory=list)
+    total_llm_calls: int = 0
+    total_tool_calls: int = 0
+    total_retries: int = 0
+    total_verifier_calls: int = 0
+    total_cost_usd: float = 0.0
+    total_latency_ms: float = 0.0
+    false_done: bool = False
+    blocked: bool = False
+    artifact_hash: Optional[str] = None
+    error: Optional[str] = None
+
+
+@dataclass 
+class RunReport:
+    config_name: str
+    config_description: str
+    results: List[TaskResult] = field(default_factory=list)
+
+    @property
+    def success_rate(self) -> float:
+        if not self.results: return 0.0
+        return sum(1 for r in self.results if r.success) / len(self.results)
+
+    @property
+    def avg_cost_per_success(self) -> float:
+        successes = [r.total_cost_usd for r in self.results if r.success]
+        if not successes: return float('inf')
+        return sum(successes) / len(successes)
+
+    @property
+    def avg_cost_per_task(self) -> float:
+        if not self.results: return 0.0
+        return sum(r.total_cost_usd for r in self.results) / len(self.results)
+
+    @property
+    def avg_llm_calls(self) -> float:
+        if not self.results: return 0.0
+        return sum(r.total_llm_calls for r in self.results) / len(self.results)
+
+    @property
+    def avg_tool_calls(self) -> float:
+        if not self.results: return 0.0
+        return sum(r.total_tool_calls for r in self.results) / len(self.results)
+
+    @property
+    def avg_retries(self) -> float:
+        if not self.results: return 0.0
+        return sum(r.total_retries for r in self.results) / len(self.results)
+
+    @property
+    def false_done_rate(self) -> float:
+        if not self.results: return 0.0
+        return sum(1 for r in self.results if r.false_done) / len(self.results)
+
+    def to_dict(self) -> Dict:
+        return {
+            "config_name": self.config_name,
+            "success_rate": self.success_rate,
+            "avg_cost_per_success": self.avg_cost_per_success,
+            "avg_cost_per_task": self.avg_cost_per_task,
+            "avg_llm_calls": self.avg_llm_calls,
+            "avg_tool_calls": self.avg_tool_calls,
+            "avg_retries": self.avg_retries,
+            "false_done_rate": self.false_done_rate,
+            "n_tasks": len(self.results),
+        }
+
+
+# =============================================================================
+# Benchmark Tasks (Synthetic to start — real traces follow)
+# =============================================================================
+
+CODING_TASKS = [
+    {"id": "code_1", "text": "Write a Python function to sort a list of dicts by a key",
+     "expected_tools": ["write_file", "run_test"], "difficulty": "easy"},
+    {"id": "code_2", "text": "Fix the race condition in this multithreaded queue",
+     "expected_tools": ["read_file", "grep", "write_file", "run_test"], "difficulty": "hard"},
+    {"id": "code_3", "text": "Add input validation to all API endpoints",
+     "expected_tools": ["grep", "read_file", "write_file", "run_test"], "difficulty": "medium"},
+    {"id": "code_4", "text": "Migrate from SQLAlchemy 1.4 to 2.0 async API",
+     "expected_tools": ["read_file", "grep", "write_file", "run_migration"], "difficulty": "hard"},
+    {"id": "code_5", "text": "Write unit tests for the User model",
+     "expected_tools": ["read_file", "write_file", "run_test"], "difficulty": "medium"},
+]
+
+RESEARCH_TASKS = [
+    {"id": "research_1", "text": "What's the SOTA for cost-aware LLM routing?",
+     "expected_tools": ["search", "read_paper", "synthesize"], "difficulty": "medium"},
+    {"id": "research_2", "text": "Find all papers about speculative decoding applied to agents",
+     "expected_tools": ["search", "citation_graph", "synthesize"], "difficulty": "hard"},
+    {"id": "research_3", "text": "How does Anthropic's prompt caching work in detail?",
+     "expected_tools": ["search", "read_docs", "synthesize"], "difficulty": "easy"},
+    {"id": "research_4", "text": "Survey test-time compute allocation methods 2023-2025",
+     "expected_tools": ["search", "read_paper", "citation_graph", "synthesize"], "difficulty": "hard"},
+    {"id": "research_5", "text": "What datasets exist for training agent cost routers?",
+     "expected_tools": ["search", "read_paper", "synthesize"], "difficulty": "medium"},
+]
+
+TOOL_USE_TASKS = [
+    {"id": "tool_1", "text": "What's the weather in Paris today?",
+     "tools_needed": False, "difficulty": "easy"},
+    {"id": "tool_2", "text": "Find me the top 5 restaurants in Tokyo and their ratings",
+     "tools_needed": True, "difficulty": "medium"},
+    {"id": "tool_3", "text": "Compare the GDP of France and Germany over the last decade",
+     "tools_needed": True, "difficulty": "medium"},
+    {"id": "tool_4", "text": "Translate 'hello' to Japanese",
+     "tools_needed": False, "difficulty": "easy"},
+    {"id": "tool_5", "text": "Calculate compound interest on $10,000 at 5% over 20 years",
+     "tools_needed": False, "difficulty": "easy"},
+]
+
+QA_TASKS = [
+    {"id": "qa_1", "text": "What is the capital of France?", "needs_verifier": False, "difficulty": "easy"},
+    {"id": "qa_2", "text": "Explain the difference between TCP and UDP",
+     "needs_verifier": False, "difficulty": "easy"},
+    {"id": "qa_3", "text": "What are the implications of Article 22 of GDPR for automated decision-making?",
+     "needs_verifier": True, "difficulty": "hard"},
+    {"id": "qa_4", "text": "Review this contract for liability clauses",
+     "needs_verifier": True, "difficulty": "hard"},
+    {"id": "qa_5", "text": "Summarize the key findings from the IPCC AR6 report",
+     "needs_verifier": True, "difficulty": "medium"},
+]
+
+LONG_HORIZON_TASKS = [
+    {"id": "long_1", "text": "Set up a complete CI/CD pipeline for a Python project",
+     "expected_steps": 8, "difficulty": "hard"},
+    {"id": "long_2", "text": "Research, design, and implement a caching layer for the API",
+     "expected_steps": 6, "difficulty": "hard"},
+    {"id": "long_3", "text": "Write a README, add tests, and set up pre-commit hooks",
+     "expected_steps": 5, "difficulty": "medium"},
+    {"id": "long_4", "text": "Migrate the database schema, update models, and add rollback scripts",
+     "expected_steps": 7, "difficulty": "hard"},
+    {"id": "long_5", "text": "Audit the codebase for security vulnerabilities and fix the top 3",
+     "expected_steps": 6, "difficulty": "hard"},
+]
+
+ALL_TASKS = {
+    "coding": CODING_TASKS,
+    "research": RESEARCH_TASKS,
+    "tool_use": TOOL_USE_TASKS,
+    "qa": QA_TASKS,
+    "long_horizon": LONG_HORIZON_TASKS,
+}
+
+
+# =============================================================================
+# Simulated Agent Harness (for benchmarking w/o real API keys)
+# =============================================================================
+
+PRICING = {
+    "gpt-4.1-mini":   {"input": 0.15, "output": 0.60},
+    "gpt-4o":          {"input": 2.50, "output": 10.00},
+    "claude-sonnet-4": {"input": 3.00, "output": 15.00},
+    "claude-haiku":    {"input": 0.80, "output": 4.00},
+    "llama-3.1-8b":    {"input": 0.06, "output": 0.06},
+    "qwen-2.5-7b":     {"input": 0.07, "output": 0.07},
+    "gemini-flash":    {"input": 0.075, "output": 0.30},
+    "local-tiny":      {"input": 0.001, "output": 0.001},
+}
+
+TIER_MODELS = {
+    "cheap":   "llama-3.1-8b",
+    "medium":  "gpt-4.1-mini",
+    "frontier": "claude-sonnet-4",
+}
+
+
+class SimulatedAgent:
+    """
+    Simulated agent with configurable success probabilities per tier.
+    Used for benchmark without real API keys.
+    """
+    
+    def __init__(self, success_rates: Optional[Dict[str, Dict[str, float]]] = None):
+        """
+        success_rates: {tier: {"coding": prob, "research": prob, ...}}
+        Defaults calibrated to realistic distributions.
+        """
+        self.success_rates = success_rates or {
+            "cheap":   {"coding": 0.55, "research": 0.40, "tool_use": 0.50, "qa": 0.60, "long_horizon": 0.30},
+            "medium":  {"coding": 0.78, "research": 0.65, "tool_use": 0.75, "qa": 0.85, "long_horizon": 0.55},
+            "frontier":{"coding": 0.92, "research": 0.85, "tool_use": 0.90, "qa": 0.95, "long_horizon": 0.80},
+        }
+        import random
+        self.rng = random.Random(42)
+    
+    def run_task(self, task: Dict, tier: str, max_steps: int = 8,
+                 use_tools: bool = True, use_verifier: bool = False) -> TaskResult:
+        """
+        Simulate running a task with given tier.
+        Returns TaskResult with simulated metrics.
+        """
+        task_type = task.get("type", "coding")
+        success_prob = self.success_rates[tier].get(task_type, 0.7)
+        difficulty = task.get("difficulty", "medium")
+        
+        # Adjust success prob by difficulty
+        diff_mult = {"easy": 1.15, "medium": 1.0, "hard": 0.80}
+        success_prob = min(0.99, success_prob * diff_mult[difficulty])
+        
+        model_id = TIER_MODELS[tier]
+        pricing = PRICING[model_id]
+        
+        # Simulated step count
+        if difficulty == "easy":
+            steps = self.rng.randint(1, 3)
+        elif difficulty == "medium":
+            steps = self.rng.randint(2, 5)
+        else:
+            steps = self.rng.randint(3, min(max_steps, 8))
+        
+        result = TaskResult(
+            task_id=task["id"],
+            task_type=task_type,
+            success=False,
+        )
+        
+        for i in range(steps):
+            in_tokens = self.rng.randint(200, 3000)
+            out_tokens = self.rng.randint(50, 1500)
+            tool_calls = self.rng.randint(0, 3) if use_tools else 0
+            tool_failures = self.rng.randint(0, 1) if tool_calls > 0 else 0
+            verifier_calls = 1 if use_verifier else 0
+            retries = self.rng.randint(0, 2) if tool_failures > 0 else 0
+            
+            cost = (in_tokens / 1_000_000 * pricing["input"] + 
+                   out_tokens / 1_000_000 * pricing["output"])
+            
+            step = StepMetrics(
+                step_num=i+1,
+                model_used=model_id,
+                provider="simulated",
+                input_tokens=in_tokens,
+                output_tokens=out_tokens,
+                context_size=in_tokens + out_tokens,
+                tool_calls=tool_calls,
+                tool_failures=tool_failures,
+                verifier_calls=verifier_calls,
+                retries=retries,
+                latency_ms=in_tokens * 2 + out_tokens * 10,
+                cost_usd=cost,
+            )
+            result.steps.append(step)
+            result.total_cost_usd += cost
+            result.total_latency_ms += step.latency_ms
+            result.total_llm_calls += 1
+            result.total_tool_calls += tool_calls
+            result.total_retries += retries
+            result.total_verifier_calls += verifier_calls
+        
+        result.success = self.rng.random() < (success_prob * (1 - 0.03 * result.total_retries))
+        
+        if result.success:
+            result.artifact_hash = hashlib.md5(task["text"].encode()).hexdigest()[:16]
+        
+        return result
+
+
+# =============================================================================
+# Benchmark Runner
+# =============================================================================
+
+class BenchmarkRunner:
+    """Runs all task types against all baseline configs and produces comparison."""
+    
+    def __init__(self, agent: SimulatedAgent):
+        self.agent = agent
+        self.reports: Dict[str, RunReport] = {}
+    
+    def run_config(self, name: str, config: Callable[[Dict], Tuple[str, bool, bool]],
+                   description: str = ""):
+        """
+        Run a configuration across all tasks.
+        
+        config(task): returns (tier, use_tools, use_verifier)
+        """
+        report = RunReport(name, description)
+        
+        for task_type, tasks in ALL_TASKS.items():
+            for task in tasks:
+                task["type"] = task_type
+                tier, use_tools, use_verifier = config(task)
+                result = self.agent.run_task(task, tier, use_tools=use_tools. use_verifier=use_verifier)
+                report.results.append(result)
+        
+        self.reports[name] = report
+        return report
+    
+    def run_all_baselines(self):
+        """Run all 4 baseline configurations."""
+        
+        # A. Always Frontier
+        print("Running A: Always Frontier...")
+        self.run_config(
+            "A_always_frontier",
+            lambda t: ("frontier", True, True),
+            "Always use frontier model with all tools and verifier"
+        )
+        
+        # B. Always Cheap
+        print("Running B: Always Cheap...")
+        self.run_config(
+            "B_always_cheap",
+            lambda t: ("cheap", False, False),
+            "Always use cheap model without tools or verifier"
+        )
+        
+        # C. Static Routing
+        print("Running C: Static Routing...")
+        self.run_config(
+            "C_static_routing",
+            lambda t: self._static_route(t),
+            "Static rules: easy→cheap, medium→medium, hard→frontier"
+        )
+        
+        # D. Prompt-Only Router
+        print("Running D: Prompt-Only Router...")
+        self.run_config(
+            "D_prompt_router",
+            lambda t: self._prompt_route(t),
+            "Heuristic router based on task type + difficulty keywords"
+        )
+        
+        # E. ACO v2 (trained router)
+        # Requires loaded models - placeholder
+    
+    def _static_route(self, task: Dict) -> Tuple[str, bool, bool]:
+        diff = task.get("difficulty", "medium")
+        if diff == "easy": return ("cheap", False, False)
+        elif diff == "medium": return ("medium", True, False)
+        else: return ("frontier", True, True)
+    
+    def _prompt_route(self, task: Dict) -> Tuple[str, bool, bool]:
+        text = task["text"].lower()
+        if any(kw in text for kw in ["fix", "debug", "migrate", "critical", "vulnerability"]):
+            return ("frontier", True, True)
+        elif any(kw in text for kw in ["test", "add", "write", "survey", "find"]):
+            return ("medium", True, False)
+        else:
+            return ("cheap", False, False)
+    
+    def compare(self) -> Dict:
+        """Produce comparison table."""
+        baseline_label = "A_always_frontier"
+        if baseline_label not in self.reports:
+            return {"error": "Run baselines first"}
+        
+        baseline = self.reports[baseline_label]
+        baseline_success = baseline.success_rate
+        baseline_cost_success = baseline.avg_cost_per_success
+        
+        comparison = {}
+        for name, report in self.reports.items():
+            sr = report.success_rate
+            cs = report.avg_cost_per_success
+            
+            if sr >= baseline_success * 0.95:  # within 5% = iso-quality
+                cost_reduction = (baseline_cost_success - cs) / baseline_cost_success * 100
+            else:
+                cost_reduction = None  # not iso-quality
+            
+            comparison[name] = {
+                "description": report.config_description,
+                "success_rate": f"{sr:.1%}",
+                "cost_per_success": f"${cs:.4f}",
+                "cost_reduction_at_iso": f"{cost_reduction:.1f}%" if cost_reduction is not None else "N/A (quality drop)",
+                "avg_llm_calls": f"{report.avg_llm_calls:.1f}",
+                "avg_tool_calls": f"{report.avg_tool_calls:.1f}",
+                "false_done_rate": f"{report.false_done_rate:.1%}",
+            }
+        
+        return comparison
+    
+    def print_comparison(self):
+        """Pretty-print comparison table."""
+        comp = self.compare()
+        if "error" in comp:
+            print(comp["error"])
+            return
+        
+        print(f"\n{'='*90}")
+        print(f"ACO BENCHMARK COMPARISON ({len(self.reports)} configs, 25 tasks)")
+        print(f"{'='*90}")
+        print(f"{'Config':<25} {'Success':>8} {'Cost/Success':>13} {'Cost Reduction':>15} {'LLM Calls':>10} {'Tools':>7} {'False-DONE':>10}")
+        print(f"{'-'*90}")
+        
+        for name, stats in comp.items():
+            print(f"{name:<25} {stats['success_rate']:>8} {stats['cost_per_success']:>13} "
+                  f"{stats['cost_reduction_at_iso']:>15} {stats['avg_llm_calls']:>10} "
+                  f"{stats['avg_tool_calls']:>7} {stats['false_done_rate']:>10}")
+        
+        print(f"{'='*90}")
+        
+        # Find best iso-quality config
+        best = None
+        best_reduction = 0
+        for name, stats in comp.items():
+            if "N/A" not in stats["cost_reduction_at_iso"]:
+                pct = float(stats["cost_reduction_at_iso"].rstrip("%"))
+                if pct > best_reduction:
+                    best_reduction = pct
+                    best = name
+        
+        if best:
+            print(f"\nBEST ISO-QUALITY CONFIG: {best} ({best_reduction:.1f}% cost reduction)")
+    
+    def save_results(self, path: str):
+        """Save all results to JSON."""
+        data = {
+            name: report.to_dict()
+            for name, report in self.reports.items()
+        }
+        with open(path, "w") as f:
+            json.dump(data, f, indent=2)
+        print(f"\nResults saved to {path}")
+
+
+# =============================================================================
+# Quick test
+# =============================================================================
+if __name__ == "__main__":
+    agent = SimulatedAgent()
+    runner = BenchmarkRunner(agent)
+    runner.run_all_baselines()
+    runner.print_comparison()
+    runner.save_results("/tmp/benchmark_results.json")