Deep Dive into Agent Lightning: Optimizing Prompts with APO

Enhancing SQL Agents with Agent Lightning Prompt Optimization

Microsoft recently released Agent Lightning, a framework that is the absolute trainer to light up AI agents. It introduces 2 algorithms:

• VERL: A Reinforcement Learning (RL) algorithm for policy-level optimization.
• APO (Automatic Prompt Optimization): A self-improving prompt tuning algorithm that learns textual “gradients” from model feedback.

This article focuses on applying APO to optimize prompts for better accuracy and consistency.

What is APO? The Math Behind Textual Gradients

Agent Lightning’s APO is based on two research papers:

• ProTeGi (EMNLP 2023)
• TextGrad (Nature 2024)

Instead of retraining a model or modifying its weights, APO uses a pair of cooperating LLMs for training:

• Critic model: analyzes execution traces and identifies what went wrong in failed cases.
• Editor model: rewrites the prompt based on the critic’s feedback to produce improved variants.

During APO training, each iteration generates new candidate prompts that are validated against your evaluation dataset, and the best ones are preserved through beam search. This process mimics gradient-based learning but entirely in text space.

Formally:

where each update refines the prompt p_t​ according to critiques derived from training data D_train​

How It Works

APO combines three key techniques: 1. Textual Gradient Computation

• The agent runs the current prompt (p) on a batch of tasks (B).
• For each input xᵢ, it generates an output fₚ(xᵢ) and scores it using a reward Rᵢ (eg, SQL correctness).
• The critic LLM reviews all (xᵢ, fₚ(xᵢ), Rᵢ) results, finds common mistakes, and produces a textual gradient. A short critique like:

The prompt doesn’t specify how to handle type mismatches in JOIN columns. When Singer_ID is INTEGER in one table but TEXT in another, the agent should apply CAST with numeric filtering. Add explicit rule: ‘If joining columns have different types, use CAST(col_text AS INTEGER) with GLOB ‘[0–9]’ filter.

• This critique g(p, B) is then passed to the editor model, which rewrites the prompt into a refined version p′.

2. Gradient-Based Editing Apply the textual gradient to generate improved prompts:

Example transformation: Before (vague):

"Be careful to not query for columns that do not exist."

After (explicit):

**Success Criteria:**
- The query must reference only tables and columns explicitly listed in {table_info}.
- Verify all referenced tables and columns exist before responding.

3. Beam Search Maintain a beam of top-k k prompts ranked by validation score:

This ensures exploration of multiple refinement directions rather than following a single, potentially suboptimal path.

Complete APO Algorithm Flow

Why APO?

APO isn’t just another prompt engineering tool; it’s a fundamentally different approach to agent optimization:

1. Data-Grounded Improvement

Unlike manual prompt engineering based on intuition, APO is grounded in real data:

• Analyzes actual failures from your training set
• Critiques are based on execution traces, not speculation
• Rewards come from your evaluation function (e.g., SQL correctness) Optimization targets your specific task distribution

Example training:

The baseline prompt said “be careful not to query for columns that do not exist”, a vague warning. After analyzing failures where the agent used Singer_ID (TEXT) to join with Singer_ID (INTEGER), APO generated this critique:

Add explicit rule for type-safe joins: When columns have different types, use CAST(col_text AS INTEGER) with numeric filtering GLOB ‘[0–9]’.

This was derived from analysing the specific failure patterns in the training data.

2. Systematic Exploration

Beam search prevents getting stuck in local optima. By maintaining multiple candidates, APO explores different improvement directions simultaneously.

Single-path (gradient descent):
p₀ → p₁ → p₂ → p₃ (stuck at local max: 85%)

Beam search:
         → p₁ (84%) → p₃ (85%) → p₅ (86%)
p₀ (84%)
         → p₂ (85%) → p₄ (87%) → p₆ (88%) ← global max!

3. Interpretable & Debuggable

Every optimization step is explainable:

• Textual gradients are human-readable critiques
• You can inspect why each change was made
• Easy to validate if improvements make sense
• Can manually override or guide the process

In my training, I could read each critique and see exactly what APO learned from failures. This transparency is invaluable for debugging and trust.

4. Not “LLM-as-a-Judge”

Instead of relying on an LLM to score outputs, APO bases its rewards on objective, measurable outcomes defined by your evaluation function.

In this setup, the LLM’s role is limited to analyzing failures and suggesting improvements, while the reward signal comes from actual execution results.

Now that we’ve explored why APO matters, let’s move from concept to application.

The Use Case: LangGraph Text-to-SQL Agent

For this experiment, I used the Spider dataset, with 50 examples for training and 50 for validation, both sampled from different parts of the development split to a total of 100 examples out of 10,181.

Agent Architecture

The SQL agent is built with LangGraph and follows a self-correction workflow:

System Architecture with Agent Lightning

Here’s how Agent Lightning integrates with the SQL agent:

You only implement the @rollout function, Agent Lightning handles everything else!

Here’s the complete APO setup

from agentlightning import Trainer
from agentlightning.algorithm.apo import APO
from openai import AsyncOpenAI

# 1. Initialize OpenAI client
openai_client = AsyncOpenAI()

# 2. Configure APO algorithm
algo = APO(
    openai_client,
    val_batch_size=10,       # Quick validation (beam selection)
    gradient_batch_size=4,   # Training batch for critique
    beam_width=2,            # Top-2 prompts in beam
    branch_factor=2,         # 2 variants per prompt
    beam_rounds=2,           # 2 optimization rounds
)

# 3. Create trainer
trainer = Trainer(
    algorithm=algo,
    n_runners=8,             # Parallel execution
    initial_resources={
        "prompt_template": prompt_template_baseline()
    }
)

# 4. Load datasets
train_dataset = load_spider_dataset("data/dev.json")[:50]
val_dataset = load_spider_dataset("data/dev.json")[50:100]

# 5. Train!
trainer.fit(
    agent=sql_agent_rollout,  # Your @rollout function
    train_dataset=train_dataset,
    val_dataset=val_dataset
)

That’s it! The @rollout function looks like this:

from agentlightning import rollout
from agentlightning.types import PromptTemplate

@rollout
def sql_agent_rollout(task, prompt_template: PromptTemplate) -> float:
    """
    Task format:
    {
        "question": "Show me singers who performed in 2014",
        "db_id": "concert_singer",
        "query": "SELECT DISTINCT s.Name FROM singer s ..."
    }
    """
    # 1. Build agent with prompt template
    agent = SQLAgent(
        db_path=f"databases/{task['db_id']}/{task['db_id']}.sqlite",
        write_prompt=prompt_template.format(
            dialect="SQLite",
            table_info=get_schema(task['db_id'])
        )
    )

    # 2. Run agent
    result = agent.run(task["question"])
    predicted_query = result["query"]

    # 3. Evaluate
    reward = evaluate_query(
        predicted_query,
        task["query"],
        task["db_id"]
    )  # Returns 1.0 or 0.0

    return reward

Baseline Prompt (v0)

The starting prompt was intentionally simple; it is a typical first draft:

You are an agent designed to interact with a SQL database.
Given an input question, create a syntactically correct {dialect} query 
to run to help find the answer.

Pay attention to use only the column names that you can see in the 
schema description. Be careful to not query for columns that do not 
exist. Also, pay attention to which column is in which table.

## Table Schema ##

Only use the following tables:
{table_info}

## Output Format ##

Respond in the following format:
```{dialect}
GENERATED QUERY

**Baseline performance:** 42/50 correct (84.0%)

## Training Results: Accuracy improved from 84% to 88%

![](https://miro.medium.com/v2/resize:fit:1176/1*rWV4sh2FzoHvH3j0MV0ZTw.png)

The SQL agent’s prompt improved steadily across two optimization rounds.

**Progress Summary :**
- Round 0 (Baseline): 84 % accuracy
- Round 1 (First optimized prompt): 86 % accuracy (+2 %)
- Round 2 (Final optimized prompt): 88 % accuracy (+4 %)

You can find the detailed training logs and intermediate evaluations in `apo.log` for full traceability of each optimization round.

Round 0: Baseline v0: 84% (baseline prompt)

Round 1: Generate 4 variants v0 → v1: 80% (added schema-first pipeline) v0 → v2: 80% (added casting rules) v0 → v3: 90% → Full: 86% ⭐ New best! v0 → v4: 90%

Select top-2 for beam: [v3, v4]

Round 2: Generate 4 more variants v3 → v7: 90% v3 → v8: 90% v4 → v5: 100% → Full: 88% 🏆 WINNER! v4 → v6: 100% → Full: 86%

Select top-2: [v5, v6]

Final: v5 = 88% accuracy

After two APO optimization rounds, the final prompt** (v5)** became far more detailed, structured, and robust compared to the baseline** (v0)**.

Key Differences:
- Accuracy: Improved from 84.0 % (42/50) → 88.0 % (44/50)
- Length: Expanded from 90 words → 360 words
- Structure: Grew from 2 informal sections → 7 clearly structured parts
- Explicit Rules: Increased from 3 vague hints → 19 concrete requirements
- Validation Logic: Added 3 checklist-style verification steps that ensure SQL correctness
- Safety Constraints: Introduced 4 explicit rules to prevent unsafe or invalid SQL queries
- Error Handling: Added a UNABLE_TO_ANSWER pattern for unresolvable cases
- Default Behaviours: Defined 2 fallback strategies for incomplete schema information

**Optimized Prompt (v5)**

You are an agent that formulates a single, syntactically and semantically correct SQL query in {dialect} to answer the given question by querying only the tables and columns provided in the schema description.

Success Criteria:

• The query must be valid {dialect} SQL syntax.
• It must reference only tables and columns explicitly listed in {table_info}.
• It must answer the question as completely as possible given the schema.
• If limiting results, include an ORDER BY clause to ensure deterministic output and avoid ties.
• Return exactly one SQL statement per response, without comments or extra content.

Handling Ambiguity or Impossible Questions:

• If the question cannot be answered with the available schema or is ambiguous, do NOT guess or ask clarifying questions.
• Instead, respond with a single query that returns 'UNABLE TO ANSWER' or an equivalent empty result in {dialect}, clearly indicating inability to provide an answer.

Rule Precedence:

• When conflicts arise, prioritize schema correctness over dialect specifics (i.e., do not reference unavailable columns even if the dialect allows flexible syntax).
• Always maintain syntactic correctness according to the specified {dialect}.

Input Schema:Only use the following tables and columns:{table_info}

If {dialect} or {table_info} are not provided, default to dialect='sqlite' (lowercase) and an empty schema (no tables).

Output Format:Respond with exactly one code block using the syntax highlighting of the {dialect} (case-insensitive, e.g., sqlite, postgresql, mysql):

<SQL QUERY>

• The code block must contain only the finalized SQL statement starting immediately (no leading blank lines or comments).
• Do not include explanations, notes, or any text outside the code block.

Additional Constraints:

• Do not perform any destructive operations (e.g., INSERT, UPDATE, DELETE, DROP).
• When constructing JOINs, ensure keys exist and types are compatible.
• Queries should be as concise as possible while correctly answering the question.
• Apply default quoting and casing conventions according to {dialect}.

Verification Steps Before Responding:

• Verify all referenced tables and columns exist in {table_info}.
• Confirm the query answers the input question correctly or return the 'UNABLE TO ANSWER' query otherwise.
• Ensure deterministic ordering when LIMIT is applied.

Table Schema

{table_info}

# Summary

Through this deep dive, we saw how Agent Lightning’s APO can systematically improve a prompt, transforming a short, informal baseline into a structured, rule-driven specification that boosted SQL accuracy from 84 % to 88 % in just two optimization rounds using 50 training examples.

I’ve shared the complete functional source code and training example on [GitHub](https://github.com/yai333/SQL-Agent-with-APO-Automatic-Prompt-Optimization) so you can experiment. Enjoy exploring Agent ⚡️!