init documentation (#33)

.github/workflows/docs.yml

@@ -0,0 +1,59 @@
+name: Deploy Documentation
+
+on:
+  push:
+    branches:
+      - main
+    tags:
+      - 'v*'
+  workflow_dispatch:
+
+permissions:
+  contents: write
+  pages: write
+  id-token: write
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+
+      - name: Install dependencies
+        run: |
+          ./scripts/setup_stable.sh
+
+      - name: Configure Git
+        run: |
+          git config --global user.name "GitHub Actions"
+          git config --global user.email "actions@github.com"
+
+      - name: Get version and commit
+        id: version
+        run: |
+          if [[ $GITHUB_REF == refs/tags/* ]]; then
+            VERSION=${GITHUB_REF#refs/tags/v}
+            SOURCE_COMMIT=${GITHUB_SHA}
+          else
+            VERSION=$(python -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['project']['version'])")
+            SOURCE_COMMIT=${GITHUB_SHA}
+          fi
+          echo "version=$VERSION" >> $GITHUB_OUTPUT
+          echo "SOURCE_COMMIT=$SOURCE_COMMIT" >> $GITHUB_ENV
+
+      - name: Deploy versioned docs
+        if: startsWith(github.ref, 'refs/tags/')
+        run: |
+          mike deploy --push --update-aliases ${{ steps.version.outputs.version }} latest
+
+      - name: Deploy dev docs
+        if: github.ref == 'refs/heads/main'
+        run: |
+          mike deploy --push dev

.github/workflows/tests.yml

@@ -33,6 +33,33 @@ jobs:
         run: |
           black --check --diff --line-length=120 .
 
+  docs:
+    name: Build documentation
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'
+      - name: Install documentation dependencies
+        run: |
+          ./scripts/setup_stable.sh
+      - name: Set source commit for docs
+        run: |
+          echo "SOURCE_COMMIT=${{ github.sha }}" >> $GITHUB_ENV
+      - name: Build documentation
+        run: |
+          mkdocs build --strict
+      - name: Upload docs artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: documentation-site
+          path: site/
+          compression-level: 6
+
   test:
     strategy:
       matrix:

README.md

@@ -1,4 +1,4 @@
-![Agent-lightning-banner](assets/readme-banner.png)
+![Agent-lightning-banner](docs/assets/readme-banner.png)
 
 # Agent Lightning⚡
 
@@ -16,7 +16,7 @@
 - **Selectively** optimize one or more agents in a multi-agent system. 🎯
 - Embraces Reinforcement Learning, Automatic Prompt Optimization and more **algorithms**. 🤗
 
-![Agent-Lightning-code-diff](assets/readme-diff.png)
+![Agent-Lightning-code-diff](docs/assets/readme-diff.png)
 
 ## ⚡ Resources
 
@@ -107,7 +107,7 @@ Currently, Agent Lightning is built around a **training server** and one or mult
 * **Agents** retrieve samples from the server, process them (which may involve interacting with the LLM), and send the results back. These results, or "trajectories," are lists of prompts and responses from the LLM.
 * The **server** then collects these trajectories and computes the losses to optimize the language models.
 
-![Agent-Lightning-architecture](assets/readme-architecture.png)
+![Agent-Lightning-architecture](docs/assets/readme-architecture.png)
 
 ## ⚡ Development Instructions
 
@@ -126,6 +126,12 @@ pre-commit install
 pre-commit run --all-files --show-diff-on-failure --color=always
 ```
 
+Serve documentation locally:
+
+```bash
+mkdocs serve
+```
+
 ## ⚡ Contributing
 
 This project welcomes contributions and suggestions. Most contributions require you to agree to a Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.

agentlightning/verl/__init__.py

@@ -0,0 +1,3 @@
+from .trainer import *
+from .daemon import *
+from .dataset import *

agentlightning/verl/daemon.py

@@ -31,8 +31,8 @@ def get_left_padded_ids_and_attention_mask(ids: List[int], max_length: int, pad_
         pad_token_id:    ID to use for padding.
 
     Returns:
-        padded_ids:      list of length == max_length.
-        attention_mask:  list of same length: 1 for non-pad tokens, 0 for pads.
+        padded_ids (any):      list of length == max_length.
+        attention_mask (any):  list of same length: 1 for non-pad tokens, 0 for pads.
     """
     seq_len = len(ids)
 
@@ -60,8 +60,8 @@ def get_right_padded_ids_and_attention_mask(ids: List[int], max_length: int, pad
         pad_token_id:   ID to use for padding.
 
     Returns:
-        padded_ids:     list of length == max_length.
-        attention_mask: list of same length: 1 for non-pad tokens, 0 for pads.
+        padded_ids (any):     list of length == max_length.
+        attention_mask (any): list of same length: 1 for non-pad tokens, 0 for pads.
     """
     seq_len = len(ids)
 

docs/deep-dive/server-client-architecture.md

@@ -0,0 +1,31 @@
+# Server-client Architecture
+
+Article to be written.
+
+```mermaid
+sequenceDiagram
+    participant RL as RL Framework
+    participant TS as Training Server
+    participant AC as Agent Client
+    participant AG as Agent
+
+    AC->>TS: Upload Dataset (1)
+    RL->>TS: Start RL Server (2)
+    TS->>RL: Latest Model (3)
+
+    loop for each batch of tasks
+        loop for each task in the batch
+            AC->>TS: Request Task (4)
+            TS->>AC: Send Task & Model API (5)
+            AC->>AG: Run Agent with Task & Model API (6)
+            loop for each LLM call
+                AC->>AG: Prompt (7)
+                AG->>AC: Response (8)
+            end
+            AG->>AC: Rewarded Trace (9)
+            AC->>TS: Send Rewarded Trace (10)
+        end
+        TS->>RL: Send Batch of Traces (11)
+        RL->>TS: Return Updated Model (12)
+    end
+```

docs/how-to/train-sql-agent.md

@@ -0,0 +1,184 @@
+# SQL Agent with Agent Lightning
+
+> This tutorial is tested with `verl==0.5.0` and `vllm==0.10.0`.
+
+This example demonstrates how to build and train a self-correcting SQL agent. It leverages [Agent Lightning]({{ config.repo_url }}) and the `verl` framework for Reinforcement Learning (RL) based training, and LangGraph to define the agent's complex, cyclical reasoning workflow. The goal is to fine-tune a Large Language Model (LLM) to accurately convert natural language questions into executable SQL queries.
+
+## SQL Agent Implementation
+
+The design of Agent-lightning **allows flexible integration with various agent frameworks**, including AutoGen, CrewAI, OpenAI Agent SDK, LangGraph, and more. It can also work without agent frameworks, allowing you to train an agent built from scratch with Python code. See [our example gallery]({{ config.repo_url }}/tree/{{ config.extra.source_commit }}/examples) for more details.
+
+The core of the agent is a state machine built with LangGraph, which allows for a robust and transparent workflow. The agent's logic, as visualized below, starts by writing a query, executes it, and then enters a refinement loop where it checks and rewrites the query until it is deemed correct or a turn limit is reached.
+
+```mermaid
+---
+config:
+  flowchart:
+    curve: linear
+---
+graph LR;
+        __start__([<p>__start__</p>]):::first
+        write_query(write_query)
+        execute_query(execute_query)
+        check_query(check_query)
+        rewrite_query(rewrite_query)
+        __end__([<p>__end__</p>]):::last
+        __start__ --> write_query;
+        check_query -.-> __end__;
+        check_query -.-> rewrite_query;
+        execute_query --> check_query;
+        rewrite_query --> execute_query;
+        write_query --> execute_query;
+        classDef default fill:#f2f2f2,line-height:1.2
+        classDef first fill-opacity:0
+        classDef last fill:#cccccc
+```
+
+This workflow is implemented in the `SQLAgent` class within `sql_agent.py`. It consists of the following key steps:
+
+1. **write_query**: Given a user's question and database schema, the agent makes an initial attempt to write a SQL query.
+2. **execute_query**: The generated query is run against the target database.
+3. **check_query**: The agent analyzes the original query and its execution result (or error) to check for mistakes. It uses a specific prompt (`CHECK_QUERY_PROMPT`) to determine if the query is correct.
+4. **rewrite_query**: If the `check_query` step finds errors, the agent enters this step. It uses the feedback from the previous step to generate a corrected SQL query. The process then loops back to `check_query` for re-evaluation.
+5. **END**: The loop terminates when `check_query` confirms the query is correct or the maximum number of turns (`max_turns`) is exceeded. One turn corresponds to a complete cycle of `write_query` (if first round), `execute_query`, `check_query`, and potentially `rewrite_query`.
+
+We aim to train **write_query** and **rewrite_query** step in the setup of this example. The **check_query** step is not trained but will share the same LLM weights as the other steps.
+
+## Client-Server Training with Agent Lightning
+
+The training process uses a distributed client-server architecture designed by Agent Lightning to efficiently fine-tune the underlying LLM. This separation allows for scalable data generation across multiple clients while centralizing the computationally intensive model training on a dedicated server with GPUs, and also provides opportunities for customizing algorithms and training strategies (like [prompt optimization]({{ config.repo_url }}/tree/{{ config.extra.source_commit }}/examples/apo)) with minimal code changes.
+
+* **Training Server (`agentlightning.verl`)**: The server, launched with the first command below, manages the core training loop. It runs an RL algorithm (with `verl` of course) and hosts an OpenAI-compatible LLM endpoint (with `verl`'s async server). The server's sole purpose is to receive interaction data from clients and update the LLM's weights to improve its performance. [This link]({{ config.repo_url }}/tree/{{ config.extra.source_commit }}/agentlightning/verl) points to the implementation of the server, which is built upon `verl`.
+* **Agent Clients (`sql_agent.py`)**: The clients run the LangGraph agent logic described above. They connect to the server to fetch tasks (natural language questions) and use the server's **OpenAI-compatible endpoint** for all generation steps (`write_query`, `check_query`, `rewrite_query`). After completing a task, the client exports its interaction traces (traced by [AgentOps](https://www.agentops.ai/) and filtered by trace hierarchy), evaluates its correctness to calculate a reward, and sends the entire interaction history (the "trajectory") back to the server for training. To adapt any agent to an "agent client", you do not need to change the agent logic, but only need to invoke the client's `run` method with `agentlightning.trainer`.
+
+![Difference between the original agent and modified agent client](../assets/sql-agent-diff.png)
+
+## Running the Example
+
+1. Prepare the dataset: download from [here](https://drive.google.com/file/d/1oi9J1jZP9TyM35L85CL3qeGWl2jqlnL6/view) and unzip it to the `data` folder. It's basically a [Spider V1](https://yale-lily.github.io/spider) dataset converted to Parquet format. The dataset contains about 8000 training samples and about 2000 test samples, from which we sampled 500 samples for evaluation.
+   ```bash
+   pip install gdown
+   gdown --fuzzy https://drive.google.com/file/d/1oi9J1jZP9TyM35L85CL3qeGWl2jqlnL6/view
+   unzip -q spider-data.zip -d data
+   rm spider-data.zip
+   ```
+
+2. Install the required dependencies:
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+3. Launch the training server:
+   ```bash
+   python -m agentlightning.verl \
+       agentlightning.port=9997 \
+       algorithm.adv_estimator=grpo \
+       data.train_files=data/train_spider.parquet \
+       data.val_files=data/test_dev_500.parquet \
+       actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+       trainer.n_gpus_per_node=1 \
+       data.train_batch_size=32 \
+       actor_rollout_ref.rollout.n=4 \
+       actor_rollout_ref.actor.ppo_mini_batch_size=32 \
+       actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
+       actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
+       actor_rollout_ref.rollout.multi_turn.format=hermes \
+       actor_rollout_ref.model.path=meta-llama/Llama-3.2-3B-Instruct \
+       data.max_prompt_length=4096 \
+       data.max_response_length=2048 \
+       data.truncation='error' \
+       trainer.val_before_train=True \
+       actor_rollout_ref.actor.optim.lr=1e-6 \
+       actor_rollout_ref.model.use_remove_padding=True \
+       actor_rollout_ref.actor.use_kl_loss=False \
+       actor_rollout_ref.actor.kl_loss_coef=0.000 \
+       actor_rollout_ref.actor.entropy_coeff=0 \
+       actor_rollout_ref.actor.clip_ratio_low=0.2 \
+       actor_rollout_ref.actor.clip_ratio_high=0.3 \
+       actor_rollout_ref.model.enable_gradient_checkpointing=True \
+       actor_rollout_ref.actor.fsdp_config.param_offload=True \
+       actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
+       actor_rollout_ref.rollout.name=vllm \
+       actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+       actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \
+       actor_rollout_ref.ref.fsdp_config.param_offload=True \
+       algorithm.use_kl_in_reward=False \
+       trainer.critic_warmup=0 \
+       trainer.logger=['console','wandb'] \
+       trainer.project_name=AgentLightning \
+       trainer.experiment_name=train_sql_agent \
+       trainer.nnodes=1 \
+       trainer.save_freq=256 \
+       trainer.test_freq=32 \
+       trainer.total_epochs=2
+    ```
+
+4. Launch agent clients that connect with the server:
+   ```bash
+   export VERL_API_BASE=http://localhost:9997/  # Same as the server port. This is used for receiving tasks and sending results.
+   python sql_agent.py \
+       --litsqlagent.trained-agents write \  # Will only train the write and rewrite agent.
+       --trainer.n-workers 16 \
+       --litsqlagent.val-temperature 0
+   ```
+
+There is no hard requirement in the launching order of the server and clients. But remember to kill the long-running agent clients after the training is done.
+
+## Debug the Agent without verl
+
+You can run the agent client alone without the `verl` server. This is useful for debugging the agent logic and SQL execution.
+
+1. Copy `.env.example` to `.env` and fill in your OpenAI API key. `VERL_API_BASE` does not really matter here because you are not connecting to the server end.
+
+2. Run the agent client:
+   ```bash
+   dotenv run python sql_agent.py \
+       --litsqlagent.trained-agents write \  # Will only select the trajectories related to write and rewrite.
+       --trainer.n-workers 1 \  # For debug, use single process.
+       --trainer.dev true  # Enable the dev debug mode.
+   ```
+
+## Evaluation
+
+The example is evaluated using Llama-3.2-Instruct models. The models are trained on the Spider dataset for 2 epochs, with evaluation performed on a randomly selected subset of 500 test samples to compute held-out accuracy. The default setup for running agent clients during evaluation is as follows:
+
+```bash
+python sql_agent.py \
+   --litsqlagent.trained-agents write \
+   --trainer.n-workers 16 \
+   --trainer.daemon true \
+   --litsqlagent.val-temperature 0 \
+   --litsqlagent.max-turns 3 \
+   --litsqlagent.table-info-truncate 2048 \
+   --litsqlagent.execution-truncate 2048
+```
+
+The setup of training server is the same as the command above.
+
+### W&B Report
+
+[link](https://api.wandb.ai/links/ultmaster/4cid500g)
+
+### Performance Metrics
+
+![](../assets/sql-agent-val-reward-curve.png)
+
+| Model         | Size   |   Context |   Max Turns | Agents                        |   Acc (Initial) |   Acc (Final) | Transitions   |   Prompt Length | Response Length   |
+|---------------|--------|-----------|-------------|-------------------------------|-----------------|---------------|---------------|-----------------|-------------------|
+| Llama3.2      | 1B     |      2048 |           3 | write&#124;rewrite            |            21   |          49.6 | 2.87 → 3.08   |           821.2 | 319.2 → 249.4     |
+| Llama3.2      | 3B     |      2048 |           3 | write&#124;rewrite            |            51.8 |          66.4 | 2.20 → 2.72   |           865.6 | 116.2 → 314.3     |
+
+**Notes:**
+
+1. **Context Length**: Controlled via `--litsqlagent.table-info-truncate <context-length>` and `--litsqlagent.execution-truncate <context-length>`
+2. **Max Turns**: Set using `--litsqlagent.max-turns <max-turns>`
+3. **Agents**: Specified with `--litsqlagent.agents <regex>` (defaults to `write`, which matches both write and rewrite agents)
+4. **Transitions**: Represents the number of prompt-response pairs traced (collected) during each rollout. Note that this differs from the turn count in the SQL agent workflow, where one turn may encompass 2-3 transitions in the check-rewrite cycle. The number of transitions is also related to which *agents* get involved in the training.
+5. **Prompt/Response Length**: Average token count per **traced** prompt/transition response.
+
+### Efficiency Metrics
+
+| Model         | Size   |   Context |   Max Turns | Agents                        |   # GPUs |   # Steps |   Time (h) |   Time/Step (s) |   Rollout Time (%) |   Update Actor Time (%) |
+|---------------|--------|-----------|-------------|-------------------------------|----------|-----------|------------|-----------------|--------------------|-------------------------|
+| Llama3.2      | 1B     |      2048 |           3 | write&#124;rewrite            |        1 |       436 |      13.06 |            98.9 |               66.7 |                    25.2 |
+| Llama3.2      | 3B     |      2048 |           3 | write&#124;rewrite            |        2 |       436 |      10.3  |           181.3 |               63.9 |                    27.9 |

docs/index.md

@@ -0,0 +1,21 @@
+# Agent Lightning
+
+Agent Lightning is the absolute trainer to light up AI agents.
+
+## Features
+
+- Turn your agent into an optimizable beast with **ZERO CODE CHANGE** (almost)! 💤
+- Build with **ANY** agent framework (LangChain, OpenAI Agent SDK, AutoGen, CrewAI, ...); or even WITHOUT agent framework (Python OpenAI). You name it! 🤖
+- **Selectively** optimize one or more agents in a multi-agent system. 🎯
+- Embraces Reinforcement Learning, Automatic Prompt Optimization and more **algorithms**. 🤗
+
+## Quick Links
+
+- [Installation](quickstart/installation.md) - Get started with Agent Lightning
+- [Quickstart](quickstart/getting-started.md) - Learn the fundamentals of Agent Lightning
+- [Train SQL Agent with RL](how-to/train-sql-agent.md) - A practical example of training a SQL agent
+- [API Reference](reference/core.md) - Complete API documentation
+
+## License
+
+See the [LICENSE](https://github.com/microsoft/agent-lightning/blob/main/LICENSE) file for details.

docs/js/favicon-theme.js

@@ -0,0 +1,82 @@
+// Dynamic favicon switcher based on system theme preference
+(function() {
+    function setFavicon(isDark) {
+        // Remove existing favicon links
+        const existingFavicons = document.querySelectorAll('link[rel*="icon"]');
+        existingFavicons.forEach(link => link.remove());
+        
+        // Create new favicon link
+        const favicon = document.createElement('link');
+        favicon.rel = 'icon';
+        favicon.type = 'image/png';
+        
+        // Get the site root by finding how many levels deep we are
+        const pathSegments = window.location.pathname.split('/').filter(s => s);
+        const siteRoot = window.location.origin + '/' + pathSegments[0] + '/';
+        
+        // Choose favicon based on theme
+        if (isDark) {
+            favicon.href = siteRoot + 'assets/logo-dark.png';
+        } else {
+            favicon.href = siteRoot + 'assets/logo-light.png';
+        }
+        
+        // Add to document head
+        document.head.appendChild(favicon);
+    }
+    
+    function updateFavicon() {
+        // Check system preference
+        const prefersDark = window.matchMedia('(prefers-color-scheme: dark)').matches;
+        
+        // Check if user has manually selected a theme
+        const palette = document.querySelector('[data-md-color-scheme]');
+        const scheme = palette ? palette.getAttribute('data-md-color-scheme') : null;
+        
+        let isDark = false;
+        if (scheme === 'slate') {
+            isDark = true;
+        } else if (scheme === 'default') {
+            isDark = false;
+        } else {
+            // Fall back to system preference
+            isDark = prefersDark;
+        }
+        
+        setFavicon(isDark);
+    }
+    
+    // Initial favicon set
+    updateFavicon();
+    
+    // Listen for system theme changes
+    window.matchMedia('(prefers-color-scheme: dark)').addEventListener('change', updateFavicon);
+    
+    // Listen for manual theme changes in MkDocs Material
+    const observer = new MutationObserver(function(mutations) {
+        mutations.forEach(function(mutation) {
+            if (mutation.type === 'attributes' && 
+                (mutation.attributeName === 'data-md-color-scheme' || 
+                 mutation.attributeName === 'data-md-color-primary')) {
+                updateFavicon();
+            }
+        });
+    });
+    
+    // Observe the document body for theme changes
+    observer.observe(document.body, {
+        attributes: true,
+        attributeFilter: ['data-md-color-scheme', 'data-md-color-primary']
+    });
+    
+    // Also listen for palette toggle clicks
+    document.addEventListener('DOMContentLoaded', function() {
+        const toggles = document.querySelectorAll('[data-md-color-scheme]');
+        toggles.forEach(function(toggle) {
+            toggle.addEventListener('click', function() {
+                // Small delay to let MkDocs Material update the scheme
+                setTimeout(updateFavicon, 50);
+            });
+        });
+    });
+})();

docs/quickstart/getting-started.md

@@ -0,0 +1,233 @@
+# Getting Started
+
+This guide walks you through building your first Agent Lightning application - a simple prompt optimization system that finds the best system prompt for an AI agent.
+
+## What You'll Build
+
+You'll create a distributed training system with a server that manages optimization algorithms and tasks, a client with multiple workers that execute tasks in parallel, and built-in telemetry for monitoring and debugging.
+
+Before starting, ensure you have Python 3.10 or later, Agent Lightning installed (`pip install agentlightning`), and an OpenAI API key. The complete code is available in the [examples/apo]({{ config.repo_url }}/tree/{{ config.extra.source_commit }}/examples/apo) directory.
+
+## Part 1: Building Your Agent
+
+Let's start by creating a simple agent that can answer questions using OpenAI's API. Your agent needs to inherit from `LitAgent` and implement a `training_rollout` method.
+
+### Step 1: Create Your Agent Class
+
+First, import the necessary dependencies and create your agent class:
+
+```python
+from agentlightning.litagent import LitAgent
+
+class SimpleAgent(LitAgent):
+    def training_rollout(self, task, rollout_id, resources):
+        """Execute a single training rollout."""
+```
+
+The `training_rollout` method is the heart of your agent. It receives three parameters: a `task` dictionary containing the work to do (like "What is the capital of France?"), a unique `rollout_id` for tracking this execution, and `resources` from the server - in our case, the system prompt we're testing.
+
+### Step 2: Execute the Task
+
+Inside the training_rollout method, extract the system prompt from resources and use it to complete the task:
+
+```python
+# Extract the system prompt being tested
+system_prompt = resources["system_prompt"].template
+
+# Call OpenAI with this prompt
+result = openai.chat.completions.create(
+    model="gpt-4o-mini",
+    messages=[
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": task["prompt"]},
+    ],
+)
+```
+
+The server sends different prompts to test, and your agent uses each one to answer the same question. This lets us compare which prompt works best.
+
+### Step 3: Return a Reward Score
+
+After executing the task, return a reward score between 0 and 1:
+
+```python
+# In real scenarios, calculate based on response quality
+return random.uniform(0, 1)
+```
+
+Higher rewards mean better performance. In a real system, you'd calculate this with rules, or even an LLM as a judge. For now, we're using random values to demonstrate the flow.
+
+### Step 4: Set Up the Trainer
+
+To run your agent with multiple workers in parallel:
+
+```python
+from agentlightning.trainer import Trainer
+
+agent = SimpleAgent()
+trainer = Trainer(n_workers=2)  # Create 2 parallel workers
+trainer.fit(agent, backend="http://127.0.0.1:9997")
+```
+
+The trainer creates separate processes for each worker, allowing them to execute tasks independently. This parallelization significantly speeds up the optimization process - with 2 workers, you can test prompts twice as fast.
+
+## Part 2: Building the Optimization Server
+
+The server coordinates the training process and implements your optimization algorithm. It manages resources, distributes tasks, and collects results.
+
+### Step 1: Initialize the Server
+
+Create an async function to run your optimization:
+
+```python
+import asyncio
+from agentlightning.server import AgentLightningServer
+from agentlightning.types import PromptTemplate
+
+async def prompt_optimization():
+    server = AgentLightningServer(host="127.0.0.1", port=9997)
+    await server.start()
+```
+
+We use async/await because the server handles multiple clients simultaneously. This allows it to queue tasks without blocking and process results as they arrive from different workers.
+
+### Step 2: Test Different Prompts
+
+Define the prompts you want to test and iterate through them:
+
+```python
+prompt_candidates = [
+    "You are a helpful assistant.",
+    "You are a knowledgeable AI.",
+    "You are a friendly chatbot.",
+]
+
+for prompt in prompt_candidates:
+    # Send this prompt to all connected clients
+    resources = {
+        "system_prompt": PromptTemplate(template=prompt, engine="f-string")
+    }
+    await server.update_resources(resources)
+```
+
+When you update resources, all connected clients immediately receive the new system prompt. The format of resources can be arbitrary. We use the key `"system_prompt"` here as an example. The resources here are exactly you would expect at the client side, who will use this prompt for the next task they process.
+
+### Step 3: Queue Tasks and Collect Results
+
+For each prompt, queue a task and wait for results. The `{"prompt": ...}` format here is exact what you would expect from the client side code.
+
+```python
+# Queue a task for clients to process
+task_id = await server.queue_task(
+    sample={"prompt": "What is the capital of France?"}, 
+    mode="train"
+)
+
+# Wait for a client to complete it (30 second timeout)
+rollout = await server.poll_completed_rollout(task_id, timeout=30)
+
+# Extract and store the reward (this comes from the return value of the client side)
+reward = rollout.final_reward
+```
+
+The server queues the same question for each prompt. The rollout object contains not just the reward, but also detailed telemetry and trace information for debugging and optimization.
+
+### Step 4: Find the Best Prompt
+
+After testing all candidates, identify the winner:
+
+```python
+best_prompt = max(prompt_and_rewards, key=lambda x: x[1])
+print(f"Best prompt: '{best_prompt[0]}' with reward {best_prompt[1]:.3f}")
+```
+
+## Running Your System
+
+The [Complete example code]({{ config.repo_url }}/tree/{{ config.extra.source_commit }}/examples/apo) can be found on the GitHub repository. To run it:
+
+Create a `.env` file with your API credentials:
+
+```bash
+OPENAI_API_KEY=your-api-key-here
+OPENAI_API_BASE=https://api.openai.com/v1  # Optional
+```
+
+Start the server first in one terminal:
+```bash
+python server.py
+```
+
+Then start the client in another terminal:
+```bash
+python client.py
+```
+
+## Understanding the Output
+
+When you run the system, you'll see detailed logs from both the client and server. Understanding these logs helps you debug issues and optimize performance.
+
+### Client Output Explained
+
+```
+2025-08-10 12:59:38,224 [INFO] Initializing Trainer...
+```
+The trainer is starting up and preparing to create workers.
+
+```
+[INFO] Starting AgentOps local server on port 52081...
+```
+A local telemetry server starts to collect metrics and traces. You can access this at `http://localhost:52081` to see detailed execution traces.
+
+```
+[INFO] Starting worker process 0...
+[INFO] Starting worker process 1...
+```
+Two separate processes are created. Each can execute tasks independently, doubling your throughput.
+
+```
+[INFO] [Task 1 Received] ID: rollout-c1eb987b...
+Resources: {'system_prompt': PromptTemplate(...)}
+```
+A worker receives a task from the server along with the current prompt to test.
+
+```
+[INFO] [Worker 0 | Rollout] Completed in 1.09s. Reward: 0.631
+```
+Worker 0 finished executing the task in 1.09 seconds and calculated a reward of 0.631. This tells you both performance (execution time) and quality (reward score).
+
+### Server Output Explained
+
+```
+[Algo] Testing prompt: 'You are a helpful assistant.'
+```
+The optimization algorithm selects the next prompt to test.
+
+```
+[Algo] Task 'rollout-c1eb987b...' is now available for clients.
+```
+The task is queued and waiting for an available worker to pick it up.
+
+```
+[Algo] Received reward: 0.631
+```
+A client completed the task and returned a performance score. The server uses this to compare prompts.
+
+```
+[Algo] Best prompt: 'You are a knowledgeable AI.' (reward: 0.925)
+```
+After testing all prompts, the server identifies which one performed best.
+
+## What's Happening Behind the Scenes
+
+Agent Lightning handles several complex operations automatically. Multiple workers process tasks simultaneously. Every API call and execution is tracked through the telemetry tracing system, providing detailed traces for debugging and optimization. If a worker fails, others continue processing, ensuring your optimization doesn't stop.
+
+The AgentOps tracer, enabled by default, collects comprehensive data about each execution, including API calls, timing information, token usage, and error traces. The data is sent to the server and can be accessed via `rollout.triplets` and `rollout.traces` at the server side to build more advanced automatic optimization algorithms.
+
+## Next Steps
+
+Now that you have a working system, how about:
+
+- Replacing the random reward with actual quality metrics based on real response accuracy?
+- Testing the system prompt on a batch of different questions to see how it performs across various tasks?
+- Making the algorithm automatically improve the system prompt based on the best-performing ones?
+- Setting up a real agent system that consists of multiple prompts, and optimizing them together?

docs/quickstart/installation.md

@@ -0,0 +1,83 @@
+# Installation
+
+## Install from PyPI
+
+### Set Up Your Environment
+
+We strongly recommend creating a new virtual environment to avoid conflicts with other packages. You can use either `conda` or `venv`. **Python 3.10 or later** is recommended.
+
+### Install Core Training Dependencies (Optional)
+
+If you are running RL with Agent-Lightning, the next step is to install the essential packages: `PyTorch`, `FlashAttention`, `vLLM` and `VERL`. The following versions and installation order have been tested and are confirmed to work.
+
+```bash
+pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128
+pip install flash-attn --no-build-isolation
+pip install vllm==0.9.2
+pip install verl==0.5.0
+```
+
+See [this script]({{ config.repo_url }}/tree/{{ config.extra.source_commit }}/scripts/setup_stable_gpu.sh) for a full installation script.
+
+### Install Agent Lightning
+
+Now, you're ready to install Agent Lightning itself.
+
+```bash
+pip install agentlightning
+```
+
+### Install Agent Frameworks (Optional)
+
+If you plan to use other agent frameworks, you can install them with the following commands. If you don't need these, feel free to skip this step.
+We recommend doing this as the final step to avoid dependency versions being overwritten by mistake.
+
+```bash
+# AutoGen (Recommended to install first)
+pip install "autogen-agentchat" "autogen-ext[openai]"
+
+# LiteLLM
+pip install "litellm[proxy]"
+
+# MCP
+pip install mcp
+
+# UV
+pip install uv
+
+# OpenAI Agents
+pip install openai-agents
+
+# LangChain
+pip install langgraph "langchain[openai]" langchain-community langchain-text-splitters
+
+# SQL-related dependencies
+pip install sqlparse nltk
+```
+
+### Shortcuts for installing Extra Dependencies
+
+For development:
+```bash
+pip install agentlightning[dev]
+```
+
+For agent support:
+```bash
+pip install agentlightning[agent]
+```
+
+## Install from Source
+
+```
+git clone {{ config.repo_url }}
+cd agent-lightning
+pip install -e .[dev]
+```
+
+Please run pre-commit hooks before checking in code:
+
+```
+pre-commit install
+pre-commit run --all-files --show-diff-on-failure --color=always
+```

docs/reference/core.md

@@ -0,0 +1,51 @@
+# Agent Lightning Core
+
+## Client Side
+
+::: agentlightning.litagent
+    options:
+      show_source: true
+
+::: agentlightning.client
+    options:
+      show_source: true
+
+::: agentlightning.runner
+    options:
+      show_source: true
+
+::: agentlightning.trainer
+    options:
+      show_source: true
+
+::: agentlightning.tracer
+    options:
+      show_source: true
+
+::: agentlightning.reward
+    options:
+      show_source: true
+
+## Server Side
+
+::: agentlightning.server
+    options:
+      show_source: true
+
+## Utilities
+
+::: agentlightning.config
+    options:
+      show_source: true
+
+::: agentlightning.types
+    options:
+      show_source: true
+
+::: agentlightning.logging
+    options:
+      show_source: true
+
+::: agentlightning.instrumentation
+    options:
+      show_source: true

docs/reference/rl.md

@@ -0,0 +1,5 @@
+# Reinforcement Learning API
+
+::: agentlightning.verl
+    options:
+      show_source: true

examples/apo/apo_client.py

@@ -36,4 +36,4 @@ if __name__ == "__main__":
     dotenv.load_dotenv()
     agent = SimpleAgent()
     trainer = Trainer(n_workers=2)
-    trainer.fit(agent, endpoint="http://127.0.0.1:9997")
+    trainer.fit(agent, backend="http://127.0.0.1:9997")

mkdocs.yml

@@ -0,0 +1,90 @@
+site_name: Agent Lightning
+site_url: https://microsoft.github.io/agent-lightning/
+repo_url: https://github.com/microsoft/agent-lightning
+repo_name: agent-lightning
+
+theme:
+  name: material
+  logo: assets/logo-light.png
+  # Don't set favicon here - we'll handle it dynamically
+  palette:
+    - scheme: default
+      primary: red
+      accent: red
+      toggle:
+        icon: material/brightness-7
+        name: Switch to dark mode
+    - scheme: slate
+      primary: red
+      accent: red
+      toggle:
+        icon: material/brightness-4
+        name: Switch to light mode
+  features:
+    - navigation.sections
+    - navigation.expand
+    - navigation.top
+    - search.suggest
+    - search.highlight
+    - content.code.copy
+
+markdown_extensions:
+  - pymdownx.highlight:
+      anchor_linenums: true
+  - pymdownx.superfences:
+      custom_fences:
+        - name: mermaid
+          class: mermaid
+          format: !!python/name:pymdownx.superfences.fence_code_format
+  - pymdownx.snippets
+  - admonition
+  - pymdownx.details
+  - pymdownx.tabbed:
+      alternate_style: true
+  - toc:
+      permalink: true
+
+plugins:
+  - search
+  - git-revision-date-localized:
+      enable_creation_date: true
+      type: timeago
+  - git-authors
+  - macros
+  - mkdocstrings:
+      handlers:
+        python:
+          options:
+            show_source: true
+            show_root_heading: true
+            show_symbol_type_heading: true
+            show_symbol_type_toc: true
+            docstring_style: google
+  - mike:
+      version_selector: true
+      css_dir: css
+      javascript_dir: js
+      canonical_version: latest
+
+extra:
+  version:
+    provider: mike
+    default: latest
+  source_commit: !ENV [SOURCE_COMMIT, 'main']
+
+extra_javascript:
+  - https://unpkg.com/mermaid@10.6.1/dist/mermaid.min.js
+  - js/favicon-theme.js
+
+nav:
+  - Home: index.md
+  - Quickstart:
+    - Installation: quickstart/installation.md
+    - Getting Started: quickstart/getting-started.md
+  - How-To Guides:
+    - Train SQL Agent: how-to/train-sql-agent.md
+  - Deep Dive:
+    - Server-Client Architecture: deep-dive/server-client-architecture.md
+  - API Reference:
+    - Core: reference/core.md
+    - RL: reference/rl.md

pyproject.toml

@@ -24,6 +24,13 @@ dev = [
   "pre-commit",
   "pytest-rerunfailures",
   "black",
+  "mkdocs",
+  "mkdocs-material",
+  "mkdocstrings[python]",
+  "mike",
+  "mkdocs-git-revision-date-localized-plugin",
+  "mkdocs-git-authors-plugin",
+  "mkdocs-macros-plugin",
 ]
 experiment = [
   "random-word",