Evaluation with Test Suites

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
 
class Difficulty(Enum):
    EASY = "easy"      # 1-2 skills, straightforward
    MEDIUM = "medium"  # 2-4 skills, some ambiguity
    HARD = "hard"      # 4+ skills, complex reasoning
 
class ScoringMethod(Enum):
    EXACT_MATCH = "exact_match"
    CONTAINS = "contains"
    LLM_JUDGE = "llm_judge"
    RUBRIC = "rubric"
 
@dataclass
class EvalTask:
    task_id: str
    input_prompt: str
    difficulty: Difficulty
    scoring_method: ScoringMethod
    scoring_config: dict = field(default_factory=dict)
    max_steps: int = 20
    max_cost_usd: float = 0.50
    tags: list[str] = field(default_factory=list)
 
EVAL_SUITE = [
    EvalTask(task_id="search_001", input_prompt="What is the population of Tokyo?",
             difficulty=Difficulty.EASY, scoring_method=ScoringMethod.CONTAINS,
             scoring_config={"must_contain": ["13", "million"]}, tags=["search"]),
    EvalTask(task_id="analysis_001",
             input_prompt="Download the CSV, calculate average sales by region, find top 3.",
             difficulty=Difficulty.MEDIUM, scoring_method=ScoringMethod.RUBRIC,
             scoring_config={"rubric": [
                 {"criterion": "Calculated averages", "points": 3},
                 {"criterion": "Correct top 3", "points": 3},
                 {"criterion": "Clear presentation", "points": 2}],
                 "passing_score": 6}, tags=["analysis", "code"]),
    EvalTask(task_id="research_001",
             input_prompt="Compare climate policies of three G7 nations.",
             difficulty=Difficulty.HARD, scoring_method=ScoringMethod.LLM_JUDGE,
             scoring_config={"judge_prompt": "Does this accurately compare 3 G7 nations "
                            "with specific emissions targets? Score 1-5.", "passing_score": 3},
             tags=["research", "synthesis"]),
]

class ContainsScorer:
    async def score(self, task: EvalTask, output: str) -> dict:
        terms = task.scoring_config["must_contain"]
        text = output.lower()
        found = [t for t in terms if t.lower() in text]
        score = len(found) / len(terms) if terms else 0.0
        return {"passed": score >= 1.0, "score": score, "missing": list(set(terms) - set(found))}
 
class LLMJudgeScorer:
    def __init__(self, llm_client):
        self.llm = llm_client
 
    async def score(self, task: EvalTask, output: str) -> dict:
        prompt = f"""Task: {task.input_prompt}\nOutput: {output[:3000]}
Criteria: {task.scoring_config['judge_prompt']}
Respond with JSON: {{"score": <1-5>, "reasoning": "<brief>"}}"""
        response = await self.llm.generate(prompt, max_tokens=200)
        result = json.loads(response.text)
        passing = task.scoring_config.get("passing_score", 3)
        return {"passed": result["score"] >= passing, "score": result["score"] / 5.0,
                "reasoning": result["reasoning"]}
 
class RubricScorer:
    def __init__(self, llm_client):
        self.llm = llm_client
 
    async def score(self, task: EvalTask, output: str) -> dict:
        rubric = task.scoring_config["rubric"]
        total = sum(c["points"] for c in rubric)
        criteria = "\n".join(f"- {c['criterion']} ({c['points']} pts)" for c in rubric)
        prompt = f"""Task: {task.input_prompt}\nOutput: {output[:3000]}
Rubric:\n{criteria}\nAward 0 to max points per criterion.
Respond with JSON: {{"scores": [...], "total": N}}"""
        result = json.loads((await self.llm.generate(prompt, max_tokens=500)).text)
        passing = task.scoring_config.get("passing_score", total * 0.7)
        return {"passed": result["total"] >= passing, "score": result["total"] / total}

import asyncio, time
 
@dataclass
class EvalResult:
    task_id: str
    passed: bool
    score: float
    steps_taken: int
    cost_usd: float
    duration_seconds: float
    error: Optional[str] = None
 
class EvaluationHarness:
    def __init__(self, agent, scorers: dict):
        self.agent = agent
        self.scorers = scorers
 
    async def run_suite(self, tasks: list[EvalTask], concurrency: int = 3) -> list[EvalResult]:
        sem = asyncio.Semaphore(concurrency)
        async def run_one(task):
            async with sem:
                return await self._evaluate_task(task)
        return await asyncio.gather(*[run_one(t) for t in tasks])
 
    async def _evaluate_task(self, task: EvalTask) -> EvalResult:
        start = time.time()
        try:
            output = await asyncio.wait_for(self.agent.execute(task.input_prompt),
                                            timeout=120)
            result = await self.scorers[task.scoring_method].score(task, output.text)
            return EvalResult(task.task_id, result["passed"], result["score"],
                              output.steps_taken, output.cost_usd, time.time() - start)
        except Exception as e:
            return EvalResult(task.task_id, False, 0.0, 0, 0.0, time.time() - start, str(e))
 
    def summarize(self, results: list[EvalResult]) -> dict:
        n = len(results)
        return {
            "completion_rate": sum(r.passed for r in results) / n,
            "avg_score": sum(r.score for r in results) / n,
            "avg_steps": sum(r.steps_taken for r in results) / n,
            "avg_cost_usd": sum(r.cost_usd for r in results) / n,
            "total_cost_usd": sum(r.cost_usd for r in results),
        }