self_correct.py

import os
import subprocess
import sys
import tempfile
from pathlib import Path
from openai import OpenAI

client = OpenAI()
MODEL = os.getenv("OPENAI_MODEL", "gpt-4.1-mini")
MAX_ATTEMPTS = 3

TASK = """
Write a Python function unique_items(items).
Remove duplicates while preserving original order.
Return only Python code. No markdown.
"""

TESTS = """
def check(name, actual, expected):
    assert actual == expected, (
        f"{name}: expected {expected!r}, got {actual!r}"
    )

check("order", unique_items([3, 1, 3, 2, 1]), [3, 1, 2])
check("empty", unique_items([]), [])
check("duplicates", unique_items([5, 5, 5]), [5])
check("strings", unique_items(["b", "a", "b"]), ["b", "a"])
check("unique", unique_items([1, 2, 3]), [1, 2, 3])
"""

def generate(prompt):
    response = client.responses.create(
        model=MODEL,
        input=prompt,
    )
    code = response.output_text.strip()

    if code.startswith("```"):
        code = "\n".join(code.splitlines()[1:-1])

    return code.strip()


def validate(code):
    with tempfile.TemporaryDirectory() as folder:
        path = Path(folder) / "solution.py"
        path.write_text(code + "\n\n" + TESTS)

        try:
            result = subprocess.run(
                [sys.executable, "-I", str(path)],
                capture_output=True,
                text=True,
                timeout=5,
            )
            feedback = result.stderr or result.stdout
            return result.returncode == 0, feedback

        except subprocess.TimeoutExpired:
            return False, "Execution timed out"


def self_correct(task):
    prompt = task

    print("=" * 52)
    print("           SELF-CORRECTING AI")
    print("=" * 52)
    print(f"Model: {MODEL}")
    print("Task: Remove duplicates while preserving order")
    print(f"Maximum attempts: {MAX_ATTEMPTS}")

    for attempt in range(1, MAX_ATTEMPTS + 1):
        print(f"\n[ATTEMPT {attempt}/{MAX_ATTEMPTS}]")
        print("[AI] Generating Python code...")

        code = generate(prompt)

        print("\n[GENERATED CODE]")
        print(code)

        print("\n[TEST] Running 5 test cases...")
        passed, feedback = validate(code)

        if passed:
            print("[PASS] All 5 tests passed!")
            print("\n[SUCCESS] Validated solution returned.")
            return code

        print("[FAIL] Validation failed.")
        print(feedback[-1200:].strip())

        if attempt < MAX_ATTEMPTS:
            print("\n[REPAIR] Sending test failure to model...")

            prompt = f"""
Original task:
{task}

Your previous implementation:
{code}

Actual test failure:
{feedback[-3000:]}

Fix the implementation based on the error.
Preserve all original requirements.
Return only executable Python code.
"""

    raise RuntimeError(
        f"Failed to produce a passing solution "
        f"after {MAX_ATTEMPTS} attempts."
    )


if __name__ == "__main__":
    try:
        self_correct(TASK)
    except Exception as error:
        print(f"\n[ERROR] {error}")
        sys.exit(1)