Agentic Code Generation

Use the PAIS API with an agentic harness to write, test, and iterate on research software — automating the inner loop of exploratory coding.

Pattern: code-and-verify loop

flowchart LR
    Spec[Task description] --> LLM[PAIS LLM]
    LLM --> Code[Generated code]
    Code --> Run[Execute in sandbox]
    Run --> Pass{Tests pass?}
    Pass -->|No| LLM
    Pass -->|Yes| Output[Final code]

This pattern is useful when you need to: - Prototype a data processing script quickly - Generate boilerplate for a standard ML training loop - Translate a pseudocode algorithm into runnable Python - Add error handling or tests to existing code

Prerequisites

pip install openai python-dotenv

Simple code-and-fix agent

# code_agent.py
"""
A minimal code-and-verify agent using PAIS.
Generates Python code for a task, runs it, fixes errors, and iterates.
"""
import os
import sys
import subprocess
import tempfile
import textwrap
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()

client = OpenAI(
    base_url=os.environ["PAIS_API_BASE"],
    api_key=os.environ["PAIS_API_KEY"],
)

SYSTEM_PROMPT = """You are an expert Python developer writing research software.
When asked to write code:
1. Write clean, runnable Python
2. Include only necessary imports
3. Add brief inline comments for non-obvious logic
4. Handle common edge cases

When given an error, diagnose it and produce corrected code.
Always output ONLY the Python code — no markdown, no explanations."""


def extract_code(text: str) -> str:
    """Strip markdown code fences if the model adds them."""
    if "```python" in text:
        return text.split("```python")[1].split("```")[0].strip()
    if "```" in text:
        return text.split("```")[1].split("```")[0].strip()
    return text.strip()


def run_code(code: str, timeout: int = 30) -> tuple[bool, str]:
    """Execute code in a subprocess and return (success, output)."""
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
        f.write(code)
        fname = f.name

    result = subprocess.run(
        [sys.executable, fname],
        capture_output=True,
        text=True,
        timeout=timeout,
    )
    os.unlink(fname)

    if result.returncode == 0:
        return True, result.stdout
    return False, result.stderr


def code_agent(task: str, max_iterations: int = 5) -> str:
    """Generate code for a task, run it, fix errors, iterate."""
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Write Python code to: {task}"},
    ]

    for iteration in range(max_iterations):
        response = client.chat.completions.create(
            model="llama-3-1-8b",
            messages=messages,
            temperature=0.1,
            max_tokens=2048,
        )
        code = extract_code(response.choices[0].message.content)

        print(f"\n--- Iteration {iteration + 1} ---")
        print(code)

        success, output = run_code(code)
        if success:
            print(f"\nOutput:\n{output}")
            print("Code runs successfully.")
            return code

        print(f"\nError:\n{output}")
        messages.append({"role": "assistant", "content": code})
        messages.append({
            "role": "user",
            "content": f"This code produced the following error:\n{output}\n\nFix it.",
        })

    raise RuntimeError(f"Code still failing after {max_iterations} iterations")


# ── Example tasks ──────────────────────────────────────────────────────────────

if __name__ == "__main__":
    # Example 1: Data processing
    task = """Load a CSV file called 'experiments.csv' with columns
    [run_id, model_name, val_loss, val_accuracy, epochs].
    Find the top 3 runs by val_accuracy for each model_name.
    Print a formatted table."""

    code = code_agent(task)

    with open("process_experiments.py", "w") as f:
        f.write(code)
    print("\nCode saved to process_experiments.py")

Research-specific tasks

The agent handles common research coding patterns well:

Training loopData pipelineAnalysis scriptAPI integration

code_agent("""
Write a PyTorch training loop for a 3-layer MLP on MNIST.
Use AdamW optimiser, cosine LR schedule, and early stopping
on validation loss (patience=5). Log loss and accuracy per epoch.
""")

code_agent("""
Write a data pipeline that:
1. Loads JSONL files from ./data/ where each line has {"text": ..., "label": ...}
2. Tokenises text using HuggingFace tokenizers (bert-base-uncased)
3. Creates a PyTorch DataLoader with batch_size=32 and padding
4. Prints dataset statistics (size, label distribution, avg token length)
""")

code_agent("""
Load MLflow runs from experiment 'bert_finetune'.
For each run, plot val_loss vs epoch using matplotlib.
Save to mlflow_runs.png. Group by the 'learning_rate' parameter.
""")

code_agent(f"""
Write a function that takes a list of scientific paper abstracts
and returns a list of (abstract, embedding_vector) tuples using
the OpenAI embeddings API at {os.environ.get('PAIS_API_BASE')}.
Use the model 'qwen3-vl-embedding-8b' and batch in groups of 32.
""")

Using Claude Code for larger projects

For multi-file projects or iterative development across sessions, use Claude Code instead of a single-script agent:

cd ~/research/my-project
claude

> Build a training pipeline for fine-tuning Llama 3.1 on my dataset in ./data/.
  Use LoRA via the PEFT library, log metrics to MLflow at $MLFLOW_TRACKING_URI,
  and save checkpoints to ./checkpoints/. The dataset is JSONL with
  {"instruction": ..., "response": ...} format.

Claude Code will scaffold the full project, write tests, install dependencies, and iterate until the code runs. See the Claude Code guide for setup.

Prompt engineering for code generation

A few patterns that improve output quality from Llama 3.1 8B:

Be explicit about imports: "Use only numpy, pandas, and matplotlib. Do not use sklearn."

Specify the interface: "Write a function process_batch(texts: list[str]) -> list[dict] that..."

Provide sample data inline: "The CSV has this structure: run_id,val_loss\n1,0.234\n2,0.198"

Request error handling explicitly: "Include error handling for missing files and API timeouts."