humaneval_agent.py

"""HumanEval solver built on the AG2 (autogen) multi-agent framework.

Three AG2 agents collaborate per problem (hard 3-min wall clock):
  1. Solver  (AssistantAgent)               — writes a Python implementation.
  2. Tester  (AssistantAgent)               — writes additional unit tests.
  3. Verifier(UserProxyAgent + executor)    — runs the code in a sandboxed
     work_dir via AG2's LocalCommandLineCodeExecutor, replies with exit status
     and output. Used twice per attempt: once for the agent-written tests,
     once for the canonical HumanEval `check()` test.
On failure, the orchestrator retries up to 3 times, feeding the verifier's
output back to the solver as feedback.

Usage:
    python humaneval_agent.py            # run all 164 problems
    python humaneval_agent.py 5          # run first 5 only
    python humaneval_agent.py 10 50      # problems [10, 60)
"""

import json
import multiprocessing as mp
import os
import re
import shutil
import sys
import tempfile
import time
import traceback
from pathlib import Path

from autogen import AssistantAgent, LLMConfig, UserProxyAgent
from autogen.coding import LocalCommandLineCodeExecutor

ROOT = Path(__file__).parent
DATA = ROOT / "HumanEval.jsonl"
OUT = ROOT / "humaneval_results.jsonl"
CONFIG_PATH = ROOT / "OAI_CONFIG_LIST"

TIME_LIMIT = 300          # seconds per problem (hard cap; raised from 180s
                          # so r1 fallback in attempt 3 has room to finish)
SUBPROC_TIMEOUT = 20      # per code-execution timeout

# Per-attempt escalation strategy:
#   1. Single sample at low temp on the cheap model.
#   2. Best-of-3 at medium temp on the cheap model — verifier picks the first
#      candidate that passes the canonical test.
#   3. Reasoning-model rescue: deepseek-r1, single sample.
ATTEMPT_STRATEGIES = [
    # Attempt 1: server-default temperature (~1.0). Empirically gives the best
    # first-attempt accuracy on this benchmark — temp 0.2 was measurably worse.
    {"model": "deepseek/deepseek-chat", "temperature": None, "n": 1},
    # Attempt 2: best-of-2 at higher temp for retry diversity.
    {"model": "deepseek/deepseek-chat", "temperature": 0.5, "n": 2},
    # Attempt 3: reasoning-model rescue.
    {"model": "deepseek/deepseek-r1",   "temperature": 0.5, "n": 1},
]
MAX_RETRIES = len(ATTEMPT_STRATEGIES)

SOLVER_SYS = """You are an expert Python programmer. You will be given a Python
function signature plus its docstring. Produce a COMPLETE, CORRECT, self-contained
implementation.

Strict output format: a single fenced ```python``` block containing the imports
(if any) and the full function definition. NO prose, NO unit tests, NO examples."""

TESTER_SYS = """You are a Python testing expert. Given a function spec (signature
+ docstring) and a candidate implementation, write 3-6 additional assert-based
unit tests covering normal cases and edge cases implied by the docstring.

Strict output format: one fenced ```python``` block defining exactly:

    def run_tests(candidate):
        assert candidate(...) == ...
        ...

Do not redefine the candidate. No prose."""


def extract_python(text: str) -> str:
    m = re.search(r"```python\s*\n(.*?)```", text, re.DOTALL)
    if m:
        return m.group(1).strip()
    m = re.search(r"```\s*\n(.*?)```", text, re.DOTALL)
    return (m.group(1) if m else text).strip()


def verifier_run(verifier: UserProxyAgent, user: UserProxyAgent, code: str):
    """Ask the Verifier agent to execute `code`. Returns (ok, output)."""
    msg = f"Please execute the following Python program and report the result:\n\n```python\n{code}\n```"
    chat = user.initiate_chat(
        verifier, message=msg, max_turns=1, clear_history=True, silent=True,
    )
    reply = chat.chat_history[-1]["content"] if chat.chat_history else ""
    # AG2's executor reply format: "exitcode: 0 (execution succeeded)\nCode output: ..."
    ok = "exitcode: 0" in reply
    # Trim — keep last ~2000 chars of output for feedback
    return ok, reply.strip()[-2000:]


def assemble_solution(prompt: str, entry_point: str, model_output: str) -> str:
    """Combine the prompt's import preamble with the model's output."""
    code = extract_python(model_output)
    if f"def {entry_point}" not in code:
        # The model only returned a body — splice it under the prompt.
        body = code
        lines = body.splitlines()
        if lines and not lines[0].startswith((" ", "\t")):
            body = "\n".join("    " + ln for ln in lines)
        return prompt.rstrip() + "\n" + body + "\n"
    # Make sure prompt-level imports are present.
    preamble = prompt.split(f"def {entry_point}")[0].strip()
    if preamble and preamble not in code:
        code = preamble + "\n\n" + code
    return code


def llm_config_for(model: str, temperature):
    """Load OAI_CONFIG_LIST, filter to the requested model, optionally inject
    temperature. Pass `temperature=None` to leave it unset (server default)."""
    raw = json.loads(CONFIG_PATH.read_text())
    filtered = [c for c in raw if c["model"] == model]
    if not filtered:
        raise ValueError(f"Model {model!r} not present in OAI_CONFIG_LIST "
                         f"(have: {[c['model'] for c in raw]})")
    cfg: dict = {"config_list": filtered}
    if temperature is not None:
        cfg["temperature"] = temperature
    return cfg


def make_solver_tester(model: str, temperature: float):
    cfg = llm_config_for(model, temperature)
    solver = AssistantAgent(
        "solver", llm_config=cfg, system_message=SOLVER_SYS,
        human_input_mode="NEVER",
    )
    tester = AssistantAgent(
        "tester", llm_config=cfg, system_message=TESTER_SYS,
        human_input_mode="NEVER",
    )
    return solver, tester


_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)

def strip_thinking(text: str) -> str:
    """Remove <think>...</think> reasoning blocks (deepseek-r1 emits these)."""
    return _THINK_RE.sub("", text).strip()


def make_verifier_and_user(work_dir: str):
    # Verifier: a code-executing agent. Receives a python code block, runs it
    # via a local sandbox, replies with exit status + output.
    verifier = UserProxyAgent(
        "verifier", human_input_mode="NEVER",
        code_execution_config={
            "executor": LocalCommandLineCodeExecutor(
                work_dir=work_dir, timeout=SUBPROC_TIMEOUT,
            ),
        },
        max_consecutive_auto_reply=1,
        default_auto_reply="",
    )
    # Orchestrator stub — sends messages, doesn't execute or auto-reply itself.
    user = UserProxyAgent(
        "user", human_input_mode="NEVER", code_execution_config=False,
        max_consecutive_auto_reply=0, default_auto_reply="",
    )
    return verifier, user


def chat_once(user, agent, message: str) -> str:
    user.initiate_chat(agent, message=message, max_turns=1, clear_history=True, silent=True)
    msg = agent.last_message()
    return msg["content"] if isinstance(msg, dict) else str(msg)


def chat_once_robust(user, agent, message: str, retries: int = 3) -> str:
    """chat_once with retry-with-backoff for transient OpenRouter/upstream errors
    (HTML error pages, rate limits, dropped connections — anything that surfaces
    as a non-JSON response or network exception)."""
    last_err = None
    for i in range(retries):
        try:
            return chat_once(user, agent, message)
        except Exception as e:
            last_err = e
            if i < retries - 1:
                time.sleep(2.0 * (i + 1))  # 2s, 4s
    raise last_err


def solve_problem(problem: dict) -> dict:
    prompt = problem["prompt"]
    entry_point = problem["entry_point"]
    canonical_test = problem["test"]
    task_id = problem["task_id"]

    work_dir = tempfile.mkdtemp(prefix=f"he_{task_id.replace('/', '_')}_")
    verifier, user = make_verifier_and_user(work_dir)

    feedback = ""
    solution_code = ""
    log = []

    def build_canonical_program(sol):
        return (sol + "\n\n" + canonical_test +
                f"\n\ncheck({entry_point})\nprint('CANONICAL_OK')\n")

    try:
      for attempt, strat in enumerate(ATTEMPT_STRATEGIES, start=1):
        model, temperature, n_samples = strat["model"], strat["temperature"], strat["n"]
        try:
            solver, tester = make_solver_tester(model, temperature)

            # 1) SOLVE — generate n_samples candidate solutions; first one whose
            #    canonical verification passes is the winner. (Best-of-N.)
            ask = f"Implement this function:\n\n{prompt}"
            if feedback:
                ask += (
                    "\n\nA previous attempt failed. Use this feedback to fix it:\n"
                    f"{feedback}\n\nReturn the corrected COMPLETE implementation."
                )

            winner_solution = None
            sample_results = []  # list of (solution_code, ok_c, out_c)
            for sample_i in range(n_samples):
                sol_raw = strip_thinking(chat_once_robust(user, solver, ask))
                cand = assemble_solution(prompt, entry_point, sol_raw)
                ok_c, out_c = verifier_run(verifier, user, build_canonical_program(cand))
                sample_results.append((cand, ok_c, out_c))
                if ok_c:
                    winner_solution = cand
                    break  # short-circuit best-of-N

            if winner_solution is not None:
                solution_code = winner_solution
                # 2) UNIT TEST (agent-written) — only run on the chosen winner.
                tmsg = (
                    f"Function spec:\n\n{prompt}\n\n"
                    f"Candidate implementation:\n```python\n{solution_code}\n```\n\n"
                    "Write the additional unit tests as `def run_tests(candidate):`."
                )
                tests_raw = strip_thinking(chat_once_robust(user, tester, tmsg))
                tests_code = extract_python(tests_raw)

                # 3a) VERIFY agent-written tests
                agent_test_program = (
                    solution_code + "\n\n" + tests_code +
                    f"\n\nrun_tests({entry_point})\nprint('AGENT_TESTS_OK')\n"
                )
                ok_t, out_t = verifier_run(verifier, user, agent_test_program)

                log.append({
                    "attempt": attempt, "model": model, "temperature": temperature,
                    "samples": sample_results.__len__(),
                    "agent_tests_passed": ok_t, "canonical_passed": True,
                    "agent_test_output": out_t if not ok_t else "",
                })
                return {
                    "task_id": task_id, "passed": True, "attempts": attempt,
                    "samples_tried": len(sample_results),
                    "agent_tests_passed": ok_t, "solution": solution_code,
                    "model": model, "log": log,
                }

            # No sample passed canonical — record + build feedback for next rung.
            solution_code = sample_results[-1][0] if sample_results else solution_code
            errs = "\n---\n".join(out for _, _, out in sample_results)
            log.append({
                "attempt": attempt, "model": model, "temperature": temperature,
                "samples": len(sample_results), "agent_tests_passed": False,
                "canonical_passed": False, "canonical_outputs": errs[-3000:],
            })
            feedback = f"[Attempt {attempt} ({model}, temp={temperature}, "
            feedback += f"{len(sample_results)} sample(s)) all failed canonical]\n"
            feedback += f"Last error:\n{sample_results[-1][2] if sample_results else '(no output)'}\n"
        except Exception as e:
            err = f"{type(e).__name__}: {e}"
            log.append({"attempt": attempt, "model": model,
                        "temperature": temperature,
                        "agent_tests_passed": False, "canonical_passed": False,
                        "exception": err})
            feedback = f"[Attempt {attempt} crashed]\n{err}\n"

      return {
        "task_id": task_id, "passed": False, "attempts": MAX_RETRIES,
        "solution": solution_code, "log": log, "final_feedback": feedback,
      }
    finally:
        shutil.rmtree(work_dir, ignore_errors=True)


def _worker(problem, q):
    try:
        q.put(solve_problem(problem))
    except Exception as e:
        q.put({
            "task_id": problem["task_id"], "passed": False,
            "error": f"Worker exception: {e}\n{traceback.format_exc()}",
        })


def solve_with_timeout(problem: dict) -> dict:
    q = mp.Queue()
    p = mp.Process(target=_worker, args=(problem, q))
    p.start()
    p.join(TIME_LIMIT)
    if p.is_alive():
        p.terminate(); p.join(2)
        if p.is_alive():
            p.kill()
        return {
            "task_id": problem["task_id"], "passed": False,
            "error": f"Wall-clock timeout after {TIME_LIMIT}s",
        }
    if not q.empty():
        return q.get()
    return {"task_id": problem["task_id"], "passed": False, "error": "No result from worker"}


def main():
    args = sys.argv[1:]
    start, count = 0, None
    if len(args) == 1:
        count = int(args[0])
    elif len(args) >= 2:
        start, count = int(args[0]), int(args[1])

    problems = [json.loads(line) for line in DATA.open()]
    if count is not None:
        problems = problems[start:start + count]
    elif start:
        problems = problems[start:]

    passed = 0
    t_start = time.time()
    with OUT.open("a") as f:
        for i, prob in enumerate(problems, 1):
            t0 = time.time()
            res = solve_with_timeout(prob)
            res["elapsed"] = round(time.time() - t0, 1)
            if res.get("passed"):
                passed += 1
            print(
                f"[{i:>3}/{len(problems)}] {res['task_id']:<14} "
                f"pass={str(res.get('passed', False)):<5} "
                f"attempts={res.get('attempts', '-')} "
                f"t={res['elapsed']:>5}s"
                + (f"  ERR={res.get('error','')[:60]}" if not res.get('passed') and res.get('error') else "")
            )
            f.write(json.dumps(res) + "\n"); f.flush()

    elapsed = time.time() - t_start
    n = len(problems)
    print(
        f"\n=== Summary: {passed}/{n} passed "
        f"({100*passed/n:.1f}%)  total {elapsed:.0f}s "
        f"({elapsed/n:.1f}s/problem) ==="
    )


if __name__ == "__main__":
    mp.set_start_method("fork", force=True)
    main()