In partnership with

You've been here before. You ask an LLM to write a function, it gives you something that looks fine, and then you run your tests and everything breaks.

So you copy the error, paste it back, get a new version, and repeat the same dance (sometimes 4-5 times if the model is feeling stubborn).

But what if the AI could run that entire loop on its own???

That's what we're building today: a small self-correcting system that writes Python code, tests it, and uses the errors from those tests to fix itself.

And we're doing it in plain Python, without LangChain or any agent framework.

But first, a word from our sponsors on taking some guesswork out of marketing.

Presented by The Marketing Millennials

A free newsletter with the marketing ideas you need

The best marketing ideas come from marketers who live it.

That’s what this newsletter delivers.

The Marketing Millennials is a look inside what’s working right now for other marketers. No theory. No fluff. Just real insights and ideas you can actually use—from marketers who’ve been there, done that, and are sharing the playbook.

Every newsletter is written by Daniel Murray, a marketer obsessed with what goes into great marketing. Expect fresh takes, hot topics, and the kind of stuff you’ll want to steal for your next campaign.

Because marketing shouldn’t feel like guesswork. And you shouldn’t have to dig for the good stuff.

Sign Up Free

Presented by Roku Ads Manager

The holiday streaming season is almost here

Are you ready? As revealed in our holiday playbook, Roku Ads Manager lets you easily set up, measure and optimize your self-serve holiday ad campaign. You can even create a shoppable ad unit that lets streamers buy directly from their remote.

Set up your campaign today.

How does self-correction work?

The idea comes from iterative refinement.

Research like Self-Refine showed that LLMs can improve their own outputs when you keep feeding them feedback (if you want to go deeper, here's their reference implementation).

We're tweaking it a bit though. 

Self-Refine has the model critique its own work, which is a bit like asking a student to grade their own exam. We're going to run proper tests on the code and let the errors do the talking.

Here's what the workflow looks like:

Whenever a test fails, the error goes back to the model, it takes another shot at the code, and the tests run again.

Let's get into it.

Step 1: Give the AI a task

We'll ask the model to write a function that removes duplicates from a list while keeping the original order intact.

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

So unique_items([3, 1, 3, 2, 1]) should give us [3, 1, 2].

The classic mistake here is list(set(items)). It does remove duplicates, but sets don't care about order, so your list can come back shuffled. Our tests are going to catch that.

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"])
"""

Four tests covering ordering, empty lists, repeated values, and strings.

Step 2: Build the self-correcting loop

Install the OpenAI SDK:

pip install openai

Set your OPENAI_API_KEY environment variable and create a main.py with the task and tests from above.

First up, a function that generates the code:

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

client = OpenAI()

def generate(prompt):
    response = client.responses.create(
        model="gpt-6-luna",
        input=prompt
    )
    code = response.output_text.strip()
    if code.startswith("```"):
        code = "\n".join(code.splitlines()[1:-1])
    return code.strip()

Next, we need something that runs that code against our tests:

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

        try:
            result = subprocess.run(
                [sys.executable, "-I", str(file)],
                capture_output=True,
                text=True,
                timeout=5
            )
            return (
                result.returncode == 0,
                result.stderr or result.stdout
            )
        except subprocess.TimeoutExpired:
            return False, "Execution timed out"

The validator hands back two things: whether the tests passed, and the error message if they didn't.

Quick warning: this runs AI-written code on your machine. Isolated mode, temp files and timeouts help, but they don't make it safe. 

Only run this in a throwaway sandbox with no network access, no secrets, and nothing you'd cry about losing.

Step 3: Let the AI fix its mistakes

Now we plug generation and validation together.

def self_correct(task, max_attempts=3):
    prompt = task

    for attempt in range(1, max_attempts + 1):
        print(f"Attempt {attempt}")

        code = generate(prompt)
        passed, error = validate(code)

        if passed:
            print("All tests passed!")
            return code

        print("Failed:", error[-500:])

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

        Previous code:
        {code}

        Test failure:
        {error[-2000:]}

        Fix the code.
        Return only executable Python code.
        """

    raise RuntimeError("All attempts failed")

print(self_correct(TASK))

And that's the whole feedback loop.

The first attempt is a normal generation. If the tests fail, the next prompt carries the broken code along with the error, so the model knows what went wrong.

It keeps going until the tests pass or it burns through 3 attempts.

Step 4: See the feedback loop

Here's what a successful repair could look like:

If you want to see the repair kick in, start with a broken version like list(set(items)) and feed its failure into the fix step.

Grab the full code

Can we use this beyond code generation?

Yes, anywhere the output can be checked automatically.

If you're generating SQL, you can check queries against your database schema.

For structured outputs, validate the JSON against a schema. And for code migrations, run your existing unit tests and send whatever fails back to the model.

Two catches though.

Passing tests only means the code passed the tests you wrote, so if your tests skip an edge case, the model will happily skip it too. And every retry costs you tokens and time.

If you want to take this further, compare one-shot generation against this loop on a fixed set of tasks and track pass rate, retry count, and latency.

My take

I think giving AI a way to catch its own mistakes will get you further than spending another hour polishing your prompt.

A simple loop with a good test suite covers the basics.

This shines for code and structured outputs, basically anything you can test.

So I'd start with a solid validator and a strict retry limit before adding anything fancier, because a weak validator just means you're paying for the model to confidently pass bad code.

AI Tools That Made Me Question My Life Choices This Week

Heal.dev

An AI bug hunter that runs your app in a sandbox, finds bugs it can reproduce, and writes failing tests so you can verify your fixes. Pretty much today's edition, packaged as a product.

Cekura

A testing platform for voice and chat agents. It simulates conversations, monitors your agents, and flags where they break.

Moli

An open-source headless browser built in Rust, made for lightweight web extraction with CDP and WebDriver automation support.

Until next time,
Vaibhav 🤝

Reply

Avatar

or to participate

Keep Reading

Read all
arrow-right
background
© 2026 Staying Ahead with AI.
beehiivPowered by beehiiv