chore: import upstream snapshot with attribution

examples/chartqa/README.md

@@ -0,0 +1,113 @@
+# ChartQA Example
+
+[![chartqa workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-chartqa.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-chartqa.yml)
+
+This example demonstrates training a visual reasoning agent on the ChartQA dataset using Agent-Lightning with the VERL algorithm and LangGraph framework. The agent answers questions about charts through a multi-step workflow with self-refinement.
+
+## Requirements
+
+This example requires a single node with at least one 40GB GPU. Install dependencies with:
+
+```bash
+uv sync --frozen \
+    --group dev \
+    --group experiment \
+    --group image \
+    --group langchain \
+    --group vllm-0-10-2 \
+    --group torch-gpu-stable
+```
+
+**Currently vLLM 0.10.2 is the only tested version. You might see issues like `cu_seqlens_q must be on CUDA` or flash-attn installation failures if you use other versions.** (See https://github.com/vllm-project/vllm/issues/27340)
+
+## Dataset
+
+Download the ChartQA dataset and prepare it for training:
+
+```bash
+cd examples/chartqa
+python prepare_data.py
+```
+
+This downloads the ChartQA dataset from HuggingFace (`HuggingFaceM4/ChartQA`), saves images locally, and creates parquet files for training/testing. No HuggingFace token is required (public dataset).
+
+**Dataset Statistics:**
+
+- Training: ~18,000 chart question-answer pairs
+- Test: ~2,500 pairs
+- Chart types: Bar, line, pie, scatter, etc.
+
+## Included Files
+
+| File/Directory | Description |
+|----------------|-------------|
+| `chartqa_agent.py` | Chart reasoning agent using LangGraph with multi-step workflow (observe → extract → calculate → check → refine) |
+| `train_chartqa_agent.py` | Training script using VERL algorithm with configurable hyperparameters (debug, qwen) |
+| `debug_chartqa_agent.py` | Debugging script to test the agent with cloud APIs or a local vLLM proxy |
+| `prepare_data.py` | Script to download ChartQA dataset from HuggingFace and prepare parquet files |
+| `prompts.py` | Prompt templates for the agent workflow |
+| `multimodal_utils.py` | Utility functions for encoding images to base64 |
+| `env_var.py` | Environment variables and configurations |
+| `data/` | Directory containing images and parquet files after download |
+
+## Running Examples
+
+### Debugging with Cloud API (Default)
+
+For quick testing with OpenAI or other cloud APIs (no local GPU required):
+
+```bash
+export OPENAI_API_KEY=<your-api-key>
+python debug_chartqa_agent.py
+```
+
+For other providers (Azure, etc.), set `OPENAI_API_BASE`:
+
+```bash
+export OPENAI_API_BASE=https://your-resource.openai.azure.com/v1
+export OPENAI_MODEL=gpt-4o
+python debug_chartqa_agent.py
+```
+
+### Debugging with Local Model (LLMProxy)
+
+To test the agent with a local vLLM server and LLMProxy:
+
+```bash
+# Start a vLLM server (specify image path for VLM)
+export CHARTQA_DATA_DIR=<path to chartqa data>
+vllm serve Qwen/Qwen2-VL-2B-Instruct \
+    --gpu-memory-utilization 0.6 \
+    --max-model-len 4096 \
+    --allowed-local-media-path $CHARTQA_DATA_DIR \
+    --enable-prefix-caching \
+    --port 8088
+
+# Run the agent with LLMProxy
+USE_LLM_PROXY=1 \
+    OPENAI_API_BASE=http://localhost:8088/v1 \
+    OPENAI_MODEL=Qwen/Qwen2-VL-2B-Instruct \
+    python debug_chartqa_agent.py
+```
+
+### Training with Local Model
+
+```bash
+python train_chartqa_agent.py debug --n-runners 2
+```
+
+You can also use an external store server (recommended for distributed setups), first start the store:
+
+```bash
+agl store --port 4747
+```
+
+Then run the training script with the external store address:
+
+```bash
+AGL_MANAGED_STORE=0 python train_chartqa_agent.py qwen --external-store-address http://localhost:4747
+```
+
+If you want to track experiments with Weights & Biases, set the `WANDB_API_KEY` environment variable before training.
+
+The script automatically launches agent workers and the training server. The agent workers execute chart reasoning rollouts using the vision-language model, while the training server applies the VERL algorithm (GRPO) to improve the model based on answer accuracy rewards.

examples/chartqa/chartqa_agent.py

@@ -0,0 +1,410 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""ChartQA agent demonstrating LangGraph-based visual reasoning with refinement.
+
+This module defines `ChartQAAgent` plus the supporting prompt utilities used by
+`debug_chartqa_agent.py` and `train_chartqa_agent.py`.
+
+1. `analyze_chart` observes and summarizes the chart.
+2. `extract_data` calls a text-only LLM to extract the requested values.
+3. `calculate_answer` runs calculations grounded in prior steps.
+4. `check_answer` verifies reasoning quality.
+5. `refine_answer` conditionally patches mistakes before responding.
+
+Example usage can be found in `debug_chartqa_agent.py` and `train_chartqa_agent.py`.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+from typing import Any, Dict, Literal, cast
+
+import env_var as chartqa_env_var
+import termcolor
+from langchain.chat_models import BaseChatModel, init_chat_model
+from langchain_core.messages import AnyMessage, BaseMessage, HumanMessage
+from langgraph.graph import END, START, MessagesState, StateGraph
+from langgraph.graph.state import CompiledStateGraph
+from multimodal_utils import encode_image_to_base64
+from prompts import (
+    ANALYZE_CHART_PROMPT,
+    CALCULATE_ANSWER_PROMPT,
+    CHECK_ANSWER_PROMPT,
+    EXTRACT_DATA_PROMPT,
+    REFINE_ANSWER_PROMPT,
+)
+
+import agentlightning as agl
+
+logger = logging.getLogger("chartqa_agent")
+
+
+class ChartState(MessagesState):
+    question: str
+    image_path: str
+    observation: str
+    extracted_data: str
+    calculation: str
+    answer: str
+    feedback: str
+    num_turns: int
+    messages: list[AnyMessage]
+
+
+class ChartQAAgent(agl.LitAgent[Dict[str, Any]]):
+    """LangGraph-powered ChartQA agent with multi-step reasoning and refinement.
+
+    The implementation shares the same [`agl.LitAgent`][agentlightning.LitAgent] interface as
+    the Calc-X sample agent but augments it with image handling and LangGraph state tracking.
+    """
+
+    def __init__(
+        self,
+        model_name: str | None = None,
+        max_turns: int = 3,
+        debug: bool = False,
+        endpoint: str | None = None,
+        temperature: float = 0.0,
+        use_base64_images: bool = False,
+    ):
+        self.debug = debug
+        self.max_turns = max_turns
+        self.use_base64_images = use_base64_images
+        self.model_name = model_name
+        self.endpoint = endpoint
+        self.temperature = temperature
+
+        self._llm: BaseChatModel | None = None
+        self._graph: CompiledStateGraph[ChartState] | None = None
+
+    def _create_llm(self) -> BaseChatModel:
+        if self.model_name is None:
+            raise ValueError("model_name is required for creating LLM")
+        return init_chat_model(
+            self.model_name,
+            model_provider="openai",
+            openai_api_base=self.endpoint,
+            openai_api_key=chartqa_env_var.OPENAI_API_KEY,
+            temperature=self.temperature,
+            max_retries=2,
+            max_tokens=1024,
+            timeout=300,
+        )
+
+    def update_llm_config(self, model_name: str, endpoint: str | None, temperature: float | None) -> None:
+        """Update the LLM configuration. Re-create the LLM if the configuration is changed."""
+        updated: bool = False
+        if model_name != self.model_name:
+            self.model_name = model_name
+            updated = True
+        if endpoint != self.endpoint:
+            self.endpoint = endpoint
+            updated = True
+        if temperature != self.temperature:
+            self.temperature = temperature
+            updated = True
+        if updated:
+            self._llm = self._create_llm()
+
+    def _ensure_llm(self) -> BaseChatModel:
+        """Ensure the LLM is created and cached."""
+        if self._llm is None:
+            self._llm = self._create_llm()
+        return self._llm
+
+    def invoke_prompt(self, prompt: Any) -> AnyMessage:
+        """Invoke LLM with prompt."""
+        if self.debug:
+            for message in prompt.messages:
+                termcolor.cprint(message.pretty_repr(), "blue")
+
+        try:
+            result = self._ensure_llm().invoke(prompt)
+        except Exception as e:
+            logger.error(f"Failed to invoke prompt: {e}")
+            result = self._ensure_llm().invoke([HumanMessage(content="Please provide a reasonable answer.")])
+
+        if self.debug:
+            termcolor.cprint(result.pretty_repr(), "green")
+
+        return result  # type: ignore
+
+    def invoke_prompt_with_image(self, prompt_text: str, image_path: str) -> str:
+        """Invoke vision-language model with image.
+
+        Handles both local vLLM (file:// URLs) and cloud APIs (base64 encoding).
+        Cloud APIs (OpenAI, Anthropic, Google, Azure, etc.) require base64 encoding.
+        """
+        # Determine image URL format based on endpoint
+        if self.use_base64_images:
+            # Cloud APIs require base64 encoding for local files
+            image_url = encode_image_to_base64(image_path)
+        else:
+            # Local vLLM supports file:// URLs
+            if not image_path.startswith("file://"):
+                image_path = f"file://{os.path.realpath(image_path)}"
+            image_url = image_path
+
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt_text},
+                    {"type": "image_url", "image_url": {"url": image_url}},
+                ],
+            }
+        ]
+
+        if self.debug:
+            termcolor.cprint(f"[VLM Call] {prompt_text[:100]}...", "blue")
+
+        try:
+            result = self._ensure_llm().invoke(messages)
+            response = result.content if hasattr(result, "content") else str(result)  # type: ignore
+        except Exception as e:
+            logger.error(f"Failed to invoke VLM: {e}")
+            response = "<observe>Unable to analyze chart</observe>"
+
+        if self.debug:
+            termcolor.cprint(f"[VLM Response] {response[:200]}...", "green")
+
+        return response  # type: ignore
+
+    def extract_content(self, text: str, tag: str) -> str:
+        """Extract content between XML-style tags."""
+        match = re.search(rf"<{tag}>(.*?)</{tag}>", text, re.DOTALL)
+        return match.group(1).strip() if match else ""
+
+    def analyze_chart(self, state: ChartState) -> ChartState:
+        """Step 1: Observe and describe the chart."""
+        prompt: Any = ANALYZE_CHART_PROMPT.invoke({"question": state["question"]})  # type: ignore
+        prompt_text = prompt.messages[1].content
+
+        result_text = self.invoke_prompt_with_image(prompt_text, state["image_path"])
+
+        observation = self.extract_content(result_text, "observe")
+        if not observation:
+            observation = result_text
+
+        return {  # type: ignore
+            **state,
+            "observation": observation,
+            "num_turns": 1,
+            "messages": [HumanMessage(content=result_text)],
+        }
+
+    def extract_data(self, state: ChartState) -> ChartState:
+        """Step 2: Extract specific data values."""
+        prompt: Any = EXTRACT_DATA_PROMPT.invoke(  # type: ignore
+            {
+                "observation": state["observation"],
+                "question": state["question"],
+            }
+        )
+        result = self.invoke_prompt(prompt)
+
+        extracted_data = self.extract_content(result.content, "extract")  # type: ignore
+        if not extracted_data:
+            extracted_data = result.content  # type: ignore
+
+        return {  # type: ignore
+            **state,
+            "extracted_data": extracted_data,  # type: ignore
+            "messages": [*state.get("messages", []), result],
+        }
+
+    def calculate_answer(self, state: ChartState) -> ChartState:
+        """Step 3: Calculate and provide answer."""
+        prompt: Any = CALCULATE_ANSWER_PROMPT.invoke(  # type: ignore
+            {
+                "extracted_data": state["extracted_data"],
+                "question": state["question"],
+            }
+        )
+        result = self.invoke_prompt(prompt)
+
+        calculation = self.extract_content(result.content, "calculate")  # type: ignore
+        answer = self.extract_content(result.content, "answer")  # type: ignore
+        if not answer:
+            answer = cast(str, result.content)  # type: ignore
+
+        return {  # type: ignore
+            **state,
+            "calculation": calculation,
+            "answer": answer,
+            "messages": [*state.get("messages", []), result],
+        }
+
+    def check_answer(self, state: ChartState) -> ChartState:
+        """Step 4: Verify answer quality."""
+        prompt: Any = CHECK_ANSWER_PROMPT.invoke(  # type: ignore
+            {
+                "observation": state["observation"],
+                "extracted_data": state["extracted_data"],
+                "question": state["question"],
+                "answer": state["answer"],
+                "calculation": state.get("calculation", "No calculation shown"),
+            }
+        )
+        result = self.invoke_prompt(prompt)
+
+        if self.debug:
+            termcolor.cprint(f"[Check] {result.content}", "yellow")  # type: ignore
+
+        return {  # type: ignore
+            **state,
+            "feedback": result.content,  # type: ignore
+            "messages": [*state.get("messages", []), *prompt.messages, result],
+        }
+
+    def refine_answer(self, state: ChartState) -> ChartState:
+        """Step 5: Refine answer based on feedback."""
+        prompt: Any = REFINE_ANSWER_PROMPT.invoke(  # type: ignore
+            {
+                "observation": state["observation"],
+                "extracted_data": state["extracted_data"],
+                "question": state["question"],
+                "answer": state["answer"],
+                "calculation": state.get("calculation", ""),
+                "feedback": state["feedback"],
+            }
+        )
+        result = self.invoke_prompt(prompt)
+        content: str = result.content  # type: ignore
+
+        new_extracted = self.extract_content(content, "extract")
+        extracted_data = new_extracted if new_extracted else state["extracted_data"]
+
+        new_calculation = self.extract_content(content, "calculate")
+
+        new_answer = self.extract_content(content, "answer")
+        if not new_answer:
+            new_answer = content
+
+        return {  # type: ignore
+            **state,
+            "extracted_data": extracted_data,
+            "calculation": new_calculation,
+            "answer": new_answer,
+            "num_turns": state.get("num_turns", 0) + 1,
+            "messages": [*prompt.messages, result],
+        }
+
+    def should_continue(self, state: ChartState) -> Literal[END, "refine_answer"]:  # type: ignore
+        """Determine if refinement is needed."""
+        if state["messages"] and isinstance(
+            state["messages"][-1], BaseMessage
+        ):  # pyright: ignore[reportUnnecessaryIsInstance]
+            last_message = state["messages"][-1]
+            if "THE ANSWER IS CORRECT" in last_message.content:  # type: ignore
+                if "THE ANSWER IS INCORRECT" in last_message.content:  # type: ignore
+                    correct_index = last_message.content.rfind("THE ANSWER IS CORRECT")  # type: ignore
+                    incorrect_index = last_message.content.rfind("THE ANSWER IS INCORRECT")  # type: ignore
+                    if correct_index > incorrect_index:
+                        return END
+                else:
+                    return END
+
+        if state.get("num_turns", 0) >= self.max_turns:
+            return END
+
+        return "refine_answer"
+
+    def graph(self) -> CompiledStateGraph[ChartState]:
+        """Build the workflow graph with refinement loop."""
+        # Check if the graph is already built
+        if self._graph is not None:
+            return self._graph
+
+        builder = StateGraph(ChartState)
+        builder.add_node(self.analyze_chart)  # type: ignore
+        builder.add_node(self.extract_data)  # type: ignore
+        builder.add_node(self.calculate_answer)  # type: ignore
+        builder.add_node(self.check_answer)  # type: ignore
+        builder.add_node(self.refine_answer)  # type: ignore
+
+        builder.add_edge(START, "analyze_chart")
+        builder.add_edge("analyze_chart", "extract_data")
+        builder.add_edge("extract_data", "calculate_answer")
+        builder.add_edge("calculate_answer", "check_answer")
+        builder.add_conditional_edges(
+            "check_answer",
+            self.should_continue,  # type: ignore
+        )
+        builder.add_edge("refine_answer", "extract_data")
+
+        self._graph = builder.compile()  # type: ignore
+        return self._graph
+
+    def rollout(self, task: Dict[str, Any], resources: agl.NamedResources, rollout: agl.Rollout) -> float | None:
+        """AgentLightning wrapper for ChartQA agent."""
+
+        question = task["question"]
+
+        rollout = cast(agl.AttemptedRollout, rollout)
+        llm = cast(agl.LLM, resources["main_llm"])
+
+        image_path = os.path.join(chartqa_env_var.CHARTQA_DATA_DIR, task["image_path"])
+        ground_truth = task["answer"]
+
+        if not os.path.exists(image_path):
+            logger.error(f"Image {image_path} does not exist. Skipping.")
+            return None
+
+        # The new rollout could have a different endpoint or temperature.
+        # Update the LLM if necessary.
+        self.update_llm_config(
+            model_name=llm.model,
+            endpoint=llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
+            temperature=llm.sampling_parameters.get("temperature", 0.0),
+        )
+
+        try:
+            handler = self.tracer.get_langchain_handler()
+            result = self.graph().invoke(  # type: ignore
+                {"question": question, "image_path": image_path},  # type: ignore
+                {"callbacks": [handler] if handler else [], "recursion_limit": 100},
+            )
+        except Exception as e:
+            error_msg = f"[Rollout {rollout.rollout_id}] Error during agent invocation: {e}"
+            logger.error(error_msg, exc_info=True)
+            # Return 0.0 as reward to indicate failure
+            return 0.0
+
+        predicted_answer = result["answer"]
+        reward = evaluate_answer(predicted_answer, ground_truth, raise_on_error=False)
+
+        return reward
+
+
+def evaluate_answer(predicted: str, ground_truth: str, raise_on_error: bool = False) -> float:
+    """Evaluate answer accuracy."""
+    try:
+        pred = predicted.lower().strip()
+        gt = ground_truth.lower().strip()
+
+        # Exact match
+        if pred == gt:
+            return 1.0
+
+        # Try numeric comparison
+        try:
+            pred_num = float(pred.replace(",", ""))
+            gt_num = float(gt.replace(",", ""))
+            if abs(pred_num - gt_num) / max(abs(gt_num), 1e-9) < 0.02:
+                return 1.0
+        except (ValueError, AttributeError):
+            pass
+
+        # Partial credit for substring match
+        if pred in gt or gt in pred:
+            return 0.5
+
+        return 0.0
+    except Exception as e:
+        if raise_on_error:
+            raise
+        logger.exception(f"Error evaluating answer: {e}")
+        return 0.0

examples/chartqa/debug_chartqa_agent.py

@@ -0,0 +1,129 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Debugging helpers for the ChartQA agent.
+
+Example usage for OpenAI API:
+
+```bash
+python debug_chartqa_agent.py
+```
+
+Example usage for self-hosted model.
+
+```
+vllm serve Qwen/Qwen2-VL-2B-Instruct \
+    --gpu-memory-utilization 0.6 \
+    --max-model-len 4096 \
+    --allowed-local-media-path $CHARTQA_DATA_DIR \
+    --enable-prefix-caching \
+    --port 8088
+USE_LLM_PROXY=1 OPENAI_API_BASE=http://localhost:8088/v1 OPENAI_MODEL=Qwen/Qwen2-VL-2B-Instruct python debug_chartqa_agent.py
+```
+
+Ensure `CHARTQA_DATA_DIR` points to a directory with the prepared parquet file by running `python prepare_data.py` beforehand.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Dict, List, cast
+
+import env_var as chartqa_env_var
+import pandas as pd
+from chartqa_agent import ChartQAAgent
+
+import agentlightning as agl
+
+logger = logging.getLogger("chartqa_agent")
+
+
+def create_llm_proxy_for_chartqa(vllm_endpoint: str, port: int = 8081) -> agl.LLMProxy:
+    """Create an LLMProxy configured for ChartQA with token ID capture.
+
+    Args:
+        vllm_endpoint: Base URL for the hosted vLLM server.
+        port: Local port where the proxy should listen.
+
+    Returns:
+        An [`LLMProxy`][agentlightning.LLMProxy] instance launched in a thread.
+    """
+    store = agl.LightningStoreThreaded(agl.InMemoryLightningStore())
+
+    llm_proxy = agl.LLMProxy(
+        port=port,
+        store=store,
+        model_list=[
+            {
+                "model_name": "Qwen/Qwen2-VL-2B-Instruct",
+                "litellm_params": {
+                    "model": "hosted_vllm/Qwen/Qwen2-VL-2B-Instruct",
+                    "api_base": vllm_endpoint,
+                },
+            }
+        ],
+        callbacks=["return_token_ids"],
+        launch_mode="thread",
+    )
+
+    return llm_proxy
+
+
+def debug_chartqa_agent(use_llm_proxy: bool = False) -> None:
+    """Debug the ChartQA agent against cloud APIs or a local vLLM proxy.
+
+    Args:
+        use_llm_proxy: When `True`, spin up an LLMProxy that points to a local vLLM endpoint.
+
+    Raises:
+        FileNotFoundError: If the prepared ChartQA parquet file is missing.
+    """
+    test_data_path = os.path.join(chartqa_env_var.CHARTQA_DATA_DIR, "test_chartqa.parquet")
+
+    if not os.path.exists(test_data_path):
+        raise FileNotFoundError(f"Test data file {test_data_path} does not exist. Please run prepare_data.py first.")
+
+    df = pd.read_parquet(test_data_path).head(10)  # type: ignore
+    test_data = cast(List[Dict[str, Any]], df.to_dict(orient="records"))  # type: ignore
+
+    model = chartqa_env_var.OPENAI_MODEL
+    endpoint = chartqa_env_var.OPENAI_API_BASE
+    logger.info(
+        "Debug data: %s samples, model: %s, endpoint: %s, llm_proxy=%s",
+        len(test_data),
+        model,
+        endpoint,
+        use_llm_proxy,
+    )
+
+    llm_endpoint = endpoint
+    trainer_kwargs: Dict[str, Any] = {}
+
+    if use_llm_proxy:
+        proxy_port = 8089
+        llm_proxy = create_llm_proxy_for_chartqa(endpoint, port=proxy_port)
+        trainer_kwargs["llm_proxy"] = llm_proxy
+        trainer_kwargs["n_workers"] = 2
+        llm_endpoint = f"http://localhost:{proxy_port}/v1"
+        agent = ChartQAAgent()
+    else:
+        trainer_kwargs["n_workers"] = 1
+        agent = ChartQAAgent(use_base64_images=True)
+
+    trainer = agl.Trainer(
+        initial_resources={
+            "main_llm": agl.LLM(
+                endpoint=llm_endpoint,
+                model=model,
+                sampling_parameters={"temperature": 0.0},
+            )
+        },
+        **trainer_kwargs,
+    )
+
+    trainer.dev(agent, test_data)
+
+
+if __name__ == "__main__":
+    agl.setup_logging(apply_to=["chartqa_agent"])
+    debug_chartqa_agent(use_llm_proxy=chartqa_env_var.USE_LLM_PROXY)

examples/chartqa/env_var.py

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import os
+
+__all__ = [
+    "CHARTQA_ROOT_DIR",
+    "CHARTQA_DATA_DIR",
+    "CHARTQA_IMAGES_DIR",
+    "USE_BASE64_IMAGES",
+    "USE_LLM_PROXY",
+    "OPENAI_API_BASE",
+    "OPENAI_API_KEY",
+    "OPENAI_MODEL",
+]
+
+CHARTQA_ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+
+CHARTQA_DATA_DIR = os.getenv("CHARTQA_DATA_DIR", os.path.realpath(os.path.join(CHARTQA_ROOT_DIR, "data")))
+
+CHARTQA_IMAGES_DIR = os.getenv("CHARTQA_IMAGES_DIR", os.path.realpath(os.path.join(CHARTQA_ROOT_DIR, "data", "images")))
+
+USE_BASE64_IMAGES = os.getenv("USE_BASE64_IMAGES", "false").lower() in ("1", "true", "yes")
+
+USE_LLM_PROXY = os.getenv("USE_LLM_PROXY", "false").lower() in ("1", "true", "yes")
+
+OPENAI_API_BASE = os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1")
+
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "token-abc123")
+
+OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4.1-mini")

examples/chartqa/multimodal_utils.py

@@ -0,0 +1,113 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Multimodal support utilities for Agent Lightning.
+
+This module provides helper functions for working with multimodal agents,
+particularly for vision-language tasks.
+"""
+
+from __future__ import annotations
+
+import base64
+from io import BytesIO
+from pathlib import Path
+from typing import Any, Union
+
+import requests
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+__all__ = [
+    "encode_image_to_base64",
+    "create_image_message",
+]
+
+
+def encode_image_to_base64(image: Union[str, Path, PILImage], max_size: int = 2048) -> str:
+    """
+    Encode an image to base64 string for multimodal LLM APIs.
+
+    Args:
+        image: Image source (file path, URL, or PIL Image object)
+        max_size: Maximum dimension for resizing
+
+    Returns:
+        Base64 encoded image string with data URI prefix
+
+    Raises:
+        ImportError: If PIL (Pillow) is not installed
+        TypeError: If image type is not supported
+
+    Examples:
+        >>> encoded = encode_image_to_base64("photo.jpg")
+        >>> encoded[:30]
+        'data:image/jpeg;base64,/9j/4A...'
+
+        >>> from PIL import Image
+        >>> img = Image.open("photo.jpg")
+        >>> encoded = encode_image_to_base64(img)
+    """
+    # Load image
+    if isinstance(image, (str, Path)):
+        image_str = str(image)
+        if image_str.startswith(("http://", "https://")):
+            response = requests.get(image_str, timeout=30)
+            response.raise_for_status()
+            img = Image.open(BytesIO(response.content))
+        else:
+            img = Image.open(image_str)
+    elif hasattr(image, "mode"):
+        # PIL Image object
+        img = image
+    else:
+        raise TypeError(f"Unsupported image type: {type(image)}")
+
+    # Convert to RGB
+    if img.mode == "RGBA":
+        background = Image.new("RGB", img.size, (255, 255, 255))
+        background.paste(img, mask=img.split()[3])
+        img = background
+    elif img.mode != "RGB":
+        img = img.convert("RGB")
+
+    # Resize if needed
+    if max(img.size) > max_size:
+        img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
+
+    # Encode
+    buffered = BytesIO()
+    img.save(buffered, format="JPEG", quality=85)
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+
+    return f"data:image/jpeg;base64,{img_str}"
+
+
+def create_image_message(text: str, image: Union[str, Path, PILImage], use_base64: bool = True) -> dict[str, Any]:
+    """
+    Create an OpenAI-compatible multimodal message.
+
+    Args:
+        text: The text prompt/question
+        image: Image source (path, URL, or PIL Image)
+        use_base64: If True, encode as base64; if False, use URL directly
+
+    Returns:
+        Message dict with role="user" and multimodal content
+
+    Examples:
+        >>> msg = create_image_message("What's in the image?", "photo.jpg")
+        >>> msg["role"]
+        'user'
+        >>> len(msg["content"])
+        2
+    """
+    content: list[dict[str, Any]] = [{"type": "text", "text": text}]
+
+    if isinstance(image, str) and image.startswith(("http://", "https://")) and not use_base64:
+        content.append({"type": "image_url", "image_url": {"url": image}})
+    else:
+        encoded = encode_image_to_base64(image)
+        content.append({"type": "image_url", "image_url": {"url": encoded}})
+
+    return {"role": "user", "content": content}

examples/chartqa/prepare_data.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Prepare ChartQA dataset from HuggingFace for training."""
+
+from pathlib import Path
+from typing import Any, Dict, List
+
+import pandas as pd
+from datasets import load_dataset  # pyright: ignore[reportUnknownVariableType]
+
+
+def prepare_chartqa():
+    """Download ChartQA and convert to parquet format."""
+    data_dir = Path("data")
+    images_dir = data_dir / "images"
+    images_dir.mkdir(parents=True, exist_ok=True)
+
+    dataset = load_dataset("HuggingFaceM4/ChartQA")
+
+    for split in ["train", "test"]:
+        tasks: List[Dict[str, Any]] = []
+        dataset_length = len(dataset[split])  # type: ignore
+        for idx, item in enumerate(dataset[split]):  # pyright: ignore[reportUnknownArgumentType]
+            if idx % 1000 == 0:
+                print(f"Processing {split} item {idx} (out of {dataset_length})")
+            image_filename = f"{split}_{idx:06d}.png"
+            image_path = images_dir / image_filename
+            if not image_path.exists():
+                item["image"].save(image_path)
+
+            tasks.append(
+                {
+                    "id": f"{split}_{idx}",
+                    "image_path": f"images/{image_filename}",
+                    "question": item["query"],
+                    "answer": str(item["label"]),
+                }
+            )
+
+        pd.DataFrame(tasks).to_parquet(data_dir / f"{split}_chartqa.parquet", index=False)  # type: ignore
+
+
+if __name__ == "__main__":
+    prepare_chartqa()

examples/chartqa/prompts.py

@@ -0,0 +1,198 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Prompts for ChartQA agent workflow."""
+
+from langchain_core.prompts import ChatPromptTemplate
+
+ANALYZE_CHART_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+You are a visual reasoning expert analyzing charts and graphs.
+Given a chart image and a question, first carefully observe and describe the chart.
+
+Instructions:
+- Identify the chart type (bar chart, line chart, pie chart, scatter plot, etc.)
+- Note the axes labels and units (if applicable)
+- Describe the data series or categories shown
+- Observe key patterns, trends, or noteworthy values
+- Pay attention to legends, titles, and annotations
+
+## Output Format ##
+
+Provide your observation inside <observe> and </observe> tags.
+
+Example:
+<observe>
+Bar chart showing GDP of 5 countries. X-axis shows country names, Y-axis shows GDP in trillions of USD.
+Data values: USA appears highest at around 25, China second at around 20, followed by India, UK, and France.
+</observe>
+""".strip(),
+        ),
+        ("user", "Question: {question}"),
+    ]
+)
+
+EXTRACT_DATA_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+Based on your observation of the chart, extract the specific data values needed to answer the question.
+
+Instructions:
+- Extract only the data relevant to the question
+- Be precise with values (read carefully from the chart)
+- Include labels/categories with each value
+- Use appropriate units
+
+## Output Format ##
+
+Provide extracted data inside <extract> and </extract> tags.
+Format: Label1: Value1, Label2: Value2, ...
+
+Example:
+<extract>
+USA: 25, China: 20, India: 15, UK: 10, France: 8
+</extract>
+""".strip(),
+        ),
+        (
+            "user",
+            """Observation: {observation}
+
+Question: {question}
+
+Please extract the relevant data values.""",
+        ),
+    ]
+)
+
+CALCULATE_ANSWER_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+Using the extracted data, perform any necessary calculations to answer the question.
+
+Instructions:
+- Show your calculation steps clearly
+- Use correct mathematical operations
+- Pay attention to the question (average, sum, difference, maximum, etc.)
+- Provide a precise numerical answer if applicable
+- Keep the answer concise (typically 1-10 words)
+
+## Output Format ##
+
+Show calculation inside <calculate> and </calculate> tags (if needed).
+Provide final answer inside <answer> and </answer> tags.
+
+Example:
+<calculate>
+Average = (25 + 20 + 15 + 10 + 8) / 5 = 78 / 5 = 15.6
+</calculate>
+<answer>
+15.6
+</answer>
+""".strip(),
+        ),
+        (
+            "user",
+            """Extracted Data: {extracted_data}
+
+Question: {question}
+
+Please calculate and provide the answer.""",
+        ),
+    ]
+)
+
+CHECK_ANSWER_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+You are a chart analysis expert with strong attention to detail.
+Review the answer for potential mistakes.
+
+Common mistakes to check:
+- Incorrect data extraction from chart (misread values)
+- Arithmetic errors in calculations
+- Misunderstanding the question type (average vs. sum vs. difference)
+- Wrong number of data points counted
+- Incorrect units or scale interpretation
+- Off-by-one errors
+
+## Chart Information ##
+
+Observation: {observation}
+Extracted Data: {extracted_data}
+
+## Output Format ##
+
+If any mistakes are found, list each error clearly.
+After listing mistakes (if any), conclude with **ONE** of the following exact phrases in all caps:
+- If mistakes are found: `THE ANSWER IS INCORRECT.`
+- If no mistakes are found: `THE ANSWER IS CORRECT.`
+
+DO NOT write the corrected answer in this response. You only need to report mistakes.
+""".strip(),
+        ),
+        (
+            "user",
+            """Question: {question}
+
+Current Answer: {answer}
+
+Calculation shown:
+{calculation}
+
+Please review this answer for correctness.""",
+        ),
+    ]
+)
+
+REFINE_ANSWER_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+You are a chart analysis agent.
+The previous answer had errors. Based on the feedback, provide a corrected answer.
+
+Instructions:
+- Re-examine the chart observation carefully
+- Correct any data extraction errors by re-extracting if needed
+- Fix calculation mistakes
+- Address all points mentioned in the feedback
+
+## Chart Observation ##
+
+{observation}
+
+## Output Format ##
+
+If you need to re-extract data, provide it inside <extract> and </extract> tags.
+Show corrected calculation inside <calculate> and </calculate> tags.
+Provide corrected answer inside <answer> and </answer> tags.
+""".strip(),
+        ),
+        (
+            "user",
+            """Question: {question}
+
+## Previous Attempt ##
+
+Extracted Data: {extracted_data}
+Calculation: {calculation}
+Answer: {answer}
+
+## Feedback ##
+
+{feedback}
+
+Please provide the corrected answer.""",
+        ),
+    ]
+)

examples/chartqa/train_chartqa_agent.py

@@ -0,0 +1,215 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Training helper for ChartQA modeled VERL workflow.
+
+Example usage:
+
+```bash
+python train_chartqa_agent.py debug --n-runners 2
+```
+
+or:
+
+```bash
+AGL_MANAGED_STORE=0 python train_chartqa_agent.py qwen --external-store-address http://localhost:9999
+```
+
+Make sure to run `python prepare_data.py` so the parquet files referenced here exist.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import uuid
+from copy import deepcopy
+from datetime import datetime
+from typing import Any, Dict, Optional, cast
+
+import env_var as chartqa_env_var
+import pandas as pd
+from chartqa_agent import ChartQAAgent
+
+import agentlightning as agl
+from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
+
+RL_CONFIG: Dict[str, Any] = {
+    "algorithm": {"adv_estimator": "grpo", "use_kl_in_reward": False},
+    "data": {
+        "image_base_dir": chartqa_env_var.CHARTQA_IMAGES_DIR,
+        "train_batch_size": 32,
+        "max_prompt_length": 4096,
+        "max_response_length": 1024,
+        "truncation": "error",
+    },
+    "actor_rollout_ref": {
+        "rollout": {
+            "tensor_model_parallel_size": 1,
+            "n": 4,
+            "log_prob_micro_batch_size_per_gpu": 1,
+            "name": "vllm",
+            "gpu_memory_utilization": 0.8,
+            "enable_prefix_caching": True,
+            "engine_kwargs": {"vllm": {"allowed_local_media_path": chartqa_env_var.CHARTQA_IMAGES_DIR}},
+        },
+        "actor": {
+            "ppo_mini_batch_size": 32,
+            "ppo_micro_batch_size_per_gpu": 4,
+            "optim": {"lr": 1e-6},
+            "use_kl_loss": False,
+            "kl_loss_coef": 0.0,
+            "entropy_coeff": 0,
+            "clip_ratio_low": 0.2,
+            "clip_ratio_high": 0.3,
+            "fsdp_config": {"param_offload": True, "optimizer_offload": True},
+        },
+        "ref": {"log_prob_micro_batch_size_per_gpu": 1, "fsdp_config": {"param_offload": True}},
+        "model": {
+            "path": "Qwen/Qwen2-VL-2B-Instruct",
+            "use_remove_padding": True,
+            "enable_gradient_checkpointing": True,
+        },
+    },
+    "trainer": {
+        "n_gpus_per_node": 1,
+        "val_before_train": False,
+        "critic_warmup": 0,
+        "logger": ["console", "wandb"],
+        "project_name": "AgentLightning",
+        "experiment_name": "chartqa",
+        "nnodes": 1,
+    },
+}
+
+
+def config_ci() -> Dict[str, Any]:
+    """Return a CI-friendly RL config for ChartQA."""
+    # For CI testing, we need to set the experiment name and project name so that
+    # they are available to subsequent steps.
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    random_suffix = uuid.uuid4().hex[:8]
+    EXPERIMENT_NAME = f"chartqa_ci_{timestamp}_{random_suffix}"
+    PROJECT_NAME = "AgentLightningCI"
+    github_output = os.getenv("GITHUB_OUTPUT")
+    if github_output:
+        with open(github_output, "a") as f:
+            f.write(f"project_name={PROJECT_NAME}\n")
+            f.write(f"run_name={EXPERIMENT_NAME}\n")
+
+    config = deepcopy(RL_CONFIG)
+    config["data"]["train_batch_size"] = 16
+    config["trainer"]["n_gpus_per_node"] = 1
+    config["trainer"]["total_training_steps"] = 4
+    config["trainer"]["val_before_train"] = True
+    config["trainer"]["test_freq"] = 2
+    config["trainer"]["experiment_name"] = EXPERIMENT_NAME
+    config["trainer"]["project_name"] = PROJECT_NAME
+    return config
+
+
+def config_debug() -> Dict[str, Any]:
+    """Return a short debugging config for smoke testing ChartQA training."""
+    config = deepcopy(RL_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.5
+    config["trainer"]["total_training_steps"] = 10
+    config["trainer"]["test_freq"] = 2
+    return config
+
+
+def config_qwen() -> Dict[str, Any]:
+    """Return a Qwen-focused config with validation before each epoch."""
+    config = deepcopy(RL_CONFIG)
+    config["trainer"]["val_before_train"] = True
+    config["trainer"]["n_gpus_per_node"] = 2
+    config["trainer"]["total_epochs"] = 2
+    config["trainer"]["test_freq"] = 32
+    return config
+
+
+def train(
+    config: Dict[str, Any],
+    train_data: agl.Dataset[Any],
+    val_data: agl.Dataset[Any],
+    external_store_address: str,
+    n_runners: int,
+    debug: bool,
+) -> None:
+    """Run VERL training for ChartQA.
+
+    Args:
+        config: VERL configuration produced by one of the helpers above.
+        train_data: Training dataset of ChartQA samples.
+        val_data: Validation dataset for periodic evaluation.
+        external_store_address: Optional address of an existing LightningStore to reuse.
+        n_runners: Number of runners passed to [`Trainer.fit`][agentlightning.Trainer.fit].
+        debug: Enables verbose logging tied to `--debug`.
+    """
+    agl.setup_logging(level="DEBUG" if debug else "INFO", apply_to=["agentlightning", __name__])
+    agent = ChartQAAgent()
+    algorithm = agl.VERL(config)
+
+    if external_store_address:
+        store: Optional[agl.LightningStore] = agl.LightningStoreClient(external_store_address)
+    else:
+        store = None
+
+    trainer = agl.Trainer(
+        n_runners=n_runners,
+        algorithm=algorithm,
+        store=store,
+    )
+
+    trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)  # type: ignore
+
+
+def main():
+    """Parse CLI arguments and kick off ChartQA training."""
+    agl.setup_logging(apply_to=["chartqa_agent"])
+    parser = argparse.ArgumentParser(description="Train ChartQA agent")
+    parser.add_argument("config", choices=["debug", "qwen", "ci"], help="Training configuration")
+    parser.add_argument("--n-runners", type=int, default=10, help="Number of runners for Trainer")
+    parser.add_argument(
+        "--external-store-address",
+        type=str,
+        default=None,
+        help="Connect to an external store instead of creating a new one in memory (e.g., http://localhost:4747)",
+    )
+    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
+    args = parser.parse_args()
+
+    if args.external_store_address:
+        print(f"Connecting to external store at: {args.external_store_address}")
+        if resolve_bool_env_var(LightningEnvVar.AGL_MANAGED_STORE, fallback=True):
+            raise ValueError(
+                "When using an external store, please set the environment variable AGL_MANAGED_STORE=0. "
+                "Otherwise the trainer will still try to manage the store lifecycle for you!"
+            )
+
+    CONFIGS = {
+        "debug": config_debug,
+        "qwen": config_qwen,
+        "ci": config_ci,
+    }
+
+    train_data_path = os.path.join(chartqa_env_var.CHARTQA_DATA_DIR, "train_chartqa.parquet")
+    val_data_path = os.path.join(chartqa_env_var.CHARTQA_DATA_DIR, "test_chartqa.parquet")
+
+    train_data = pd.read_parquet(train_data_path).to_dict(orient="records")  # type: ignore
+
+    if args.config in ["debug", "ci"]:
+        val_data = pd.read_parquet(val_data_path).sample(n=100, random_state=42).to_dict(orient="records")  # type: ignore
+    else:
+        val_data = pd.read_parquet(val_data_path).to_dict(orient="records")  # type: ignore
+
+    train(
+        config=CONFIGS[args.config](),
+        train_data=cast(agl.Dataset[Any], train_data),
+        val_data=cast(agl.Dataset[Any], val_data),
+        external_store_address=args.external_store_address,
+        n_runners=args.n_runners,
+        debug=args.debug,
+    )
+
+
+if __name__ == "__main__":
+    main()