feat: add pass-at-k code evaluation

code_agent/__init__.py

@@ -1,6 +1,12 @@
 """Utilities for execution-feedback code generation experiments."""
 
 from .reward import CodeReward, score_verification
+from .evaluation import (
+    EvaluationReport,
+    aggregate_reports,
+    estimate_pass_at_k,
+    evaluate_candidates,
+)
 from .verifier import (
     CodeTask,
     TestCase,
@@ -13,9 +19,13 @@ from .verifier import (
 __all__ = [
     "CodeReward",
     "CodeTask",
+    "EvaluationReport",
     "TestCase",
     "VerificationResult",
     "VerificationStatus",
+    "aggregate_reports",
+    "estimate_pass_at_k",
+    "evaluate_candidates",
     "extract_python_code",
     "score_verification",
     "verify_code",

code_agent/evaluation.py

@@ -0,0 +1,116 @@
+"""Evaluation metrics for execution-verified code generation."""
+
+from __future__ import annotations
+
+import math
+from collections import Counter
+from dataclasses import asdict, dataclass
+from typing import Any, Iterable, Sequence
+
+from .verifier import CodeTask, VerificationResult, VerificationStatus, verify_code
+
+
+def estimate_pass_at_k(num_samples: int, num_correct: int, k: int) -> float:
+    """Unbiased pass@k estimator used by common code-generation benchmarks."""
+    if num_samples <= 0:
+        raise ValueError("num_samples must be positive")
+    if not 0 <= num_correct <= num_samples:
+        raise ValueError("num_correct must be between 0 and num_samples")
+    if not 1 <= k <= num_samples:
+        raise ValueError("k must be between 1 and num_samples")
+    if num_samples - num_correct < k:
+        return 1.0
+    return 1.0 - math.comb(num_samples - num_correct, k) / math.comb(num_samples, k)
+
+
+@dataclass(frozen=True)
+class EvaluationReport:
+    task_id: str
+    candidate_count: int
+    correct_count: int
+    compile_rate: float
+    execution_rate: float
+    test_pass_rate: float
+    pass_at_k: dict[str, float]
+    status_counts: dict[str, int]
+    results: tuple[VerificationResult, ...]
+
+    def to_dict(self, *, include_results: bool = True) -> dict[str, Any]:
+        payload = asdict(self)
+        if include_results:
+            payload["results"] = [result.to_dict() for result in self.results]
+        else:
+            payload.pop("results", None)
+        return payload
+
+
+_COMPILED_STATUSES = {
+    VerificationStatus.PASSED,
+    VerificationStatus.RUNTIME_ERROR,
+    VerificationStatus.TIMEOUT,
+    VerificationStatus.WRONG_ANSWER,
+}
+_EXECUTED_STATUSES = {
+    VerificationStatus.PASSED,
+    VerificationStatus.RUNTIME_ERROR,
+    VerificationStatus.WRONG_ANSWER,
+}
+
+
+def evaluate_candidates(
+    task: CodeTask,
+    candidates: Sequence[str],
+    *,
+    ks: Iterable[int] = (1, 5, 10),
+    timeout_seconds: float = 2.0,
+    memory_mb: int = 256,
+) -> EvaluationReport:
+    if not candidates:
+        raise ValueError("at least one candidate is required")
+    results = tuple(
+        verify_code(candidate, task, timeout_seconds=timeout_seconds, memory_mb=memory_mb)
+        for candidate in candidates
+    )
+    count = len(results)
+    correct = sum(result.status == VerificationStatus.PASSED for result in results)
+    status_counts = Counter(result.status.value for result in results)
+    total_tests = sum(result.total_tests for result in results)
+    passed_tests = sum(result.passed_tests for result in results)
+    requested_ks = sorted(set(int(k) for k in ks if 1 <= int(k) <= count))
+    pass_at_k = {f"pass@{k}": estimate_pass_at_k(count, correct, k) for k in requested_ks}
+    return EvaluationReport(
+        task_id=task.task_id,
+        candidate_count=count,
+        correct_count=correct,
+        compile_rate=sum(result.status in _COMPILED_STATUSES for result in results) / count,
+        execution_rate=sum(result.status in _EXECUTED_STATUSES for result in results) / count,
+        test_pass_rate=passed_tests / total_tests if total_tests else 0.0,
+        pass_at_k=pass_at_k,
+        status_counts=dict(sorted(status_counts.items())),
+        results=results,
+    )
+
+
+def aggregate_reports(reports: Sequence[EvaluationReport]) -> dict[str, Any]:
+    if not reports:
+        raise ValueError("at least one evaluation report is required")
+    status_counts: Counter[str] = Counter()
+    for report in reports:
+        status_counts.update(report.status_counts)
+    metric_names = sorted({name for report in reports for name in report.pass_at_k})
+    pass_at_k = {
+        name: sum(report.pass_at_k[name] for report in reports if name in report.pass_at_k)
+        / sum(name in report.pass_at_k for report in reports)
+        for name in metric_names
+    }
+    candidate_count = sum(report.candidate_count for report in reports)
+    return {
+        "task_count": len(reports),
+        "candidate_count": candidate_count,
+        "correct_count": sum(report.correct_count for report in reports),
+        "compile_rate": sum(report.compile_rate * report.candidate_count for report in reports) / candidate_count,
+        "execution_rate": sum(report.execution_rate * report.candidate_count for report in reports) / candidate_count,
+        "test_pass_rate": sum(report.test_pass_rate * report.candidate_count for report in reports) / candidate_count,
+        "pass_at_k": pass_at_k,
+        "status_counts": dict(sorted(status_counts.items())),
+    }

code_agent/verifier.py

@@ -21,7 +21,7 @@ import time
 from dataclasses import asdict, dataclass, field
 from enum import Enum
 from pathlib import Path
-from typing import Any, Iterable
+from typing import Any
 
 
 class VerificationStatus(str, Enum):
@@ -44,6 +44,14 @@ class TestCase:
     def to_json(self) -> dict[str, Any]:
         return {"args": list(self.args), "kwargs": self.kwargs, "expected": self.expected}
 
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> "TestCase":
+        return cls(
+            args=tuple(payload.get("args", ())),
+            kwargs=dict(payload.get("kwargs", {})),
+            expected=payload.get("expected"),
+        )
+
 
 @dataclass(frozen=True)
 class CodeTask:
@@ -60,6 +68,23 @@ class CodeTask:
         # Fail early instead of discovering non-serializable tests in a worker.
         json.dumps([case.to_json() for case in self.tests], ensure_ascii=False)
 
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "task_id": self.task_id,
+            "prompt": self.prompt,
+            "entry_point": self.entry_point,
+            "tests": [case.to_json() for case in self.tests],
+        }
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> "CodeTask":
+        return cls(
+            task_id=str(payload["task_id"]),
+            prompt=str(payload["prompt"]),
+            entry_point=str(payload["entry_point"]),
+            tests=tuple(TestCase.from_dict(case) for case in payload["tests"]),
+        )
+
 
 @dataclass(frozen=True)
 class VerificationResult:

scripts/eval_codegen.py

@@ -0,0 +1,66 @@
+"""Evaluate JSONL code-generation candidates with deterministic tests."""
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from code_agent import CodeTask, aggregate_reports, evaluate_candidates
+
+
+def read_jsonl(path: str):
+    with open(path, "r", encoding="utf-8") as handle:
+        for line_number, line in enumerate(handle, start=1):
+            if line.strip():
+                try:
+                    yield json.loads(line)
+                except json.JSONDecodeError as exc:
+                    raise ValueError(f"Invalid JSON at {path}:{line_number}: {exc}") from exc
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Execution-based evaluation for code generation")
+    parser.add_argument("--tasks", required=True, help="JSONL file containing CodeTask records")
+    parser.add_argument("--predictions", required=True, help="JSONL records with task_id and candidates")
+    parser.add_argument("--output", help="Optional JSON report path")
+    parser.add_argument("--ks", type=int, nargs="+", default=[1, 5, 10], help="pass@k values")
+    parser.add_argument("--timeout", type=float, default=2.0, help="Timeout per candidate in seconds")
+    parser.add_argument("--memory-mb", type=int, default=256, help="POSIX worker memory limit")
+    args = parser.parse_args()
+
+    tasks = {record["task_id"]: CodeTask.from_dict(record) for record in read_jsonl(args.tasks)}
+    predictions = {record["task_id"]: record["candidates"] for record in read_jsonl(args.predictions)}
+    missing = sorted(set(tasks) - set(predictions))
+    unknown = sorted(set(predictions) - set(tasks))
+    if missing or unknown:
+        raise ValueError(f"Task/prediction mismatch: missing={missing}, unknown={unknown}")
+
+    reports = [
+        evaluate_candidates(
+            task,
+            predictions[task_id],
+            ks=args.ks,
+            timeout_seconds=args.timeout,
+            memory_mb=args.memory_mb,
+        )
+        for task_id, task in tasks.items()
+    ]
+    output = {
+        "summary": aggregate_reports(reports),
+        "tasks": [report.to_dict() for report in reports],
+        "safety_note": (
+            "The local verifier is for controlled experiments, not hostile code. "
+            "Use a dedicated container or micro-VM for untrusted submissions."
+        ),
+    }
+    rendered = json.dumps(output, ensure_ascii=False, indent=2)
+    if args.output:
+        Path(args.output).write_text(rendered + "\n", encoding="utf-8")
+    print(rendered)
+
+
+if __name__ == "__main__":
+    main()

tests/fixtures/code_predictions.jsonl

@@ -0,0 +1 @@
+{"task_id":"add-two","candidates":["def add(a, b):\n    return a + b","def add(a, b):\n    return a - b"]}

tests/fixtures/code_tasks.jsonl

@@ -0,0 +1 @@
+{"task_id":"add-two","prompt":"Implement add(a, b) and return the sum.","entry_point":"add","tests":[{"args":[1,2],"expected":3},{"args":[-5,5],"expected":0}]}

tests/test_code_evaluation.py

@@ -0,0 +1,54 @@
+import pytest
+
+from code_agent import (
+    CodeTask,
+    TestCase as CodeTestCase,
+    VerificationStatus,
+    aggregate_reports,
+    estimate_pass_at_k,
+    evaluate_candidates,
+)
+
+
+TASK = CodeTask(
+    task_id="square",
+    prompt="Implement square(x).",
+    entry_point="square",
+    tests=(
+        CodeTestCase(args=(2,), expected=4),
+        CodeTestCase(args=(-3,), expected=9),
+    ),
+)
+
+
+def test_pass_at_k_estimator():
+    assert estimate_pass_at_k(10, 0, 1) == 0.0
+    assert estimate_pass_at_k(10, 1, 10) == 1.0
+    assert estimate_pass_at_k(10, 2, 1) == pytest.approx(0.2)
+    with pytest.raises(ValueError):
+        estimate_pass_at_k(2, 1, 3)
+
+
+def test_evaluation_reports_execution_metrics():
+    candidates = [
+        "def square(x):\n    return x * x",
+        "def square(x):\n    return x + x",
+        "def square(x)\n    return x * x",
+    ]
+    report = evaluate_candidates(TASK, candidates, ks=(1, 3))
+
+    assert report.correct_count == 1
+    assert report.compile_rate == pytest.approx(2 / 3)
+    assert report.execution_rate == pytest.approx(2 / 3)
+    assert report.test_pass_rate == pytest.approx(0.5)
+    assert report.pass_at_k["pass@1"] == pytest.approx(1 / 3)
+    assert report.pass_at_k["pass@3"] == 1.0
+    assert report.results[2].status == VerificationStatus.SYNTAX_ERROR
+
+    summary = aggregate_reports([report])
+    assert summary["candidate_count"] == 3
+    assert summary["status_counts"] == {
+        "passed": 1,
+        "syntax_error": 1,
+        "wrong_answer": 1,
+    }