/agent-lightning

name: agent-lightning
description: Train and optimize AI agents using Microsoft's Agent Lightning framework with reinforcement learning. Use when setting up agent training, instrumenting agents with tracing, configuring LightningStore, implementing reward functions, or optimizing prompts with RL/APO algorithms.
domain: engineering

Agent Lightning

Microsoft's framework for training AI agents with reinforcement learning, automatic prompt optimization, and supervised fine-tuning.

Quick Start

Installation

pip install agentlightning

For nightly builds:

pip install --upgrade --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ --pre agentlightning

Minimal Integration (Zero Code Change)

Add `agl.emit_xxx()` helpers to your existing agent:

import agentlightning as agl

# Your existing agent code
def my_agent(task):
    agl.emit_input(task)  # Track input
    
    response = llm.generate(task)
    agl.emit_output(response)  # Track output
    
    reward = evaluate(response)
    agl.emit_reward(reward)  # Track reward
    
    return response

Core Concepts

Architecture Flow

Agent (your code) → agl.emit_xxx() → Spans → LightningStore → Algorithm → Updated Resources

Key Components

| Component | Purpose | |-----------|---------| | `LightningStore` | Central hub for traces, tasks, and resources | | `Tracer` | Collects spans from agent execution | | `Algorithm` | Consumes traces, produces improvements | | `Trainer` | Orchestrates training loop |

Instrumentation

Emit Functions

import agentlightning as agl

# Basic emissions
agl.emit_input(prompt)           # Track input to agent
agl.emit_output(response)        # Track agent output
agl.emit_reward(score)           # Track reward signal
agl.emit_tool_call(name, args)   # Track tool usage
agl.emit_tool_result(result)     # Track tool results

Tracer Context

from agentlightning import Tracer

tracer = Tracer(store=store)

with tracer.trace_context(task_id="task-123"):
    # All emissions within this context are grouped
    result = agent.run(task)
    
# Retrieve trace after execution
trace = tracer.get_last_trace()

OpenTelemetry Integration

Agent Lightning integrates with OpenTelemetry:

from agentlightning.utils.otel import get_tracer

tracer = get_tracer()  # Returns OTel tracer for "agentlightning"

LightningStore

In-Memory Store (Development)

from agentlightning.store.memory import InMemoryLightningStore

store = InMemoryLightningStore()

Client-Server Store (Production)

from agentlightning.store.client_server import (
    LightningStoreServer,
    LightningStoreClient
)

# Server side
server = LightningStoreServer(store, host="0.0.0.0", port=8080)
await server.start()

# Client side
client = LightningStoreClient("http://localhost:8080")

Store Operations

# Add rollouts (tasks for the agent)
await store.enqueue_rollout(task=task, config=RolloutConfig())

# Query rollouts
rollouts = await store.query_rollouts(status_in=["completed"])

# Add resources (updated prompts, weights)
await store.add_resources(resources)

# Get latest resources
resources = await store.get_latest_resources()

Training

Basic Trainer Setup

import agentlightning as agl

trainer = agl.Trainer(
    n_runners=8,           # Parallel rollout workers
    algorithm=algorithm,   # Your chosen algorithm
    store=store           # Optional, creates InMemory if not provided
)

trainer.run()

Custom Algorithm

from agentlightning import LightningStore
from agentlightning.types import ExecutionEvent

async def my_algorithm(store: LightningStore, event: ExecutionEvent):
    # Fetch completed rollouts
    rollouts = await store.query_rollouts(status_in=["completed"])
    
    # Process traces, compute gradients, etc.
    new_resources = optimize(rollouts)
    
    # Push updated resources
    await store.add_resources(new_resources)

Runner Function

async def my_runner(store: LightningStore, worker_id: int, event: ExecutionEvent):
    while not event.is_set():
        rollout = await store.dequeue_rollout()
        if rollout:
            result = execute_task(rollout.task)
            await store.update_rollout(
                rollout_id=rollout.id,
                status="completed",
                result=result
            )

Algorithms

Reinforcement Learning (GRPO/PPO)

For RL training with vLLM backend:

from agentlightning.algorithm.verl import VeRLAlgorithm

algorithm = VeRLAlgorithm(
    model="your-model",
    learning_rate=1e-5,
    batch_size=32
)

Automatic Prompt Optimization (APO)

from agentlightning.algorithm.apo import APOAlgorithm

algorithm = APOAlgorithm(
    optimizer_model="gpt-4",
    target_model="gpt-3.5-turbo"
)

Framework Adapters

LangChain

from agentlightning.instrumentation.langchain import instrument_langchain

instrument_langchain()  # Auto-traces all LangChain calls

OpenAI SDK

from agentlightning.instrumentation.openai import instrument_openai

instrument_openai()  # Auto-traces OpenAI API calls

vLLM

from agentlightning.instrumentation.vllm import instrument_vllm

instrument_vllm()  # Instrument vLLM for token-level tracing

Logging & Debugging

Configure Logging

from agentlightning import setup_logging

setup_logging(
    level="DEBUG",
    submodule_levels={
        "agentlightning.store": "INFO",
        "agentlightning.tracer": "DEBUG"
    }
)

Metrics

Agent Lightning emits Prometheus-compatible metrics:

• `agl.store.total` - Store operation counts
• `agl.store.latency` - Store operation latencies
• `agl.rollouts.total` - Rollout counts by status
• `agl.rollouts.duration` - Rollout execution times

Common Patterns

Re