evaluation/benchmark.py

"""
evaluation/benchmark.py
------------------------
Built-in benchmark tasks and runner for AutoDevAgent.

Provides 5 standardised tasks — 3 Python and 2 SQL — that can be
run against the full pipeline to measure agent quality. Results are
logged to W&B and exported for the README and LinkedIn posts.

Benchmark tasks are fixed so results are comparable across runs,
model changes, and prompt iterations. The same tasks always run in
the same order so W&B charts show meaningful trends over time.

Tasks:
  Python:
    1. Reverse a string             (basic, 1 function)
    2. Fibonacci sequence           (recursion / iteration)
    3. Find duplicates in a list    (data manipulation, edge cases)

  SQL:
    4. Top N customers by revenue   (aggregation, ORDER BY, LIMIT)
    5. Department headcount query   (GROUP BY, COUNT, JOIN)

Design:
  - Each BenchmarkTask is a frozen dataclass — immutable, hashable,
    and safe to share across threads.
  - BenchmarkRunner.run_all() runs every task through the full
    pipeline graph and collects BenchmarkResult objects.
  - Results are passed to metrics.py for aggregation and to
    wandb_tracker.py for logging.

Usage:
    from evaluation.benchmark import BenchmarkRunner

    runner  = BenchmarkRunner()
    results = runner.run_all()
    for r in results:
        print(r.task_name, r.success, r.iterations, r.total_tokens)
"""

import logging
import time
from dataclasses import dataclass
from typing import Any

from config import settings

logger = logging.getLogger(__name__)


# ------------------------------------------------------------------ #
#  Task definitions                                                    #
# ------------------------------------------------------------------ #

@dataclass(frozen=True)
class BenchmarkTask:
    """
    A single immutable benchmark task.

    Attributes:
        name:        Short display name shown in W&B and README tables.
        description: The task description sent to the pipeline as input.
        language:    "python" or "sql".
        category:    Broad category for grouping in W&B charts.
    """
    name:        str
    description: str
    language:    str
    category:    str


# The 5 canonical benchmark tasks — do not change without versioning
BENCHMARK_TASKS: list[BenchmarkTask] = [

    BenchmarkTask(
        name="Reverse string",
        description=(
            "Write a Python function called reverse_string that takes a "
            "string as input and returns it reversed. Handle edge cases: "
            "empty string, single character, and string with spaces. "
            "Include a main guard that prints example outputs."
        ),
        language="python",
        category="string manipulation",
    ),

    BenchmarkTask(
        name="Fibonacci sequence",
        description=(
            "Write a Python function called fibonacci that takes an integer n "
            "and returns a list of the first n Fibonacci numbers. "
            "Handle edge cases: n=0 returns empty list, n=1 returns [0], n=2 returns [0,1]. "
            "Use an iterative approach (not recursive). "
            "Include a main guard that prints fibonacci(10)."
        ),
        language="python",
        category="algorithms",
    ),

    BenchmarkTask(
        name="Find duplicates",
        description=(
            "Write a Python function called find_duplicates that takes a list "
            "and returns a sorted list of all elements that appear more than once. "
            "Each duplicate should appear only once in the result even if it "
            "appears many times in the input. Handle empty list and list with no "
            "duplicates (return empty list). "
            "Include a main guard with example outputs."
        ),
        language="python",
        category="data manipulation",
    ),

    BenchmarkTask(
        name="Top customers by revenue",
        description=(
            "Write a SQL query to find the top 5 customers by total order revenue. "
            "The result should show customer name and their total revenue, "
            "ordered from highest to lowest revenue."
        ),
        language="sql",
        category="aggregation",
    ),

    BenchmarkTask(
        name="Department headcount",
        description=(
            "Write a SQL query that shows each department name, the number of "
            "employees in that department, and the average salary. "
            "Order results by headcount descending. "
            "Only include departments with more than 1 employee."
        ),
        language="sql",
        category="grouping and filtering",
    ),
]


# ------------------------------------------------------------------ #
#  Result dataclass                                                    #
# ------------------------------------------------------------------ #

@dataclass
class BenchmarkResult:
    """
    Result of running a single benchmark task through the pipeline.

    Attributes:
        task_name:        The BenchmarkTask.name.
        language:         "python" or "sql".
        category:         Task category for W&B grouping.
        success:          True if the pipeline reached SUCCESS status.
        iterations:       Number of debug iterations required.
        total_tokens:     Total Groq tokens consumed.
        exec_time:        Wall-clock seconds for the full pipeline run.
        final_status:     The pipeline's final PipelineStatus value.
        error_message:    Last error if failed, else empty string.
        tests_passed:     True if test suite passed (if applicable).
        total_tests:      Total number of tests run.
    """
    task_name:     str   = ""
    language:      str   = "python"
    category:      str   = ""
    success:       bool  = False
    iterations:    int   = 0
    total_tokens:  int   = 0
    exec_time:     float = 0.0
    final_status:  str   = ""
    error_message: str   = ""
    tests_passed:  bool  = False
    total_tests:   int   = 0

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dict for W&B logging."""
        return {
            "task_name":    self.task_name,
            "language":     self.language,
            "category":     self.category,
            "success":      self.success,
            "iterations":   self.iterations,
            "total_tokens": self.total_tokens,
            "exec_time":    self.exec_time,
            "final_status": self.final_status,
            "tests_passed": self.tests_passed,
            "total_tests":  self.total_tests,
        }


# ------------------------------------------------------------------ #
#  Runner                                                              #
# ------------------------------------------------------------------ #

class BenchmarkRunner:
    """
    Runs all 5 benchmark tasks through the full AutoDevAgent pipeline.

    Each task is run independently with a fresh pipeline state. Results
    are collected into a list of BenchmarkResult objects and passed to
    the metrics module for aggregation and W&B for logging.

    Attributes:
        tasks:    The list of BenchmarkTask objects to run.
        results:  Populated after run_all() completes.
    """

    def __init__(
        self,
        tasks: list[BenchmarkTask] | None = None,
    ) -> None:
        """
        Initialise with the default benchmark tasks or a custom list.

        Args:
            tasks: Optional custom task list. Defaults to BENCHMARK_TASKS.
        """
        self.tasks:   list[BenchmarkTask]   = tasks or BENCHMARK_TASKS
        self.results: list[BenchmarkResult] = []

    def run_all(
        self,
        on_progress: Any = None,
    ) -> list[BenchmarkResult]:
        """
        Run all benchmark tasks sequentially and collect results.

        Args:
            on_progress: Optional callable(task_name, index, total)
                         called before each task starts. Used to update
                         a Gradio progress bar from app.py.

        Returns:
            List of BenchmarkResult objects, one per task.
        """
        from pipeline.graph import run_pipeline

        self.results = []
        total = len(self.tasks)

        logger.info("BenchmarkRunner: starting %d tasks", total)

        for i, task in enumerate(self.tasks):
            if on_progress:
                try:
                    on_progress(task.name, i, total)
                except Exception:
                    pass

            logger.info(
                "BenchmarkRunner: running task %d/%d — '%s' (%s)",
                i + 1, total, task.name, task.language,
            )

            result = self._run_task(task)
            self.results.append(result)

            logger.info(
                "BenchmarkRunner: task '%s' — success=%s iterations=%d tokens=%d",
                task.name, result.success, result.iterations, result.total_tokens,
            )

        logger.info(
            "BenchmarkRunner: completed %d/%d tasks successfully",
            sum(1 for r in self.results if r.success),
            total,
        )

        return self.results

    def _run_task(self, task: BenchmarkTask) -> BenchmarkResult:
        """
        Run a single benchmark task through the full pipeline.

        Calls run_pipeline() with the task description and language,
        extracts metrics from the final state, and returns a
        BenchmarkResult.

        Args:
            task: The BenchmarkTask to run.

        Returns:
            Populated BenchmarkResult for this task.
        """
        from pipeline.graph import run_pipeline
        from pipeline.state import PipelineStatus

        start = time.perf_counter()

        try:
            final_state = run_pipeline(
                task=task.description,
                language=task.language,
            )
            elapsed = round(time.perf_counter() - start, 2)

            # Extract test result info if available
            tests_passed = False
            total_tests  = 0
            if final_state.test_result:
                tests_passed = final_state.test_result.passed
                total_tests  = final_state.test_result.total_tests

            return BenchmarkResult(
                task_name    = task.name,
                language     = task.language,
                category     = task.category,
                success      = (final_state.status == PipelineStatus.SUCCESS),
                iterations   = final_state.debug_iterations,
                total_tokens = final_state.token_usage.total_tokens,
                exec_time    = elapsed,
                final_status = final_state.status.value,
                error_message = final_state.latest_error()[:200],
                tests_passed = tests_passed,
                total_tests  = total_tests,
            )

        except Exception as e:
            elapsed = round(time.perf_counter() - start, 2)
            logger.error(
                "BenchmarkRunner: task '%s' crashed: %s", task.name, e
            )
            return BenchmarkResult(
                task_name     = task.name,
                language      = task.language,
                category      = task.category,
                success       = False,
                exec_time     = elapsed,
                final_status  = "crashed",
                error_message = str(e)[:200],
            )