src/agent/repair_loop.py

"""Phase 6: agentic generate -> execute -> repair loop.

Model-agnostic: takes a `generate_fn(intent, feedback) -> code` callable, so it
works with CodeAssistant or any other generator (and is unit-testable with a mock).
"""
from __future__ import annotations

import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

sys.path.append(str(Path(__file__).resolve().parents[2]))
from src.eval.sandbox import run_code  # noqa: E402


@dataclass
class RepairTrace:
    final_code: str
    success: bool
    iterations: int
    history: list = field(default_factory=list)  # list of (code, error)


def repair_loop(
    intent: str,
    generate_fn: Callable[[str, str | None], str],
    check_program_fn: Callable[[str], str],
    max_iters: int = 3,
    timeout: float = 8.0,
) -> RepairTrace:
    """Iteratively generate and self-correct.

    generate_fn(intent, feedback)  -> candidate code
        feedback is None on the first call, else the previous error string.
    check_program_fn(code)         -> a runnable program string (code + a smoke
        test or the harness test) used to decide pass/fail.
    """
    feedback = None
    history = []
    for i in range(1, max_iters + 1):
        code = generate_fn(intent, feedback)
        result = run_code(check_program_fn(code), timeout=timeout)
        history.append((code, result.error))
        if result.ok:
            return RepairTrace(code, True, i, history)
        feedback = result.error  # feed the traceback back in
    return RepairTrace(history[-1][0], False, max_iters, history)


def make_repair_generator(assistant):
    """Adapt a CodeAssistant into a generate_fn for the repair loop."""
    def generate_fn(intent: str, feedback: str | None) -> str:
        if feedback:
            intent = (f"{intent}\n\n# Your previous attempt failed with this error:\n"
                      f"# {feedback}\n# Fix it and return the corrected function.")
        return assistant.generate(intent, mode="rag")
    return generate_fn