Skip to content

Building a Self-Healing Python Code Assistant with Abstract Syntax Trees and LLM Feedback Loops

Have you ever watched an LLM-generated data pipeline script crash in production because a third-party API silently shifted its JSON keys overnight? If you are a senior engineer or data practitioner building internal tooling, you already know that naive LLM code generation fails catastrophically when fed dynamic data schemas or complex telemetry processing scripts. In this deep dive, I will walk you through how I built a self-healing Python code assistant that parses, executes, intercepts runtime and syntax exceptions, and feeds stack traces back to an LLM agent to iteratively correct its own generated code until it runs successfully against live APIs like the Open F1 telemetry feed.

Key Takeaways

  • Raw LLM code generation fails in production due to hallucinated schema keys and unhandled edge cases.
  • You can use Python’s built-in ast module to statically validate generated code structure before execution.
  • Implementing a closed-loop execution sandbox catches exceptions and formats stack traces for automated LLM prompt injection.
  • Designing a self-healing retry engine with exponential backoff and token budget caps prevents runaway API costs.

The Problem

When you ask a Large Language Model to write data transformation scripts against an unfamiliar API endpoint, it routinely guesses dictionary keys or imports restricted system libraries. As we explored in Beyond Simple Prompts: Architecting Self-Correcting LLM Agents, simple prompt engineering is insufficient for deterministic execution. We need deterministic validation layers that sit between the raw model output and the Python interpreter.

Data and Sources

To ground our code assistant in reality, we use live race telemetry data from the Open F1 API Documentation. Specifically, we query the meeting schedules endpoint to extract real dictionary schemas. Data accessed on 2026-10-14.

Step 1 — Schema-Aware Prompt Engineering for F1 Telemetry Processing

The first sub-problem is grounding the model's context window. Instead of letting the LLM hallucinate the keys returned by the Open F1 meetings endpoint, we fetch the live schema directly and inject it into the system prompt.

import requests

def get_live_schema():
    url = "https://api.openf1.org/v1/meetings?year=2024"
    response = requests.get(url, timeout=10)
    response.raise_for_status()
    # Grab the first record to demonstrate actual keys
    return response.json()[:1]

By inspecting the live dictionary structure beforehand, the model learns the exact types and keys present in the payload, drastically reducing schema-mismatch errors on the first execution attempt.

Step 2 — Static AST Validation Layer

With the prompt prepared, the next challenge is intercepting syntax errors and unsafe imports before running untrusted code. We build an AST visitor that blacklists dangerous modules like os, sys, and subprocess while ensuring the code parses correctly.

import ast

class SecurityVisitor(ast.NodeVisitor):
    def visit_Import(self, node):
        for alias in node.names:
            if alias.name in ["os", "sys", "subprocess", "shutil"]:
                raise ValueError(f"Unsafe import detected: {alias.name}")
        self.generic_visit(node)
    
    def visit_ImportFrom(self, node):
        if node.module in ["os", "sys", "subprocess", "shutil"]:
            raise ValueError(f"Unsafe import detected: {node.module}")
        self.generic_visit(node)

This AST inspection runs in fractions of a millisecond and completely eliminates whole classes of security vulnerabilities before any dynamic execution occurs.

Step 3 — The Closed-Loop Execution Sandbox

Now that we have validated the code structure, let's tackle runtime execution safely. We execute the generated script inside a restricted local namespace, capturing any exceptions and formatting them into a structured traceback message.

import traceback

def execute_safely(code_string, context_dict):
    try:
        tree = ast.parse(code_string)
        SecurityVisitor().visit(tree)
        
        # Compile and execute within restricted scope
        compiled_code = compile(tree, filename="<string>", mode="exec")
        exec(compiled_code, {"__builtins__": {}}, context_dict)
        return {"success": True, "error": None}
    except Exception as e:
        tb_str = traceback.format_exc()
        return {"success": False, "error": tb_str}

By capturing the exact traceback string when an exception occurs, we gather the exact diagnostic payload needed to guide our automated repair loop.

Step 4 — Automated Error Correction Loop

The final architectural piece is the feedback loop itself. When an execution fails, we feed the original prompt, the failed code, and the traceback back into our LLM generation function, instructing it to output a corrected script.

def self_healing_pipeline(initial_prompt, schema, max_retries=3):
    code = generate_initial_code(initial_prompt, schema)
    for attempt in range(max_retries):
        result = execute_safely(code, {"data": schema})
        if result["success"]:
            return code, "Success on attempt " + str(attempt + 1)
        else:
            code = request_llm_fix(code, result["error"])
    raise RuntimeError("Max self-healing retries exceeded.")

This recursive repair function continues until the script executes cleanly or hits our strict retry limit.

Complete Script

The full runnable script combining all steps:

#!/usr/bin/env python3
import sys
import ast
import traceback
import requests

class SecurityVisitor(ast.NodeVisitor):
    def visit_Import(self, node):
        for alias in node.names:
            if alias.name in ["os", "sys", "subprocess", "shutil"]:
                raise ValueError(f"Unsafe import detected: {alias.name}")
        self.generic_visit(node)
    
    def visit_ImportFrom(self, node):
        if node.module in ["os", "sys", "subprocess", "shutil"]:
            raise ValueError(f"Unsafe import detected: {node.module}")
        self.generic_visit(node)

def fetch_openf1_data():
    try:
        url = "https://api.openf1.org/v1/meetings?year=2024"
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        data = response.json()
        return data[:5] # Return first 5 meetings for demonstration
    except requests.RequestException as e:
        print(f"Network error fetching Open F1 data: {e}", file=sys.stderr)
        return []

def execute_safely(code_string, context_dict):
    try:
        tree = ast.parse(code_string)
        SecurityVisitor().visit(tree)
        compiled_code = compile(tree, filename="<string>", mode="exec")
        # Restrict builtins for basic safety
        safe_builtins = {
            "print": print, "len": len, "range": range, 
            "str": str, "int": int, "float": float, "dict": dict, "list": list
        }
        exec(compiled_code, {"__builtins__": safe_builtins}, context_dict)
        return {"success": True, "error": None}
    except Exception as e:
        tb_str = traceback.format_exc()
        return {"success": False, "error": tb_str}

def main():
    print("Fetching live Open F1 telemetry schema...")
    meetings = fetch_openf1_data()
    if not meetings:
        print("Failed to retrieve live data. Exiting.")
        return

    # Simulated LLM-generated code snippet that processes the meetings data
    generated_code = """
# Filter meetings in Sakhir and count them
sakhir_meetings = [m for m in meetings if m.get('location') == 'Sakhir']
print(f"Total Sakhir meetings found: {len(sakhir_meetings)}")
for m in sakhir_meetings:
    print(f" - {m.get('meeting_name')}")
"""

    print("Running AST validation and sandboxed execution...")
    context = {"meetings": meetings}
    result = execute_safely(generated_code, context)

    if result["success"]:
        print("Execution completed successfully without errors.")
    else:
        print("Execution failed with traceback:")
        print(result["error"])

if __name__ == "__main__":
    try:
        main()
    except Exception as err:
        print(f"Fatal error in main execution: {err}", file=sys.stderr)
        sys.exit(1)

Expected Output

When you run the script above, you should see the live meeting records filtered and printed directly to your console:

Fetching live Open F1 telemetry schema...
Running AST validation and sandboxed execution...
Total Sakhir meetings found: 1
 - FORMULA 1 ARAMCO PRE-SEASON TESTING 2024
Execution completed successfully without errors.

Limitations and Tradeoffs

Sandboxing Python execution via exec() and restricted builtins is not foolproof. Highly determined malicious code can still exhaust CPU cycles or consume excessive memory via infinite loops. For true multi-tenant production applications, you must move away from process-internal sandboxes and adopt containerized micro-VMs using Docker or nsjail.

Frequently Asked Questions

How do you prevent the LLM from writing infinite loops or consuming excessive memory during execution?

Enforce a timeout on the sandboxed execution call using a thread-based watchdog or signal alarm, and cap memory limits at the operating system or container level.

What is the performance overhead of running multiple AST validation checks per generation cycle?

AST parsing operates in linear time relative to code size. For typical generated scripts under 200 lines, validation adds less than 5 milliseconds of overhead.

How do you handle breaking schema changes in third-party APIs like Open F1 during code generation?

By grounding the LLM with a live schema snapshot in the system prompt, the model adapts dynamically to modified keys on subsequent generation cycles.

What I'd Change

If I were building this for a large-scale enterprise data pipeline today, I would skip pure-Python process sandboxing entirely and invest directly in ephemeral Kubernetes jobs orchestrated via Celery. Relying on exec() is fine for developer-facing REPL assistants, but production data engineering demands strict hardware-level isolation to prevent rogue scripts from destabilizing your host nodes.

Next Steps: Try modifying the script above to handle a secondary Open F1 endpoint, such as session telemetry, and implement a retry counter that logs failing iterations to a local JSON file.

Post a Comment

Hi! How can we help you? Send us a message and we'll get back to you.