loop.py

"""The engineered loop.

This is the whole point of the project. A naive "AI agent" is just the *body*
of this loop: act -> observe -> repeat. Loop ENGINEERING is everything wrapped
around that body:

  - ground-truth feedback   : pytest's exit code decides success, not the model
  - the maker/checker split  : the model that writes code does not grade it
  - four independent exits    : SUCCESS / BLOCKED / CAPPED / STUCK

Run it:
    .venv/bin/python loop.py solvable                 # converges -> SUCCESS
    .venv/bin/python loop.py unsolvable               # guardrails ON  -> STUCK fast
    .venv/bin/python loop.py unsolvable --no-guardrails   # ablation: spins to CAPPED

The --no-guardrails ablation is the teaching moment: same loop, same task, but
with no-progress detection turned off it wastes its entire iteration budget
re-trying a doomed fix instead of bailing early.
"""

import argparse
import hashlib
import re
import shutil
import subprocess
import sys
from pathlib import Path

from agent import CHECKER_MODEL, MAKER_MODEL, chat, extract_code

ROOT = Path(__file__).parent

# Each task: a buggy implementation (reset from .seed each run) + a fixed test
# file the agent is NOT allowed to touch (that would be cheating the oracle).
TASKS = {
    "solvable": {
        "impl": ROOT / "tasks/solvable/roman_numeral.py",
        "seed": ROOT / "tasks/solvable/roman_numeral.seed.py",
        "test": ROOT / "tasks/solvable/test_roman_numeral.py",
    },
    "unsolvable": {
        "impl": ROOT / "tasks/unsolvable/widget.py",
        "seed": ROOT / "tasks/unsolvable/widget.seed.py",
        "test": ROOT / "tasks/unsolvable/test_widget.py",
    },
}

MAKER_SYSTEM = (
    "You are a Python engineer. You will be given the full contents of a module "
    "and the failing-test output for it. Fix the module so every test passes. "
    "Do NOT change the tests. Reply with the COMPLETE updated module as a single "
    "```python code block and nothing else."
)

CHECKER_SYSTEM = (
    "You are a strict, adversarial code reviewer. The test suite already passes. "
    "Your job is to find what is still WRONG: missing edge cases, spec violations, "
    "or code that games the tests instead of solving the problem. Be concise. "
    "If it is genuinely correct, say so plainly."
)


def run_pytest(test_file):
    """Run the suite. Returns (passed, signature, failure_feedback).

    `signature` is a hash of the failing assertions — the loop's memory of
    "what does broken look like right now", used for no-progress detection.
    `failure_feedback` is the one-line-per-failure summary we feed the maker:
    structured feedback, not a raw multi-thousand-line log dump.
    """
    proc = subprocess.run(
        [sys.executable, "-m", "pytest", str(test_file), "-q", "--tb=short", "-ra"],
        capture_output=True,
        text=True,
        cwd=ROOT,
    )
    passed = proc.returncode == 0
    # pytest's "short test summary info" prints one stable line per failure:
    #   FAILED path/test_x.py::test_name - assert 6 == 4
    # We drop the path dirs (keeping the filename) so the signature is stable
    # across machines, and use the test name + assertion as structured feedback.
    failed = sorted(
        re.sub(r"^FAILED\s+\S*/", "FAILED ", ln)
        for ln in proc.stdout.splitlines()
        if ln.startswith("FAILED")
    )
    if failed:
        feedback = "\n".join(failed)
    else:
        # No parseable failures (e.g. an import/collection error) — fall back to
        # the raw tail so the signature still reflects *what* broke.
        feedback = proc.stdout.strip()[-1500:]
    signature = hashlib.sha1(feedback.encode()).hexdigest()[:10]
    return passed, signature, feedback


def make(impl_file, feedback):
    """Maker step: ask the coding model to rewrite the module."""
    code = chat(
        MAKER_MODEL,
        MAKER_SYSTEM,
        f"Module `{impl_file.name}`:\n```python\n{impl_file.read_text()}\n```\n\n"
        f"Failing tests:\n{feedback}\n\nReturn the complete fixed module.",
    )
    impl_file.write_text(extract_code(code) + "\n")


def check(impl_file):
    """Checker step: a *different* model grades the maker's work adversarially."""
    return chat(
        CHECKER_MODEL,
        CHECKER_SYSTEM,
        f"Review this module that just passed its tests:\n"
        f"```python\n{impl_file.read_text()}\n```",
    ).strip()


def run(task_name, guardrails=True, max_iters=8, run_checker=False):
    task = TASKS[task_name]
    shutil.copy(task["seed"], task["impl"])  # reset to the buggy starting point

    print(f"\n=== loop: task={task_name}  guardrails={'ON' if guardrails else 'OFF'} "
          f"max_iters={max_iters} ===")

    seen = set()  # every failure signature we've encountered (no-progress memory)
    outcome = "CAPPED"

    for i in range(1, max_iters + 1):
        passed, sig, feedback = run_pytest(task["test"])

        # --- EXIT 1: SUCCESS -------------------------------------------------
        if passed:
            print(f"[iter {i}] tests PASS  -> SUCCESS")
            outcome = "SUCCESS"
            break

        print(f"[iter {i}] tests FAIL  sig={sig}  ({feedback.splitlines()[0] if feedback else '?'})")

        # --- EXIT 4: STUCK (no-progress detection) -------------------------
        # The single most important guardrail. If we've seen this exact failure
        # state before, the agent is re-trying a doomed approach. Bail now
        # instead of burning the rest of the budget. Disabled by --no-guardrails.
        if guardrails and sig in seen:
            print(f"[iter {i}] failure signature already seen -> STUCK "
                  f"(no progress, exiting early)")
            outcome = "STUCK"
            break
        seen.add(sig)

        # --- loop body: act -------------------------------------------------
        print(f"[iter {i}] maker ({MAKER_MODEL}) attempting a fix...")
        make(task["impl"], feedback)

    # --- EXIT 3: CAPPED is the fall-through when the for-loop ends ----------
    if outcome == "CAPPED":
        print(f"[done] hit max_iters={max_iters} without passing -> CAPPED")

    print(f"=== outcome: {outcome} ===")

    # Maker/checker split: only meaningful once tests are green.
    if outcome == "SUCCESS" and run_checker:
        print(f"\n--- checker ({CHECKER_MODEL}) adversarial review ---")
        print(check(task["impl"]))

    return outcome


def main():
    ap = argparse.ArgumentParser(description="An engineered agent loop.")
    ap.add_argument("task", choices=list(TASKS))
    ap.add_argument("--no-guardrails", action="store_true",
                    help="disable no-progress detection (the ablation)")
    ap.add_argument("--max-iters", type=int, default=8)
    ap.add_argument("--checker", action="store_true",
                    help="run the adversarial checker after SUCCESS")
    args = ap.parse_args()
    run(args.task, guardrails=not args.no_guardrails,
        max_iters=args.max_iters, run_checker=args.checker)


if __name__ == "__main__":
    main()