update

.github/workflows/examples.yml

@@ -129,6 +129,17 @@ jobs:
           OPENAI_BASE_URL: http://localhost:12306/
           OPENAI_API_KEY: dummy
 
+      - name: APO built-in algorithm
+        run: |
+          set -ex
+          . .venv/bin/activate
+          cd examples/apo
+          python room_selector_apo.py
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+        if: success() || failure()
+
       - name: Spider sanity check
         run: |
           set -ex
@@ -139,6 +150,7 @@ jobs:
           VERL_API_BASE: http://localhost:9999/
           OPENAI_API_BASE: http://localhost:12306/
           OPENAI_API_KEY: dummy
+        if: success() || failure()
       - name: Calc-X MCP sanity check
         run: |
           set -ex

a.py

@@ -0,0 +1,22 @@
+import asyncio
+
+
+async def a():
+    print("a")
+    b()
+    print("finish")
+
+
+def b():
+    print("b")
+    loop = asyncio.get_running_loop()
+    fut = asyncio.run_coroutine_threadsafe(c(), loop)
+    fut.result(timeout=5.0)
+
+
+async def c():
+    print("c")
+    await asyncio.sleep(0.1)
+
+
+asyncio.run(a())

agentlightning/__init__.py

@@ -2,22 +2,18 @@
 
 __version__ = "0.2.0"
 
-from .client import AgentLightningClient, DevTaskLoader
-from .config import lightning_cli
+from .adapter import *
+from .algorithm import *
+from .client import AgentLightningClient, DevTaskLoader  # deprecated  # type: ignore
+from .config import *
+from .emitter import *
+from .execution import *
 from .litagent import *
-from .logging import configure_logger
-from .reward import reward
-from .server import AgentLightningServer
-from .trainer import Trainer
+from .llm_proxy import *
+from .logging import *
+from .runner import *
+from .server import AgentLightningServer  # deprecated  # type: ignore
+from .store import *
+from .tracer import *
+from .trainer import *
 from .types import *
-
-__all__ = [
-    "AgentLightningClient",
-    "DevTaskLoader",
-    "lightning_cli",
-    "configure_logger",
-    "reward",
-    "AgentLightningServer",
-    "Trainer",
-    "__version__",
-]

agentlightning/adapter/__init__.py

@@ -1,6 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 from .base import Adapter, TraceAdapter
-from .triplet import BaseTraceTripletAdapter, LlmProxyTripletAdapter, TraceTripletAdapter
+from .messages import TraceToMessages
+from .triplet import LlmProxyTraceToTriplet, TracerTraceToTriplet, TraceToTripletBase
 
-__all__ = ["TraceAdapter", "Adapter", "BaseTraceTripletAdapter", "TraceTripletAdapter", "LlmProxyTripletAdapter"]
+__all__ = [
+    "TraceAdapter",
+    "Adapter",
+    "TraceToTripletBase",
+    "TracerTraceToTriplet",
+    "LlmProxyTraceToTriplet",
+    "TraceToMessages",
+]

agentlightning/adapter/base.py

@@ -14,14 +14,17 @@ class Adapter(Generic[T_from, T_to]):
     """Base class for synchronous adapters that convert data from one format to another.
 
     This class defines a simple protocol for transformation:
+
     - The `__call__` method makes adapters callable, so they can be used like functions.
     - Subclasses must implement the `adapt` method to define the actual conversion logic.
 
     Type parameters:
-        T_from: The source data type (input).
-        T_to: The target data type (output).
+
+    - T_from: The source data type (input).
+    - T_to: The target data type (output).
 
     Example:
+
         >>> class IntToStrAdapter(Adapter[int, str]):
         ...     def adapt(self, source: int) -> str:
         ...         return str(source)
@@ -56,10 +59,6 @@ class Adapter(Generic[T_from, T_to]):
 
         Returns:
             Data converted to the target format.
-
-        Raises:
-            NotImplementedError: If the method is not implemented
-            in a subclass.
         """
         raise NotImplementedError("Adapter.adapt() is not implemented")
 

agentlightning/adapter/messages.py

@@ -2,23 +2,24 @@
 
 import json
 from collections import defaultdict
-from typing import Any, Dict, Generator, List, Optional, Sequence, TypedDict, Union, cast
+from typing import Any, Dict, Generator, Iterable, List, Optional, TypedDict, Union, cast
 
-from openai.types.chat.chat_completion_function_tool_param import ChatCompletionFunctionToolParam
-from openai.types.chat.chat_completion_message import ChatCompletionMessage
-from openai.types.chat.chat_completion_message_function_tool_call import ChatCompletionMessageFunctionToolCall, Function
-from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
-from openai.types.shared_params import FunctionDefinition
-from pydantic import BaseModel, TypeAdapter
+from openai.types.chat import (
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionFunctionToolParam,
+    ChatCompletionMessageFunctionToolCallParam,
+    ChatCompletionMessageParam,
+)
+from pydantic import TypeAdapter
 
 from agentlightning.types import Span
 
 from .base import TraceAdapter
 
 
-class OpenAIMessages(BaseModel):
-    messages: List[Union[ChatCompletionMessage, ChatCompletionMessageParam]]
-    tools: Optional[List[ChatCompletionFunctionToolParam]] = None
+class OpenAIMessages(TypedDict):
+    messages: List[ChatCompletionMessageParam]
+    tools: Optional[List[ChatCompletionFunctionToolParam]]
 
 
 class _RawSpanInfo(TypedDict):
@@ -26,6 +27,7 @@ class _RawSpanInfo(TypedDict):
     completion: List[Dict[str, Any]]
     request: Dict[str, Any]
     response: Dict[str, Any]
+    tools: List[Dict[str, Any]]
 
 
 def group_genai_dict(data: Dict[str, Any], prefix: str) -> Union[Dict[str, Any], List[Any]]:
@@ -77,22 +79,15 @@ def group_genai_dict(data: Dict[str, Any], prefix: str) -> Union[Dict[str, Any],
     return result
 
 
-def convert_to_openai_messages(
-    prompt_completion_list: List[_RawSpanInfo], tool_requests: List[Dict[str, Any]]
-) -> Generator[OpenAIMessages, None, None]:
+def convert_to_openai_messages(prompt_completion_list: List[_RawSpanInfo]) -> Generator[OpenAIMessages, None, None]:
     """
     Convert raw tool call traces + prompt/completion list
     into OpenAI fine-tuning JSONL format (tool calling style).
 
-    Since promopt-completions sometimes do not contain the generated tool calls,
-    the tool call requests need to be provided separately.
-    The tool calls are then matched in a first-come-first-served basis to the tool call requests.
-
     https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/fine-tuning-functions
     """
     for pc_entry in prompt_completion_list:
-        messages: List[Union[ChatCompletionMessage, ChatCompletionMessageParam]] = []
-        tools: List[ChatCompletionFunctionToolParam] = []
+        messages: List[ChatCompletionMessageParam] = []
 
         # Extract messages
         for msg in pc_entry["prompt"]:
@@ -100,17 +95,18 @@ def convert_to_openai_messages(
 
             if role == "assistant" and "tool_calls" in msg:
                 # Use the tool_calls directly
-                tool_calls: Sequence[ChatCompletionMessageFunctionToolCall] = []
-                for call in msg["tool_calls"]:
-                    function = Function(name=call["name"], arguments=call["arguments"])
-                    tool_calls.append(
-                        ChatCompletionMessageFunctionToolCall(
-                            id=call["id"],
-                            type="function",
-                            function=function,
-                        )
+                # This branch is usually not used in the wild.
+                tool_calls: List[ChatCompletionMessageFunctionToolCallParam] = [
+                    ChatCompletionMessageFunctionToolCallParam(
+                        id=call["id"],
+                        type="function",
+                        function={"name": call["name"], "arguments": call["arguments"]},
                     )
-                messages.append(ChatCompletionMessage(role="assistant", tool_calls=list(tool_calls)))
+                    for call in msg["tool_calls"]
+                ]
+                messages.append(
+                    ChatCompletionAssistantMessageParam(role="assistant", content=None, tool_calls=tool_calls)
+                )
             else:
                 # Normal user/system/tool content
                 message = cast(
@@ -124,49 +120,43 @@ def convert_to_openai_messages(
         # Extract completions (assistant outputs after tool responses)
         for comp in pc_entry["completion"]:
             if comp.get("role") == "assistant":
-                if comp.get("content"):
-                    message = ChatCompletionMessage(role="assistant", content=comp["content"])
-                    messages.append(message)
-                elif comp.get("finish_reason") == "tool_calls":
-                    if len(tool_requests) == 0:
-                        raise ValueError("No tool requests available for tool_calls completion")
-                    tool_req = tool_requests.pop(0)
-                    # TODO: this is a hack because tracing frameworks did not report the tool call properly (?)
-                    message = ChatCompletionMessage(
-                        role="assistant",
-                        tool_calls=[
-                            ChatCompletionMessageFunctionToolCall(
-                                id=tool_req["call"]["id"],
-                                type=tool_req["call"]["type"],
-                                function=Function(name=tool_req["name"], arguments=tool_req["parameters"]),
-                            )
-                        ],
+                content = comp.get("content")
+                if pc_entry["tools"]:
+                    tool_calls = [
+                        ChatCompletionMessageFunctionToolCallParam(
+                            id=tool["call"]["id"],
+                            type=tool["call"]["type"],
+                            function={"name": tool["name"], "arguments": tool["parameters"]},
+                        )
+                        for tool in pc_entry["tools"]
+                    ]
+                    messages.append(
+                        ChatCompletionAssistantMessageParam(role="assistant", content=content, tool_calls=tool_calls)
                     )
-                    messages.append(message)
                 else:
-                    raise ValueError(f"Unsupported assistant completion: {comp}")
+                    messages.append(ChatCompletionAssistantMessageParam(role="assistant", content=content))
 
         # Build tools definitions (if available)
         if "functions" in pc_entry["request"]:
-            for fn in pc_entry["request"]["functions"]:
-                tools.append(
-                    ChatCompletionFunctionToolParam(
-                        type="function",
-                        function=FunctionDefinition(
-                            name=fn["name"],
-                            description=fn.get("description", ""),
-                            parameters=(
-                                json.loads(fn["parameters"]) if isinstance(fn["parameters"], str) else fn["parameters"]
-                            ),
+            tools = [
+                ChatCompletionFunctionToolParam(
+                    type="function",
+                    function={
+                        "name": fn["name"],
+                        "description": fn.get("description", ""),
+                        "parameters": (
+                            json.loads(fn["parameters"]) if isinstance(fn["parameters"], str) else fn["parameters"]
                         ),
-                    )
+                    },
                 )
+                for fn in pc_entry["request"]["functions"]
+            ]
             yield OpenAIMessages(messages=messages, tools=tools)
         else:
-            yield OpenAIMessages(messages=messages, tools=tools)
+            yield OpenAIMessages(messages=messages, tools=None)
 
 
-class TraceMessagesAdapter(TraceAdapter[List[OpenAIMessages]]):
+class TraceToMessages(TraceAdapter[List[OpenAIMessages]]):
     """
     Adapter that converts OpenTelemetry trace spans into OpenAI-compatible message format.
 
@@ -180,25 +170,29 @@ class TraceMessagesAdapter(TraceAdapter[List[OpenAIMessages]]):
     - Extracting and matching tool calls with their corresponding requests
     - Building proper OpenAI ChatCompletionMessage objects with roles, content, and tool calls
     - Generating function definitions for tools used in conversations
-
-    Returns:
-        List[OpenAIMessages]: A list of structured message conversations with associated tools
     """
 
+    def get_tool_calls(self, completion: Span, all_spans: List[Span], /) -> Iterable[Dict[str, Any]]:
+        """Find tool calls in the trace. Returns a dict with the tool call id, name, and arguments.
+
+        The spans that are direct children of the completion span are the tool calls.
+        """
+        # Get all the spans that are children of the completion span
+        children = [span for span in all_spans if span.parent_id == completion.span_id]
+        # Get the tool calls from the children
+        for maybe_tool_call in children:
+            tool_call = group_genai_dict(maybe_tool_call.attributes, "tool")
+            if not isinstance(tool_call, dict):
+                raise ValueError(f"Extracted tool call from trace is not a dict: {tool_call}")
+            if tool_call:
+                yield tool_call
+
     def adapt(self, source: List[Span], /) -> List[OpenAIMessages]:
-        raw_tool_calls: List[Dict[str, Any]] = []
         raw_prompt_completions: List[_RawSpanInfo] = []
 
         for span in source:
             attributes = {k: v for k, v in span.attributes.items()}
 
-            # Otherwise we strip all the tool calls and prompts and responses
-            tool_call = group_genai_dict(dict(attributes), "tool")
-            if not isinstance(tool_call, dict):
-                raise ValueError(f"Extracted tool call from trace is not a dict: {tool_call}")
-            if tool_call:
-                raw_tool_calls.append(tool_call)
-
             # Get all related information from the trace span
             prompt = group_genai_dict(attributes, "gen_ai.prompt") or []
             completion = group_genai_dict(attributes, "gen_ai.completion") or []
@@ -213,8 +207,11 @@ class TraceMessagesAdapter(TraceAdapter[List[OpenAIMessages]]):
             if not isinstance(response, dict):
                 raise ValueError(f"Extracted response from trace is not a dict: {response}")
             if prompt or completion or request or response:
+                tools = list(self.get_tool_calls(span, source)) or []
                 raw_prompt_completions.append(
-                    _RawSpanInfo(prompt=prompt or [], completion=completion, request=request, response=response)
+                    _RawSpanInfo(
+                        prompt=prompt or [], completion=completion, request=request, response=response, tools=tools
+                    )
                 )
 
-        return list(convert_to_openai_messages(raw_prompt_completions, raw_tool_calls))
+        return list(convert_to_openai_messages(raw_prompt_completions))

agentlightning/adapter/triplet.py

@@ -521,13 +521,13 @@ class TraceTree:
         )
 
 
-class BaseTraceTripletAdapter(TraceAdapter[List[Triplet]]):
+class TraceToTripletBase(TraceAdapter[List[Triplet]]):
     """
     Base class for trace triplet adapters.
     """
 
 
-class TraceTripletAdapter(BaseTraceTripletAdapter):
+class TracerTraceToTriplet(TraceToTripletBase):
     """
     An adapter to convert OpenTelemetry spans to triplet data.
 
@@ -600,10 +600,10 @@ class TraceTripletAdapter(BaseTraceTripletAdapter):
         return trajectory
 
 
-class LlmProxyTripletAdapter(BaseTraceTripletAdapter):
+class LlmProxyTraceToTriplet(TraceToTripletBase):
     """
     Converting telemetry data emitted by the LLM Proxy to triplet data.
-    This adapter is very experimental. Should only be used when the TraceTripletAdapter does not work at all.
+    This adapter is very experimental. Should only be used when the TracerTraceToTriplet does not work at all.
 
     IMPORTANT: Do NOT rely on timestamps here. Proxy spans can be emitted from different
     machines with unsynchronized clocks. We therefore treat `sequence_id` as the only

agentlightning/algorithm/__init__.py

@@ -1,5 +1,29 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .base import BaseAlgorithm
+from __future__ import annotations
 
-__all__ = ["BaseAlgorithm"]
+from typing import TYPE_CHECKING, Any
+
+from .base import BaseAlgorithm
+from .decorator import algo
+from .fast import Baseline, FastAlgorithm
+
+if TYPE_CHECKING:
+    from .apo import APO as APOType
+    from .verl import VERL as VERLType
+
+__all__ = ["BaseAlgorithm", "algo", "FastAlgorithm", "Baseline", "APO", "VERL"]
+
+# Shortcuts for usages like algo.APO(...)
+
+
+def APO(*args: Any, **kwargs: Any) -> APOType[Any]:
+    from .apo import APO as APOImplementation
+
+    return APOImplementation(*args, **kwargs)
+
+
+def VERL(*args: Any, **kwargs: Any) -> VERLType:
+    from .verl import VERL as VERLImplementation
+
+    return VERLImplementation(*args, **kwargs)

agentlightning/algorithm/apo/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .apo import APO
+
+__all__ = ["APO"]

agentlightning/algorithm/apo/apo.py

@@ -0,0 +1,891 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+APO with textual gradients that read rollout spans and outputs to modify the prompt.
+
+- algo: beam search with span-aware textual gradients -> apply_edit via LLM
+- rollout: same pattern as your example, but task is a dict (T_task)
+"""
+
+import asyncio
+import logging
+import random
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Counter, Dict, Generic, Iterator, List, Optional, Sequence, Set, Tuple, TypedDict, TypeVar, cast
+
+import poml
+from openai import AsyncOpenAI
+
+from agentlightning.adapter.messages import TraceToMessages
+from agentlightning.algorithm.base import BaseAlgorithm
+from agentlightning.reward import find_final_reward
+from agentlightning.types import Dataset, NamedResources, PromptTemplate, Rollout, RolloutMode, RolloutStatus
+
+logger = logging.getLogger(__name__)
+
+T_task = TypeVar("T_task")
+
+
+class RolloutResultForAPO(TypedDict):
+    """This must be all JSON serializable to be processable by POML."""
+
+    status: RolloutStatus
+    final_reward: Optional[float]
+    spans: List[Dict[str, Any]]
+    messages: List[Any]
+
+
+@dataclass
+class VersionedPromptTemplate:
+    version: str
+    prompt_template: PromptTemplate
+    score: Optional[float] = None
+
+
+GRADIENT_PROMPT_FILES = [
+    Path(__file__).parent / "prompts" / "text_gradient_variant01.poml",
+    Path(__file__).parent / "prompts" / "text_gradient_variant02.poml",
+    Path(__file__).parent / "prompts" / "text_gradient_variant03.poml",
+]
+
+APPLY_EDIT_PROMPT_FILES = [
+    Path(__file__).parent / "prompts" / "apply_edit_variant01.poml",
+    Path(__file__).parent / "prompts" / "apply_edit_variant02.poml",
+]
+
+
+def batch_iter_over_dataset(dataset: Dataset[T_task], batch_size: int) -> Iterator[Sequence[T_task]]:
+    """
+    Create an infinite iterator that yields batches from the dataset.
+
+    When batch_size >= dataset size, yields the entire shuffled dataset repeatedly.
+    When batch_size < dataset size, yields batches of the specified size, reshuffling
+    after each complete pass through the dataset.
+
+    Args:
+        dataset: The dataset to iterate over.
+        batch_size: The desired batch size.
+
+    Yields:
+        Sequences of tasks from the dataset. Each task appears at most once per epoch.
+    """
+    if batch_size >= len(dataset):
+        while True:
+            dataset_copy = [dataset[i] for i in range(len(dataset))]
+            random.shuffle(dataset_copy)
+            yield dataset_copy
+
+    else:
+        current_batch: List[int] = []
+        while True:
+            indices = list(range(len(dataset)))
+            random.shuffle(indices)
+            for index in indices:
+                if index in current_batch:
+                    continue
+                current_batch.append(index)
+                if len(current_batch) == batch_size:
+                    yield [dataset[index] for index in current_batch]
+                    current_batch = []
+
+
+class APO(BaseAlgorithm, Generic[T_task]):
+    """Automatic Prompt Optimization (APO) algorithm using textual gradients and beam search.
+
+    APO is an iterative prompt optimization algorithm that uses LLM-generated textual gradients
+    to improve prompts through a beam search process. It evaluates prompts on rollouts,
+    computes critiques based on the results, and applies edits to generate improved prompts.
+
+    The algorithm operates in rounds, where each round:
+    1. Samples parent prompts from the current beam
+    2. Generates new prompts by computing textual gradients and applying edits
+    3. Evaluates all candidates on a validation set
+    4. Selects the top-k prompts for the next round
+
+    Based on the ideas from:
+    - ProTeGi: https://aclanthology.org/2023.emnlp-main.494.pdf
+    - TextGrad: https://github.com/zou-group/textgrad
+    """
+
+    def __init__(
+        self,
+        async_openai_client: AsyncOpenAI,
+        *,
+        gradient_model: str = "gpt-5-mini",
+        apply_edit_model: str = "gpt-4.1-mini",
+        diversity_temperature: float = 1.0,
+        gradient_batch_size: int = 4,
+        val_batch_size: int = 16,
+        beam_width: int = 4,
+        branch_factor: int = 4,
+        beam_rounds: int = 3,
+        rollout_batch_timeout: float = 3600.0,
+        run_initial_validation: bool = True,
+        # Internal flags for debugging
+        _poml_trace: bool = False,
+    ):
+        """
+        Initialize the APO algorithm with configuration parameters.
+
+        Args:
+            async_openai_client: AsyncOpenAI client for making LLM API calls.
+            gradient_model: Model name for computing textual gradients (critiques).
+            apply_edit_model: Model name for applying edits based on critiques.
+            diversity_temperature: Temperature parameter for LLM calls to control diversity.
+            gradient_batch_size: Number of rollout results to sample for gradient computation.
+            val_batch_size: Number of validation examples to use for evaluation.
+            beam_width: Number of top-scoring prompts to keep in the beam at each round.
+            branch_factor: Number of new prompt candidates to generate from each parent prompt
+                by applying textual gradient edits. This controls the expansion of the search tree.
+            beam_rounds: Number of beam search rounds to perform.
+            rollout_batch_timeout: Maximum time in seconds to wait for rollout batch completion.
+            run_initial_validation: If True, runs validation on the seed prompt before starting
+                optimization to establish a baseline score. Defaults to True.
+        """
+        self.async_openai_client = async_openai_client
+        self.gradient_model = gradient_model
+        self.apply_edit_model = apply_edit_model
+        self.diversity_temperature = diversity_temperature
+        self.gradient_batch_size = gradient_batch_size
+        self.val_batch_size = val_batch_size
+        self.beam_width = beam_width
+        self.branch_factor = branch_factor
+        self.beam_rounds = beam_rounds
+        self.rollout_batch_timeout = rollout_batch_timeout
+        self.run_initial_validation = run_initial_validation
+
+        self._history_best_prompt: Optional[PromptTemplate] = None
+        self._history_best_score: float = float("-inf")
+        self._history_best_version: Optional[str] = None
+
+        self._version_counter: int = 0
+
+        self._poml_trace = _poml_trace
+
+    def _create_versioned_prompt(
+        self,
+        prompt_template: PromptTemplate,
+        *,
+        score: Optional[float] = None,
+    ) -> VersionedPromptTemplate:
+        """
+        Wrap a prompt template with a new monotonically increasing version identifier.
+        """
+        version = f"v{self._version_counter}"
+        self._version_counter += 1
+        return VersionedPromptTemplate(version=version, prompt_template=prompt_template, score=score)
+
+    def _format_log_prefix(
+        self,
+        *,
+        round_num: Optional[int] = None,
+        beam_idx: Optional[int] = None,
+        branch_idx: Optional[int] = None,
+        prompt_version: Optional[str] = None,
+    ) -> str:
+        """
+        Construct the standardized log prefix.
+        """
+        parts: List[str] = []
+        if round_num is not None:
+            parts.append(f"Round {round_num:02d}")
+        if beam_idx is not None:
+            parts.append(f"Beam {beam_idx:02d}")
+        if branch_idx is not None:
+            parts.append(f"Branch {branch_idx:02d}")
+        if prompt_version is not None:
+            parts.append(f"Prompt {prompt_version}")
+        if not parts:
+            return ""
+        return f"[{' | '.join(parts)}]"
+
+    def _log(self, level: int, message: str, *, prefix: Optional[str] = None) -> None:
+        """
+        Log a message with an optional standardized prefix.
+        """
+        effective_prefix = prefix
+        if effective_prefix:
+            logger.log(level, f"{effective_prefix} {message}")
+        else:
+            logger.log(level, message)
+
+    def get_seed_prompt_template(self) -> Tuple[str, PromptTemplate]:
+        """
+        Extract the initial prompt template from the algorithm's resources.
+
+        Returns:
+            A tuple of (resource_name, prompt_template) representing the seed prompt.
+
+        Raises:
+            ValueError: If initial_resources is not set or no PromptTemplate is found.
+        """
+        initial_resources = self.get_initial_resources()
+        if initial_resources is None:
+            raise ValueError(
+                "initial_resources are not set for APO algorithm. "
+                "Use algorithm.set_initial_resources() to set initial resources or set it in Trainer()"
+            )
+        for name, resource in initial_resources.items():
+            if isinstance(resource, PromptTemplate):
+                return name, resource
+        raise ValueError("No prompt template resource found in initial_resources")
+
+    def get_adapter(self) -> TraceToMessages:
+        """
+        Get the adapter for converting spans to messages.
+
+        Returns:
+            The TraceToMessages instance for this algorithm.
+
+        Raises:
+            ValueError: If the adapter is not a TraceToMessages.
+        """
+        adapter = super().get_adapter()
+        if not isinstance(adapter, TraceToMessages):
+            raise ValueError("Adapter must be a TraceToMessages for APO algorithm")
+        return adapter
+
+    def get_best_prompt(self) -> PromptTemplate:
+        """
+        Retrieve the best prompt discovered during optimization.
+
+        Returns:
+            The prompt template with the highest validation score found so far.
+
+        Raises:
+            ValueError: If no best prompt has been found yet (run() not called).
+        """
+        if self._history_best_prompt is None:
+            raise ValueError("No best prompt found")
+        return self._history_best_prompt
+
+    async def compute_textual_gradient(
+        self,
+        current_prompt: VersionedPromptTemplate,
+        rollout_results: List[RolloutResultForAPO],
+        *,
+        prefix: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Compute a textual gradient (critique) for the current prompt based on rollout results.
+
+        This method samples rollout results, sends them to an LLM along with the current prompt,
+        and generates a critique describing how the prompt could be improved.
+
+        Args:
+            current_prompt: The prompt template to critique.
+            rollout_results: List of rollout results containing spans, messages, and rewards.
+
+        Returns:
+            A textual critique generated by the LLM, or None if generation fails.
+        """
+        tg_template = random.choice(GRADIENT_PROMPT_FILES)
+
+        if len(rollout_results) < self.gradient_batch_size:
+            self._log(
+                logging.WARNING,
+                f"Only {len(rollout_results)} rollouts available, but {self.gradient_batch_size} are needed. Using all rollouts.",
+                prefix=prefix,
+            )
+            sampled_rollout_results = rollout_results
+        else:
+            sampled_rollout_results = random.sample(rollout_results, self.gradient_batch_size)
+
+        self._log(
+            logging.INFO,
+            f"Gradient will be computed with {self.gradient_model} for {len(sampled_rollout_results)} rollouts with template: {tg_template.name}",
+            prefix=prefix,
+        )
+
+        tg_msg = poml.poml(  # type: ignore
+            tg_template,
+            context={
+                "experiments": sampled_rollout_results,
+                "prompt_template": current_prompt.prompt_template.template,
+            },
+            format="openai_chat",
+        )
+        self._log(
+            logging.DEBUG,
+            f"Gradient computed with {self.gradient_model} prompt: {tg_msg}",
+            prefix=prefix,
+        )
+        critique_response = await self.async_openai_client.chat.completions.create(
+            model=self.gradient_model,
+            messages=tg_msg["messages"],  # type: ignore
+            temperature=self.diversity_temperature,
+        )
+        critique_text = critique_response.choices[0].message.content
+        self._log(
+            logging.INFO,
+            f"Gradient computed with {self.gradient_model} has result: {critique_text}",
+            prefix=prefix,
+        )
+
+        return critique_text
+
+    async def textual_gradient_and_apply_edit(
+        self,
+        current_prompt: VersionedPromptTemplate,
+        rollout: List[RolloutResultForAPO],
+        *,
+        prefix: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Generate an improved prompt by computing a textual gradient and applying an edit.
+
+        This is the main optimization step that:
+        1. Computes a critique (textual gradient) based on rollout performance
+        2. Uses another LLM to apply the critique and generate an improved prompt
+
+        Args:
+            current_prompt: The current prompt template to improve.
+            rollout: List of rollout results to base the critique on.
+
+        Returns:
+            The improved prompt text, or the original prompt if gradient computation fails.
+        """
+        # 1) Critique
+        critique_text = await self.compute_textual_gradient(
+            current_prompt,
+            rollout,
+            prefix=prefix,
+        )
+        if not critique_text:
+            self._log(
+                logging.ERROR,
+                "Failed to compute critique for prompt.",
+                prefix=prefix,
+            )
+            return current_prompt.prompt_template.template
+
+        # 2) Apply edit
+        ae_template = random.choice(APPLY_EDIT_PROMPT_FILES)
+        self._log(
+            logging.INFO,
+            f"Edit will be generated by {self.apply_edit_model} with template: {ae_template.name}",
+            prefix=prefix,
+        )
+        ae_msg = poml.poml(  # type: ignore
+            ae_template,
+            context={
+                "prompt_template": current_prompt.prompt_template.template,
+                "critique": critique_text,
+            },
+            format="openai_chat",
+        )
+
+        ae_response = await self.async_openai_client.chat.completions.create(
+            model=self.apply_edit_model,
+            messages=ae_msg["messages"],  # type: ignore
+            temperature=self.diversity_temperature,
+        )
+        new_prompt = ae_response.choices[0].message.content
+        if new_prompt:
+            self._log(
+                logging.INFO,
+                f"Edit generated by {self.apply_edit_model}: {new_prompt[:50]}...",
+                prefix=prefix,
+            )
+        return new_prompt
+
+    async def get_rollout_results(
+        self,
+        rollout: List[Rollout],
+        *,
+        prefix: Optional[str] = None,
+    ) -> List[RolloutResultForAPO]:
+        """
+        Convert completed rollouts to APO-compatible result format.
+
+        Fetches spans for each rollout, adapts them to messages, and packages them
+        with rewards and status information for gradient computation.
+
+        Args:
+            rollout: List of completed rollout metadata.
+
+        Returns:
+            List of rollout results formatted for APO processing.
+        """
+        rollout_results: List[RolloutResultForAPO] = []
+        store = self.get_store()
+        adapter = self.get_adapter()
+        for r in rollout:
+            spans = await store.query_spans(r.rollout_id)
+            messages = adapter.adapt(spans)
+            rollout_result = RolloutResultForAPO(
+                status=r.status,
+                final_reward=find_final_reward(spans),
+                spans=[span.model_dump() for span in spans],
+                messages=messages,
+            )
+            self._log(
+                logging.DEBUG,
+                f"Rollout result for {r.rollout_id}: status {rollout_result['status']} with final reward {rollout_result['final_reward']}. "
+                f"{len(rollout_result['spans'])} spans and {len(rollout_result['messages'])} messages.",
+                prefix=prefix,
+            )
+            rollout_results.append(rollout_result)
+        return rollout_results
+
+    async def evaluate_prompt_on_batch(
+        self,
+        prompt: VersionedPromptTemplate,
+        resource_name: str,
+        dataset: Sequence[T_task],
+        mode: RolloutMode,
+        *,
+        prefix: Optional[str] = None,
+    ) -> Tuple[List[RolloutResultForAPO], float]:
+        """
+        Evaluate a prompt on a batch of tasks by running rollouts and computing average reward.
+
+        This method:
+        1. Adds the prompt as a named resource to the store
+        2. Enqueues rollouts for each task in the dataset
+        3. Waits for rollouts to complete (with timeout)
+        4. Computes and returns the average reward
+
+        Args:
+            prompt: The prompt template string to evaluate.
+            resource_name: The name to register the prompt under in the store.
+            dataset: Sequence of tasks to evaluate the prompt on.
+            mode: Rollout mode ("train" or "val") for logging/tracking.
+
+        Returns:
+            A tuple of (rollout_results, average_reward) where rollout_results contains
+            detailed information for each rollout and average_reward is the mean final reward.
+        """
+        store = self.get_store()
+        preview = prompt.prompt_template.template[:50]
+        self._log(
+            logging.INFO,
+            f'Evaluating prompt "{preview}..." on {len(dataset)} tasks in {mode} mode',
+            prefix=prefix,
+        )
+
+        # Install prompt as named resource
+        resources: NamedResources = {resource_name: prompt.prompt_template}
+        resource_update = await store.update_resources(prompt.version, resources)
+
+        rollout_ids: List[str] = []
+        for t in dataset:
+            r = await store.enqueue_rollout(input=t, mode=mode, resources_id=resource_update.resources_id)
+            rollout_ids.append(r.rollout_id)
+
+        deadline = time.time() + self.rollout_batch_timeout
+        finished: List[Rollout] = []
+        while time.time() < deadline:
+            finished = await store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=0.0)
+            if len(finished) >= len(rollout_ids):
+                self._log(
+                    logging.INFO,
+                    f"All {len(rollout_ids)} rollouts finished within timeout.",
+                    prefix=prefix,
+                )
+                break
+            else:
+                self._log(
+                    logging.DEBUG,
+                    f"Only {len(finished)} rollouts finished within timeout. Waiting for remaining {len(rollout_ids) - len(finished)} rollouts.",
+                    prefix=prefix,
+                )
+                # Sleep to avoid busy-waiting
+                await asyncio.sleep(2.0)
+
+        rollout_results = await self.get_rollout_results(
+            finished,
+            prefix=prefix,
+        )
+        final_rewards = [rr["final_reward"] for rr in rollout_results]
+
+        avg = float(sum([r or 0.0 for r in final_rewards]) / max(1, len(final_rewards)))
+        status_counter = Counter([rr["status"] for rr in rollout_results])
+
+        self._log(
+            logging.INFO,
+            f"Evaluated {len(rollout_results)} rollouts. Statuses: {status_counter}. Rewards: {final_rewards}, average is {avg}",
+            prefix=prefix,
+        )
+        return rollout_results, avg
+
+    def _initialize_beam(
+        self,
+        train_dataset: Optional[Dataset[T_task]],
+        val_dataset: Optional[Dataset[T_task]],
+    ) -> Tuple[str, PromptTemplate, Iterator[Sequence[T_task]], Iterator[Sequence[T_task]]]:
+        """
+        Initialize the beam search with seed prompt and dataset iterators.
+
+        Args:
+            train_dataset: Dataset for computing gradients.
+            val_dataset: Dataset for evaluating prompts.
+
+        Returns:
+            Tuple of (resource_name, seed_prompt, grad_iterator, val_iterator).
+
+        Raises:
+            ValueError: If either dataset is None.
+        """
+        resource_name, seed_prompt = self.get_seed_prompt_template()
+
+        if train_dataset is None:
+            raise ValueError("train_dataset is required for APO algorithm")
+        if val_dataset is None:
+            raise ValueError("val_dataset is required for APO algorithm")
+
+        grad_dataset_iterator = batch_iter_over_dataset(train_dataset, self.gradient_batch_size)
+        val_dataset_iterator = batch_iter_over_dataset(val_dataset, self.val_batch_size)
+
+        # Initialize history tracking
+        self._history_best_prompt = seed_prompt
+        self._history_best_score = float("-inf")
+
+        return resource_name, seed_prompt, grad_dataset_iterator, val_dataset_iterator
+
+    def _sample_parent_prompts(
+        self,
+        beam: List[VersionedPromptTemplate],
+        round_num: int,
+    ) -> List[Tuple[int, VersionedPromptTemplate]]:
+        """
+        Sample parent prompts from the current beam for generating new candidates.
+
+        If the beam has fewer prompts than beam_width, replicates existing prompts.
+        Otherwise, randomly samples beam_width prompts.
+
+        Args:
+            beam: Current list of prompt templates in the beam.
+            round_num: Current round number (for logging, 0-indexed).
+
+        Returns:
+            List of parent prompts to generate children from.
+        """
+        display_round = round_num + 1
+        if len(beam) < self.beam_width:
+            prefix = self._format_log_prefix(round_num=display_round)
+            self._log(
+                logging.WARNING,
+                f"Beam width is currently {self.beam_width}, but only {len(beam)} prompts in beam. Replicating all prompts.",
+                prefix=prefix,
+            )
+            return [(i % len(beam), beam[i % len(beam)]) for i in range(self.beam_width)]
+
+        selected_indices = random.sample(range(len(beam)), self.beam_width)
+        return [(idx, beam[idx]) for idx in selected_indices]
+
+    async def _generate_candidate_prompts(
+        self,
+        parent_prompts: List[Tuple[int, VersionedPromptTemplate]],
+        resource_name: str,
+        grad_dataset_iterator: Iterator[Sequence[T_task]],
+        round_num: int,
+    ) -> List[VersionedPromptTemplate]:
+        """
+        Generate new candidate prompts from parents using textual gradients.
+
+        For each parent prompt, generates branch_factor new candidates by:
+        1. Evaluating the parent on a training batch
+        2. Computing textual gradient
+        3. Applying edit to generate improved prompt
+
+        Args:
+            parent_prompts: List of parent prompts to generate children from.
+            resource_name: Name to register prompts under in the store.
+            grad_dataset_iterator: Iterator over training data batches.
+            round_num: Current round number (for logging, 0-indexed).
+
+        Returns:
+            List of newly generated prompt templates.
+        """
+        display_round = round_num + 1
+        round_prefix = self._format_log_prefix(round_num=display_round)
+        self._log(
+            logging.INFO,
+            f"Applying {self.branch_factor} edits to each of the {len(parent_prompts)} parents based on "
+            "gradients computed on training dataset",
+            prefix=round_prefix,
+        )
+
+        parent_prompts_str = [
+            f"{p.version}:{p.score:.3f}" if p.score is not None else p.version for _, p in parent_prompts
+        ]
+        self._log(
+            logging.INFO,
+            f"Parent prompts: {', '.join(parent_prompts_str)}",
+            prefix=round_prefix,
+        )
+
+        candidates: List[VersionedPromptTemplate] = []
+        used_beam_indices: Set[int] = set()
+        for real_beam_idx, (beam_idx, prompt) in enumerate(parent_prompts):
+            if beam_idx in used_beam_indices:
+                beam_prefix = self._format_log_prefix(
+                    round_num=display_round,
+                    beam_idx=beam_idx + 1,
+                    prompt_version=prompt.version,
+                )
+                self._log(
+                    logging.WARNING,
+                    "Duplicated beam index found. Might be caused by beam_width too high. "
+                    + f"The real index of this beam is {real_beam_idx + 1}.",
+                    prefix=beam_prefix,
+                )
+            else:
+                used_beam_indices.add(beam_idx)
+            for branch_idx in range(self.branch_factor):
+                parent_prefix = self._format_log_prefix(
+                    round_num=display_round,
+                    beam_idx=beam_idx + 1,
+                    branch_idx=branch_idx + 1,
+                    prompt_version=prompt.version,
+                )
+                baseline_score = f"{prompt.score:.3f}" if prompt.score is not None else "N/A"
+                self._log(
+                    logging.INFO,
+                    f"Use parent prompt {prompt.version} as a baseline to generate a new prompt. Baseline score: {baseline_score}",
+                    prefix=parent_prefix,
+                )
+                grad_samples = next(grad_dataset_iterator)
+                rollout_results, _ = await self.evaluate_prompt_on_batch(
+                    prompt,
+                    resource_name,
+                    grad_samples,
+                    mode="train",
+                    prefix=parent_prefix,
+                )
+                new_prompt = await self.textual_gradient_and_apply_edit(
+                    prompt,
+                    rollout_results,
+                    prefix=parent_prefix,
+                )
+                if not new_prompt:
+                    self._log(
+                        logging.ERROR,
+                        f"Failed to compute edit for prompt: {prompt.prompt_template.template}",
+                        prefix=parent_prefix,
+                    )
+                    continue
+                new_prompt_template = PromptTemplate(template=new_prompt, engine="f-string")
+                versioned_candidate = self._create_versioned_prompt(new_prompt_template)
+                self._log(
+                    logging.INFO,
+                    f"New prompt template created from parent {prompt.version}: {versioned_candidate.version}",
+                    prefix=parent_prefix,
+                )
+                candidate_prefix = self._format_log_prefix(
+                    round_num=display_round, prompt_version=versioned_candidate.version
+                )
+                self._log(
+                    logging.INFO,
+                    f"New prompt template created from parent {prompt.version}:\n```\n{new_prompt}\n```",
+                    prefix=candidate_prefix,
+                )
+                candidates.append(versioned_candidate)
+
+        return candidates
+
+    async def _evaluate_and_select_beam(
+        self,
+        candidates: List[VersionedPromptTemplate],
+        resource_name: str,
+        val_dataset_iterator: Iterator[Sequence[T_task]],
+        round_num: int,
+    ) -> List[VersionedPromptTemplate]:
+        """
+        Evaluate all candidate prompts on validation data and select top-k for the beam.
+
+        Args:
+            candidates: List of candidate prompts to evaluate.
+            resource_name: Name to register prompts under in the store.
+            val_dataset_iterator: Iterator over validation data batches.
+            round_num: Current round number (for logging, 0-indexed).
+
+        Returns:
+            List of top beam_width prompts sorted by validation score (best first).
+
+        Raises:
+            ValueError: If no candidates remain after evaluation.
+        """
+        display_round = round_num + 1
+        round_prefix = self._format_log_prefix(round_num=display_round)
+        self._log(
+            logging.INFO,
+            f"Evaluating {len(candidates)} candidates on validation dataset",
+            prefix=round_prefix,
+        )
+
+        val_batch = next(val_dataset_iterator)
+
+        for prompt in candidates:
+            candidate_prefix = self._format_log_prefix(
+                round_num=display_round,
+                prompt_version=prompt.version,
+            )
+            _, score = await self.evaluate_prompt_on_batch(
+                prompt,
+                resource_name,
+                val_batch,
+                mode="val",
+                prefix=candidate_prefix,
+            )
+            prompt.score = score
+            self._log(
+                logging.INFO,
+                f"Candidate score: {score:.3f}",
+                prefix=candidate_prefix,
+            )
+
+        # Sort by score (descending) and select top beam_width
+        sorted_prompts = [p for p in sorted(candidates, key=lambda x: cast(float, x.score), reverse=True)]
+        selected_prompts = sorted_prompts[: self.beam_width]
+        selected_versions = [
+            f"{prompt.version}:{prompt.score:.3f}" if prompt.score is not None else prompt.version
+            for prompt in selected_prompts
+        ]
+        self._log(
+            logging.INFO,
+            f"Top {len(selected_prompts)} candidates on validation dataset: {selected_versions}",
+            prefix=round_prefix,
+        )
+
+        if len(selected_prompts) == 0:
+            raise ValueError("No beam candidates any more")
+
+        return selected_prompts
+
+    async def _update_best_prompt(
+        self,
+        beam: List[VersionedPromptTemplate],
+        resource_name: str,
+        val_dataset: Dataset[T_task],
+        round_num: int,
+    ) -> None:
+        """
+        Evaluate the best prompt in the beam on the full validation set and update history.
+
+        Args:
+            beam: Current beam of prompts (sorted, best first).
+            resource_name: Name to register prompts under in the store.
+            val_dataset: Full validation dataset.
+            round_num: Current round number (for logging, 0-indexed).
+        """
+        display_round = round_num + 1
+        best_prompt = beam[0]
+        prefix = self._format_log_prefix(round_num=display_round, prompt_version=best_prompt.version)
+        _, best_score = await self.evaluate_prompt_on_batch(
+            best_prompt,
+            resource_name,
+            cast(Sequence[T_task], val_dataset),
+            mode="val",
+            prefix=prefix,
+        )
+        self._log(
+            logging.INFO,
+            f"Beam leader score: {best_score:.3f}",
+            prefix=prefix,
+        )
+
+        if best_score > self._history_best_score:
+            prev = self._history_best_score
+            self._log(
+                logging.INFO,
+                f"Best prompt updated. New best score: {best_score:.3f} (prev: {prev:.3f})",
+                prefix=prefix,
+            )
+            self._history_best_prompt = best_prompt.prompt_template
+            self._history_best_score = best_score
+            self._history_best_version = best_prompt.version
+        else:
+            self._log(
+                logging.WARNING,
+                f"Best prompt not updated. Current score: {best_score:.3f} vs. history best: {self._history_best_score:.3f})",
+                prefix=prefix,
+            )
+
+    async def run(
+        self,
+        train_dataset: Optional[Dataset[T_task]] = None,
+        val_dataset: Optional[Dataset[T_task]] = None,
+    ) -> None:
+        """
+        Execute the APO algorithm to optimize prompts through beam search with textual gradients.
+
+        The algorithm performs iterative prompt optimization over multiple rounds:
+        - Each round: samples parent prompts, generates new candidates via textual gradients,
+          evaluates all candidates on validation data, and keeps the top performers
+        - Tracks the historically best prompt across all rounds
+        - Uses different training data samples for each gradient computation to ensure diversity
+
+        Args:
+            train_dataset: Dataset of tasks for computing textual gradients. Required.
+            val_dataset: Dataset of tasks for evaluating and selecting prompts. Required.
+
+        Raises:
+            ValueError: If train_dataset or val_dataset is None, or if resources are not set.
+        """
+        # Initialize beam search
+        resource_name, seed_prompt, grad_iterator, val_iterator = self._initialize_beam(train_dataset, val_dataset)
+
+        if self._poml_trace:
+            poml.set_trace(trace_dir="pomltrace")
+
+        # Validation datasets are guaranteed to be non-None after initialization
+        assert val_dataset is not None
+
+        # Start with seed prompt in the beam
+        seed_versioned = self._create_versioned_prompt(seed_prompt)
+        beam: List[VersionedPromptTemplate] = [seed_versioned]
+        self._history_best_prompt = seed_prompt
+        self._history_best_version = seed_versioned.version
+
+        # Optionally evaluate seed prompt on validation set to establish baseline
+        if self.run_initial_validation:
+            seed_prefix = self._format_log_prefix(round_num=0, prompt_version=seed_versioned.version)
+            self._log(
+                logging.INFO,
+                "Evaluating seed prompt on validation dataset before optimization...",
+                prefix=seed_prefix,
+            )
+            _, seed_score = await self.evaluate_prompt_on_batch(
+                seed_versioned,
+                resource_name,
+                cast(Sequence[T_task], val_dataset),
+                mode="val",
+                prefix=seed_prefix,
+            )
+            self._log(
+                logging.INFO,
+                f"Seed prompt baseline score: {seed_score:.3f}",
+                prefix=seed_prefix,
+            )
+            self._history_best_prompt = seed_prompt
+            self._history_best_score = seed_score
+            self._history_best_version = seed_versioned.version
+
+        # Run beam search for specified number of rounds
+        for rnd in range(self.beam_rounds):
+            display_round = rnd + 1
+            round_prefix = self._format_log_prefix(round_num=display_round)
+            self._log(
+                logging.INFO,
+                f"Round {display_round}/{self.beam_rounds}...",
+                prefix=round_prefix,
+            )
+
+            # Sample parent prompts from current beam
+            parent_prompts = self._sample_parent_prompts(beam, rnd)
+
+            # Generate new candidate prompts from parents
+            new_candidates = await self._generate_candidate_prompts(parent_prompts, resource_name, grad_iterator, rnd)
+
+            # Combine existing beam with new candidates
+            all_candidates = [*beam, *new_candidates]
+
+            # Evaluate and select top-k prompts for next beam
+            beam = await self._evaluate_and_select_beam(all_candidates, resource_name, val_iterator, rnd)
+
+            # Update historically best prompt if improved
+            await self._update_best_prompt(beam, resource_name, val_dataset, rnd)

agentlightning/algorithm/apo/prompts/apply_edit_variant01.poml

@@ -0,0 +1,22 @@
+<poml>
+  <p>Revise the given prompt template using the critique as constraints and improvement guide.</p>
+  <cp caption="Revision Rules">
+    <list listStyle="decimal">
+      <item>Rewrite or restructure the prompt if critique implies it.</item>
+      <item>Explicitly include any requested output format, structure, or word limit, if requested by the critique.</item>
+      <item>Prioritize mechanism-first phrasing: define what to do, then how to do it.</item>
+      <item>Preserve placeholder variables inside curly brackets.</item>
+    </list>
+  </cp>
+  <output-format>
+    Return only the improved prompt template with placeholders intact. Do not include other explanations on how you did it, or headers and introductory texts.
+  </output-format>
+  <human-msg>
+    <cp caption="Prompt Template">
+      <text whiteSpace="pre">{{ prompt_template }}</text>
+    </cp>
+    <cp caption="Critique">
+      <text whiteSpace="pre">{{ critique }}</text>
+    </cp>
+  </human-msg>
+</poml>

agentlightning/algorithm/apo/prompts/apply_edit_variant02.poml

@@ -0,0 +1,18 @@
+<!-- Conservative Edit Prompt -->
+
+<poml>
+  <p>Revise the prompt to address ONE critique point clearly and effectively. Preserve all variable names in curly-brackets.</p>
+  <p>Do not address more than one critique point. Focus on the single most critical issue.</p>
+  <p>Keep the new prompt close in tone, length, and structure to the original.</p>
+  <output-format>
+    Return only the revised full prompt. Do not include explanations, comparisons, or other text.
+  </output-format>
+  <human-msg>
+    <cp caption="PROMPT" level="3">
+      <text whiteSpace="pre">{{ prompt_template }}</text>
+    </cp>
+    <cp caption="CRITIQUE" level="3">
+      <text whiteSpace="pre">{{ critique }}</text>
+    </cp>
+  </human-msg>
+</poml>

agentlightning/algorithm/apo/prompts/text_gradient_variant01.poml

@@ -0,0 +1,18 @@
+<poml>
+  <p>You optimize a prompt template.</p>
+  <cp caption="Original Prompt Template">
+    <text whiteSpace="pre">{{ prompt_template }}</text>
+  </cp>
+  <cp caption="Experiments with Original Prompt Template">
+    <cp for="experiment in experiments" caption="Experiment {{ loop.index + 1 }}">
+      <p>This experiment has {{ experiment.status }}. It gets a final reward: {{ experiment.final_reward }}</p>
+      <cp caption="Rollout Traces (Chat Messages, Grader Requests included)">
+        <object data="{{ experiment.messages }}" />
+      </cp>
+    </cp>
+  </cp>
+  <cp caption="Your Task">
+    Produce a brief critique listing specific causes for the error or ways to raise reward next time.
+    Return a bullet list with concrete, testable changes (format, constraints, ordering, definitions).
+  </cp>
+</poml>

agentlightning/algorithm/apo/prompts/text_gradient_variant02.poml

@@ -0,0 +1,16 @@
+<poml>
+  <role>You are a prompt engineer.</role>
+  <task>Analyze where the current prompt failed to elicit the right mechanism.</task>
+  <cp caption="Current Prompt Template">
+    <text whiteSpace="pre">{{ prompt_template }}</text>
+  </cp>
+  <cp caption="Sample Runs with Current Prompt Template">
+    <p>The following are the OpenTelemetry spans collected from the sample runs with the current prompt template. They should contain both prompt, responses and rewards.</p>
+    <cp for="experiment in experiments" caption="Sample Run #{{ loop.index + 1 }} Diagnostics">
+      <object for="span in experiment.spans" data="{{ span }}" />
+    </cp>
+  </cp>
+  <output-format>
+    Write 3-5 short bullets titled 'Critique:' focusing on missing constraints, ordering, or formatting.
+  </output-format>
+</poml>

agentlightning/algorithm/apo/prompts/text_gradient_variant03.poml

@@ -0,0 +1,107 @@
+<poml>
+
+  <role>You are an expert prompt engineer.</role>
+
+  <task>Your task is to analyze the prompt and provide a critique of the prompt. Follow the steps below to create the critique.
+
+  <cp caption="1. Structural Issues">
+    <p>These flaws block clarity and logic. Always check them first.</p>
+
+    <list>
+      <item><b>Missing goal</b>: The prompt never defines what success looks like. Ask: <i>Can I summarize its output goal in one line?</i></item>
+      <item><b>Contradictions</b>: Two or more instructions conflict. Search for words like *never*, *always*, *except*, *but also*.</item>
+      <item><b>Circular dependencies</b>: The model is told to do A before B and B before A.</item>
+      <item><b>No stop condition</b>: The prompt doesn’t say when the task is done. Flag any open-ended verbs: <i>explore,</i> <i>analyze further,</i> <i>continue indefinitely.</i></item>
+    </list>
+  </cp>
+
+  <cp caption="2. Instruction Quality">
+    <p>Examine how the instructions are stated and ordered to ensure clarity and enforceability.</p>
+    <list>
+      <item><b>Vague verbs</b>: Avoid terms like <i>optimize,</i> <i>improve,</i> and <i>ensure.</i> Use precise, measurable instructions.</item>
+      <item><b>Lack of hierarchy</b>: All rules appear equally important, making conflict resolution impossible. Clarify rule precedence.</item>
+      <item><b>Mixed abstraction</b>: High-level policies are interleaved with implementation details. Keep principles separate from step-by-step actions.</item>
+      <item><b>Overlapping scope</b>: Similar instructions appear in several sections with minor changes. Identify and consolidate duplicates.</item>
+    </list>
+  </cp>
+
+  <cp caption="3. Control and Behavior">
+    <p>Review boundaries on model autonomy, tool use, and communication style.</p>
+    <list>
+      <item><b>No tool limits</b>: Limits on tool calls, retries, or time not specified. Define boundaries for operations.</item>
+      <item><b>Unclear uncertainty handling</b>: Conflicting instructions regarding clarifying uncertainties vs. never asking users. Select one behavior.</item>
+      <item><b>Verbosity confusion</b>: Some parts demand detailed answers, others specify brevity. Highlight and resolve inconsistency.</item>
+      <item><b>Feedback omission</b>: No plan for progress reporting or preamble during multi-step operations.</item>
+    </list>
+  </cp>
+
+  <cp caption="4. Input and Output Specification">
+    <p>Assess if required data and expected output formats are clearly defined.</p>
+    <list>
+      <item><b>No input defaults</b>: What should happen if a needed value is absent or invalid isn’t explained.</item>
+      <item><b>Output schema missing</b>: Expected response format or sections are not spelled out.</item>
+      <item><b>Format inconsistency</b>: Output style (Markdown, JSON, XML, etc.) shifts mid-prompt. Ensure format requirements are stable.</item>
+      <item><b>No validation</b>: Lacks steps like <i>verify results before submitting</i> or <i>summarize at end.</i></item>
+    </list>
+  </cp>
+
+  <cp caption="5. Scope and Safety">
+    <p>Ensure prompt actions remain within safe, authorized boundaries.</p>
+    <list>
+      <item><b>Scope creep</b>: Open-ended statements such as <i>feel free to enhance</i> can justify unrelated changes.</item>
+      <item><b>Unsafe actions</b>: Allows deletions or modifications without explicit user approval.</item>
+      <item><b>No error handling</b>: What happens if a tool call fails or data is missing is not addressed.</item>
+      <item><b>User authority ambiguity</b>: Model may act for multiple users or perform irreversible actions without checks.</item>
+    </list>
+  </cp>
+
+  <cp caption="6. Efficiency and Maintainability">
+    <p>Consider the prompt’s length, redundancy, and future comprehensibility.</p>
+    <list>
+      <item><b>Overexplained</b>: Verbose explanations where concise, numbered steps suffice.</item>
+      <item><b>Redundancy</b>: Similar rules scattered in multiple aliases; centralize and summarize them.</item>
+      <item><b>Hidden assumptions</b>: Implicit defaults (like timezone, language) are not stated.</item>
+      <item><b>Poor auditability</b>: Lacks section markers (e.g., <code>&lt;policy&gt;</code>, <code>&lt;procedure&gt;</code>). Structure prompt for easy review.</item>
+    </list>
+  </cp>
+
+  <cp caption="7. Testing Method">
+    <p>Methodical approach for reviewing a prompt:</p>
+    <list>
+      <item>Read the prompt fully; highlight all unclear or contradictory instructions.</item>
+      <item>For each main area, answer:
+      <list listStyle="decimal">
+        <item>What is the intended outcome?</item>
+        <item>What is the stop or completion condition?</item>
+        <item>How are conflicts between rules resolved?</item>
+        <item>What are the explicit limits (tools, run time, tokens)?</item>
+        <item>What should the output format be?</item>
+      </list>
+      </item>
+      <item>Rate each section: <i>clear</i>, <i>incomplete</i>, <i>contradictory</i>, or <i>redundant</i>.</item>
+      <item>Summarize findings under categories: structure, control, scope, format, safety.</item>
+    </list>
+    <p>This method surfaces issues such as ambiguity, contradiction, missing boundaries, and output uncertainty—core failure modes in prompting identified by the GPT-5 prompting guide.</p>
+  </cp>
+  </task>
+
+  <output-format>
+    Respond with a complete analysis and critique of the prompt. Be concise and direct. Less than 350 words.
+  </output-format>
+
+  <human-msg>
+    <cp caption="Prompt">
+      <text whiteSpace="pre">{{ prompt_template }}</text>
+    </cp>
+    <cp caption="Sample Runs of the Prompts (Historical Messages and Rewards)">
+      <cp for="experiment in experiments" caption="Sample Run #{{ loop.index + 1 }}">
+        <cp caption="Overall Status">
+          This run has {{ experiment.status }}. The final score is {{ experiment.final_reward }}.
+        </cp>
+        <cp caption="Messages">
+          <object data="{{ experiment.messages }}" />
+        </cp>
+      </cp>
+    </cp>
+  </human-msg>
+</poml>

agentlightning/algorithm/base.py

@@ -2,22 +2,14 @@
 
 from __future__ import annotations
 
-import functools
 import inspect
 import weakref
 from typing import (
     TYPE_CHECKING,
     Any,
     Awaitable,
-    Dict,
-    Generic,
-    Literal,
     Optional,
-    Protocol,
-    TypeVar,
     Union,
-    cast,
-    overload,
 )
 
 from agentlightning.adapter import TraceAdapter
@@ -168,239 +160,3 @@ class BaseAlgorithm:
             The AgentLightningClient instance associated with this algorithm.
         """
         raise NotImplementedError("Subclasses must implement get_client().")
-
-
-class FastAlgorithm(BaseAlgorithm):
-    """Algorithm that can run fast and qualify for dev mode.
-
-    Fast algorithms enable agent developers to quickly iterate on agent development
-    without waiting for a long training to complete.
-    """
-
-
-# Algorithm function signature types
-# We've missed a lot of combinations here.
-# Let's add them in future.
-
-
-class AlgorithmFuncSyncFull(Protocol):
-    def __call__(
-        self,
-        *,
-        store: LightningStore,
-        train_dataset: Optional[Dataset[Any]],
-        val_dataset: Optional[Dataset[Any]],
-        llm_proxy: Optional[LLMProxy],
-        adapter: Optional[TraceAdapter[Any]],
-        initial_resources: Optional[NamedResources],
-    ) -> None: ...
-
-
-class AlgorithmFuncSyncOnlyStore(Protocol):
-    def __call__(self, *, store: LightningStore) -> None: ...
-
-
-class AlgorithmFuncSyncOnlyDataset(Protocol):
-    def __call__(self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None: ...
-
-
-class AlgorithmFuncAsyncFull(Protocol):
-    def __call__(
-        self,
-        *,
-        store: LightningStore,
-        train_dataset: Optional[Dataset[Any]],
-        val_dataset: Optional[Dataset[Any]],
-        llm_proxy: Optional[LLMProxy],
-        adapter: Optional[TraceAdapter[Any]],
-        initial_resources: Optional[NamedResources],
-    ) -> Awaitable[None]: ...
-
-
-class AlgorithmFuncAsyncOnlyStore(Protocol):
-    def __call__(self, *, store: LightningStore) -> Awaitable[None]: ...
-
-
-class AlgorithmFuncAsyncOnlyDataset(Protocol):
-    def __call__(
-        self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]
-    ) -> Awaitable[None]: ...
-
-
-AlgorithmFuncAsync = Union[AlgorithmFuncAsyncOnlyStore, AlgorithmFuncAsyncOnlyDataset, AlgorithmFuncAsyncFull]
-
-AlgorithmFuncSync = Union[AlgorithmFuncSyncOnlyStore, AlgorithmFuncSyncOnlyDataset, AlgorithmFuncSyncFull]
-
-
-class AlgorithmFuncSyncFallback(Protocol):
-    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
-
-
-class AlgorithmFuncAsyncFallback(Protocol):
-    def __call__(self, *args: Any, **kwargs: Any) -> Awaitable[Any]: ...
-
-
-AlgorithmFuncSyncLike = Union[AlgorithmFuncSync, AlgorithmFuncSyncFallback]
-AlgorithmFuncAsyncLike = Union[AlgorithmFuncAsync, AlgorithmFuncAsyncFallback]
-
-AlgorithmFunc = Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]
-
-
-AsyncFlag = Literal[True, False]
-AF = TypeVar("AF", bound=AsyncFlag)
-
-
-class FunctionalAlgorithm(BaseAlgorithm, Generic[AF]):
-    """A BaseAlgorithm that wraps a function-based algorithm implementation.
-
-    This class allows users to define algorithm behavior using a simple function
-    that takes train_dataset and val_dataset parameters, rather than implementing
-    a full BaseAlgorithm subclass.
-    """
-
-    @overload
-    def __init__(self: "FunctionalAlgorithm[Literal[False]]", algorithm_func: AlgorithmFuncSyncLike) -> None: ...
-
-    @overload
-    def __init__(self: "FunctionalAlgorithm[Literal[True]]", algorithm_func: AlgorithmFuncAsyncLike) -> None: ...
-
-    def __init__(self, algorithm_func: Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]) -> None:
-        """
-        Initialize the FunctionalAlgorithm with an algorithm function.
-
-        Args:
-            algorithm_func: A function that defines the algorithm's behavior.
-                           Can be sync or async with signature:
-                           (train_dataset, val_dataset) -> None
-        """
-        super().__init__()
-        self._algorithm_func = algorithm_func
-        self._sig = inspect.signature(algorithm_func)
-        self._is_async = inspect.iscoroutinefunction(algorithm_func)
-
-        # Copy function metadata to preserve type hints and other attributes
-        functools.update_wrapper(self, algorithm_func)  # type: ignore
-
-    def is_async(self) -> bool:
-        return self._is_async
-
-    @overload
-    def run(
-        self: "FunctionalAlgorithm[Literal[False]]",
-        train_dataset: Optional[Dataset[Any]] = None,
-        val_dataset: Optional[Dataset[Any]] = None,
-    ) -> None: ...
-
-    @overload
-    def run(
-        self: "FunctionalAlgorithm[Literal[True]]",
-        train_dataset: Optional[Dataset[Any]] = None,
-        val_dataset: Optional[Dataset[Any]] = None,
-    ) -> Awaitable[None]: ...
-
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        return self._algorithm_func(*args, **kwargs)  # type: ignore
-
-    def run(
-        self,
-        train_dataset: Optional[Dataset[Any]] = None,
-        val_dataset: Optional[Dataset[Any]] = None,
-    ) -> Union[None, Awaitable[None]]:
-        """Execute the algorithm using the wrapped function.
-
-        Args:
-            train_dataset: The dataset to train on.
-            val_dataset: The dataset to validate on.
-
-        Returns:
-            None or Awaitable[None] if the function is async.
-        """
-        kwargs: Dict[str, Any] = {}
-        if "store" in self._sig.parameters:
-            kwargs["store"] = self.get_store()
-        if "adapter" in self._sig.parameters:
-            kwargs["adapter"] = self.get_adapter()
-        if "llm_proxy" in self._sig.parameters:
-            kwargs["llm_proxy"] = self.get_llm_proxy()
-        if "initial_resources" in self._sig.parameters:
-            kwargs["initial_resources"] = self.get_initial_resources()
-        if "train_dataset" in self._sig.parameters:
-            kwargs["train_dataset"] = train_dataset
-        elif train_dataset is not None:
-            raise TypeError(
-                f"train_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
-            )
-        if "val_dataset" in self._sig.parameters:
-            kwargs["val_dataset"] = val_dataset
-        elif val_dataset is not None:
-            raise TypeError(
-                f"val_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
-            )
-        # both sync and async functions can be called with the same signature
-        result = self._algorithm_func(**kwargs)  # type: ignore[misc]
-        if self._is_async:
-            return cast(Awaitable[None], result)
-        return None
-
-
-@overload
-def algo(func: AlgorithmFuncAsync) -> FunctionalAlgorithm[Literal[True]]: ...
-
-
-@overload
-def algo(func: AlgorithmFuncAsyncFallback) -> FunctionalAlgorithm[Any]: ...
-
-
-@overload
-def algo(func: AlgorithmFuncSync) -> FunctionalAlgorithm[Literal[False]]: ...
-
-
-@overload
-def algo(func: AlgorithmFuncSyncFallback) -> FunctionalAlgorithm[Any]: ...
-
-
-def algo(
-    func: Union[
-        AlgorithmFuncSync,
-        AlgorithmFuncAsync,
-        AlgorithmFuncSyncFallback,
-        AlgorithmFuncAsyncFallback,
-    ],
-) -> Union[FunctionalAlgorithm[Literal[False]], FunctionalAlgorithm[Literal[True]]]:
-    """Create a BaseAlgorithm from a function.
-
-    This decorator allows you to define an algorithm using a simple function
-    instead of creating a full BaseAlgorithm subclass. The returned FunctionalAlgorithm
-    instance is callable, preserving the original function's behavior.
-
-    Args:
-        func: A function that defines the algorithm's behavior with signature:
-              (train_dataset, val_dataset) -> None
-              Can be sync or async.
-
-    Returns:
-        A callable FunctionalAlgorithm instance that preserves the original function's
-        type hints and behavior while providing all algorithm functionality.
-
-    Example:
-        @algo
-        def my_algorithm(train_dataset, val_dataset):
-            # Algorithm logic here
-            for task in train_dataset:
-                # Process training tasks
-                pass
-
-        @algo
-        async def my_async_algorithm(train_dataset, val_dataset):
-            # Async algorithm logic here
-            async for task in train_dataset:
-                # Process training tasks asynchronously
-                pass
-
-        # Function is still callable with original behavior
-        my_algorithm(train_data, val_data)
-
-        # Algorithm methods are also available
-        my_algorithm.run(train_data, val_data)
-    """
-    return FunctionalAlgorithm(func)

agentlightning/algorithm/decorator.py

@@ -0,0 +1,256 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import functools
+import inspect
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Dict,
+    Generic,
+    Literal,
+    Optional,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from agentlightning.adapter import TraceAdapter
+from agentlightning.store.base import LightningStore
+from agentlightning.types import Dataset, NamedResources
+
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+
+from .base import BaseAlgorithm
+
+# Algorithm function signature types
+# We've missed a lot of combinations here.
+# Let's add them in future.
+
+
+class AlgorithmFuncSyncFull(Protocol):
+    def __call__(
+        self,
+        *,
+        store: LightningStore,
+        train_dataset: Optional[Dataset[Any]],
+        val_dataset: Optional[Dataset[Any]],
+        llm_proxy: Optional[LLMProxy],
+        adapter: Optional[TraceAdapter[Any]],
+        initial_resources: Optional[NamedResources],
+    ) -> None: ...
+
+
+class AlgorithmFuncSyncOnlyStore(Protocol):
+    def __call__(self, *, store: LightningStore) -> None: ...
+
+
+class AlgorithmFuncSyncOnlyDataset(Protocol):
+    def __call__(self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None: ...
+
+
+class AlgorithmFuncAsyncFull(Protocol):
+    def __call__(
+        self,
+        *,
+        store: LightningStore,
+        train_dataset: Optional[Dataset[Any]],
+        val_dataset: Optional[Dataset[Any]],
+        llm_proxy: Optional[LLMProxy],
+        adapter: Optional[TraceAdapter[Any]],
+        initial_resources: Optional[NamedResources],
+    ) -> Awaitable[None]: ...
+
+
+class AlgorithmFuncAsyncOnlyStore(Protocol):
+    def __call__(self, *, store: LightningStore) -> Awaitable[None]: ...
+
+
+class AlgorithmFuncAsyncOnlyDataset(Protocol):
+    def __call__(
+        self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]
+    ) -> Awaitable[None]: ...
+
+
+AlgorithmFuncAsync = Union[AlgorithmFuncAsyncOnlyStore, AlgorithmFuncAsyncOnlyDataset, AlgorithmFuncAsyncFull]
+
+AlgorithmFuncSync = Union[AlgorithmFuncSyncOnlyStore, AlgorithmFuncSyncOnlyDataset, AlgorithmFuncSyncFull]
+
+
+class AlgorithmFuncSyncFallback(Protocol):
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+class AlgorithmFuncAsyncFallback(Protocol):
+    def __call__(self, *args: Any, **kwargs: Any) -> Awaitable[Any]: ...
+
+
+AlgorithmFuncSyncLike = Union[AlgorithmFuncSync, AlgorithmFuncSyncFallback]
+AlgorithmFuncAsyncLike = Union[AlgorithmFuncAsync, AlgorithmFuncAsyncFallback]
+
+AlgorithmFunc = Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]
+
+
+AsyncFlag = Literal[True, False]
+AF = TypeVar("AF", bound=AsyncFlag)
+
+
+class FunctionalAlgorithm(BaseAlgorithm, Generic[AF]):
+    """A BaseAlgorithm that wraps a function-based algorithm implementation.
+
+    This class allows users to define algorithm behavior using a simple function
+    that takes train_dataset and val_dataset parameters, rather than implementing
+    a full BaseAlgorithm subclass.
+    """
+
+    @overload
+    def __init__(self: "FunctionalAlgorithm[Literal[False]]", algorithm_func: AlgorithmFuncSyncLike) -> None: ...
+
+    @overload
+    def __init__(self: "FunctionalAlgorithm[Literal[True]]", algorithm_func: AlgorithmFuncAsyncLike) -> None: ...
+
+    def __init__(self, algorithm_func: Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]) -> None:
+        """
+        Initialize the FunctionalAlgorithm with an algorithm function.
+
+        Args:
+            algorithm_func: A function that defines the algorithm's behavior.
+                           Can be sync or async with signature:
+                           (train_dataset, val_dataset) -> None
+        """
+        super().__init__()
+        self._algorithm_func = algorithm_func
+        self._sig = inspect.signature(algorithm_func)
+        self._is_async = inspect.iscoroutinefunction(algorithm_func)
+
+        # Copy function metadata to preserve type hints and other attributes
+        functools.update_wrapper(self, algorithm_func)  # type: ignore
+
+    def is_async(self) -> bool:
+        return self._is_async
+
+    @overload
+    def run(
+        self: "FunctionalAlgorithm[Literal[False]]",
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> None: ...
+
+    @overload
+    def run(
+        self: "FunctionalAlgorithm[Literal[True]]",
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> Awaitable[None]: ...
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self._algorithm_func(*args, **kwargs)  # type: ignore
+
+    def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> Union[None, Awaitable[None]]:
+        """Execute the algorithm using the wrapped function.
+
+        Args:
+            train_dataset: The dataset to train on.
+            val_dataset: The dataset to validate on.
+
+        Returns:
+            None or Awaitable[None] if the function is async.
+        """
+        kwargs: Dict[str, Any] = {}
+        if "store" in self._sig.parameters:
+            kwargs["store"] = self.get_store()
+        if "adapter" in self._sig.parameters:
+            kwargs["adapter"] = self.get_adapter()
+        if "llm_proxy" in self._sig.parameters:
+            kwargs["llm_proxy"] = self.get_llm_proxy()
+        if "initial_resources" in self._sig.parameters:
+            kwargs["initial_resources"] = self.get_initial_resources()
+        if "train_dataset" in self._sig.parameters:
+            kwargs["train_dataset"] = train_dataset
+        elif train_dataset is not None:
+            raise TypeError(
+                f"train_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
+            )
+        if "val_dataset" in self._sig.parameters:
+            kwargs["val_dataset"] = val_dataset
+        elif val_dataset is not None:
+            raise TypeError(
+                f"val_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
+            )
+        # both sync and async functions can be called with the same signature
+        result = self._algorithm_func(**kwargs)  # type: ignore[misc]
+        if self._is_async:
+            return cast(Awaitable[None], result)
+        return None
+
+
+@overload
+def algo(func: AlgorithmFuncAsync) -> FunctionalAlgorithm[Literal[True]]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncAsyncFallback) -> FunctionalAlgorithm[Any]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncSync) -> FunctionalAlgorithm[Literal[False]]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncSyncFallback) -> FunctionalAlgorithm[Any]: ...
+
+
+def algo(
+    func: Union[
+        AlgorithmFuncSync,
+        AlgorithmFuncAsync,
+        AlgorithmFuncSyncFallback,
+        AlgorithmFuncAsyncFallback,
+    ],
+) -> Union[FunctionalAlgorithm[Literal[False]], FunctionalAlgorithm[Literal[True]]]:
+    """Create a BaseAlgorithm from a function.
+
+    This decorator allows you to define an algorithm using a simple function
+    instead of creating a full BaseAlgorithm subclass. The returned FunctionalAlgorithm
+    instance is callable, preserving the original function's behavior.
+
+    Args:
+        func: A function that defines the algorithm's behavior with signature:
+              (train_dataset, val_dataset) -> None
+              Can be sync or async.
+
+    Returns:
+        A callable FunctionalAlgorithm instance that preserves the original function's
+        type hints and behavior while providing all algorithm functionality.
+
+    Example:
+        @algo
+        def my_algorithm(train_dataset, val_dataset):
+            # Algorithm logic here
+            for task in train_dataset:
+                # Process training tasks
+                pass
+
+        @algo
+        async def my_async_algorithm(train_dataset, val_dataset):
+            # Async algorithm logic here
+            async for task in train_dataset:
+                # Process training tasks asynchronously
+                pass
+
+        # Function is still callable with original behavior
+        my_algorithm(train_data, val_data)
+
+        # Algorithm methods are also available
+        my_algorithm.run(train_data, val_data)
+    """
+    return FunctionalAlgorithm(func)

agentlightning/algorithm/fast.py

@@ -5,21 +5,31 @@ from __future__ import annotations
 import asyncio
 import logging
 from datetime import datetime
-from typing import Any, List, Optional
+from typing import Any, List, Literal, Optional
 
 from agentlightning.llm_proxy import ModelConfig
-from agentlightning.types import Dataset, RolloutStatus, RolloutV2
+from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span
 
-from .base import FastAlgorithm
+from .base import BaseAlgorithm
 
 logger = logging.getLogger(__name__)
 
+__all__ = ["FastAlgorithm", "Baseline"]
+
+
+class FastAlgorithm(BaseAlgorithm):
+    """Algorithm that can run fast and qualify for dev mode.
+
+    Fast algorithms enable agent developers to quickly iterate on agent development
+    without waiting for a long training to complete.
+    """
+
 
 def _timestamp_to_iso_str(timestamp: float) -> str:
     return datetime.fromtimestamp(timestamp).isoformat()
 
 
-class MockAlgorithm(FastAlgorithm):
+class Baseline(FastAlgorithm):
     """A dummy implementation of algorithm interface that puts all dataset into the queue, and waits for all rollouts to complete.
 
     Logs all collected spans and rewards.
@@ -42,18 +52,39 @@ class MockAlgorithm(FastAlgorithm):
         train_split: float = 0.5,
         polling_interval: float = 5.0,
         max_queue_length: int = 4,
+        span_verbosity: Literal["keys", "key_values", "none"] = "keys",
     ) -> None:
         super().__init__()
         self.n_epochs = n_epochs
         self.train_split = train_split
         self.polling_interval = polling_interval
         self.max_queue_length = max_queue_length
+        self.span_verbosity = span_verbosity
         if not (0.0 < self.train_split < 1.0):
             raise ValueError("train_split must be between 0 and 1.")
 
         self._finished_rollout_count = 0
 
-    async def _handle_rollout_finish(self, rollout: RolloutV2) -> None:
+    def _span_to_string(self, rollout_id: str, attempt: Attempt, span: Span) -> str:
+        if self.span_verbosity == "none":
+            return ""
+
+        prefix_msg = f"[Rollout {rollout_id} | Attempt {attempt.attempt_id} | Span {span.span_id}] #{span.sequence_id} ({span.name}) "
+        elapsed = f"{span.end_time - span.start_time:.2f}" if span.start_time and span.end_time else "unknown"
+
+        msg = (
+            prefix_msg
+            + f"From {_timestamp_to_iso_str(span.start_time) if span.start_time else 'unknown'}, "
+            + f"to {_timestamp_to_iso_str(span.end_time) if span.end_time else 'unknown'}, "
+            + f"{elapsed} seconds. "
+        )
+        if self.span_verbosity == "key_values":
+            msg += f"Attributes: {span.attributes}"
+        else:
+            msg += f"Attribute keys: {list(span.attributes.keys())}"
+        return msg
+
+    async def _handle_rollout_finish(self, rollout: Rollout) -> None:
         store = self.get_store()
 
         rollout_id = rollout.rollout_id
@@ -70,14 +101,8 @@ class MockAlgorithm(FastAlgorithm):
             )
             spans = await store.query_spans(rollout_id=rollout_id)
             for span in spans:
-                prefix_msg = f"[Rollout {rollout_id} | Attempt {attempt.attempt_id} | Span {span.span_id}] #{span.sequence_id} ({span.name}) "
-                elapsed = f"{span.end_time - span.start_time:.2f}" if span.start_time and span.end_time else "unknown"
-                logger.info(
-                    prefix_msg
-                    + f"From {_timestamp_to_iso_str(span.start_time) if span.start_time else 'unknown'}, "
-                    + f"to {_timestamp_to_iso_str(span.end_time) if span.end_time else 'unknown'}, "
-                    + f"{elapsed} seconds. Attributes: {span.attributes}"
-                )
+                if self.span_verbosity != "none":
+                    logger.info(self._span_to_string(rollout.rollout_id, attempt, span))
 
         # Attempts to adapt the spans using the adapter if provided
         try:
@@ -178,5 +203,6 @@ class MockAlgorithm(FastAlgorithm):
                     await asyncio.sleep(self.polling_interval)
 
             # Wait for all harvest tasks to complete
+            print(f"Waiting for {len(harvest_tasks)} harvest tasks to complete...")
             if len(harvest_tasks) > 0:
                 await asyncio.gather(*harvest_tasks)

agentlightning/algorithm/verl/interface.py

@@ -12,6 +12,18 @@ from agentlightning.verl.entrypoint import run_ppo  # type: ignore
 
 
 class VERL(BaseAlgorithm):
+    """Algorithm leveraging VERL as the backend framework.
+
+    **Note on Customization:**
+
+    At present, we recommend copying the source code from VERL and modifying it as needed to suit your requirements.
+    Native support for customizing training logic will be provided in future releases.
+
+    Args:
+        config: The VERL configuration, matching what is typically provided when running VERL via the command line.
+            This config will be merged with VERL's base configuration and processed by Hydra.
+    """
+
     def __init__(self, config: dict[str, Any]):
         super().__init__()
 

agentlightning/client.py

@@ -6,12 +6,13 @@ import asyncio
 import logging
 import time
 import urllib.parse
+import warnings
 from typing import Any, Dict, List, Optional, Union
 
 import aiohttp
 import requests
 
-from .types import NamedResources, ResourcesUpdate, Rollout, Task, TaskIfAny, TaskInput
+from .types import NamedResources, ResourcesUpdate, RolloutLegacy, Task, TaskIfAny, TaskInput
 
 logger = logging.getLogger(__name__)
 
@@ -39,6 +40,9 @@ class AgentLightningClient:
             poll_interval: The interval in seconds to wait between polling for new tasks.
             timeout: The timeout in seconds for HTTP requests.
         """
+        warnings.warn(
+            "AgentLightningClient is deprecated. Please use LightningStoreClient instead.", DeprecationWarning
+        )
         self.endpoint = endpoint
         self.task_count = 0
         self.poll_interval = poll_interval
@@ -140,7 +144,7 @@ class AgentLightningClient:
             return resources_update
         return None
 
-    async def post_rollout_async(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
+    async def post_rollout_async(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
         """Posts a completed rollout to the server asynchronously.
 
         Args:
@@ -242,7 +246,7 @@ class AgentLightningClient:
             return resources_update
         return None
 
-    def post_rollout(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
+    def post_rollout(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
         """Posts a completed rollout to the server synchronously.
 
         Args:
@@ -280,6 +284,7 @@ class DevTaskLoader(AgentLightningClient):
             resources: Either NamedResources or ResourcesUpdate object.
             **kwargs: Additional arguments passed to the parent AgentLightningClient.
         """
+        warnings.warn("DevTaskLoader is deprecated. Please use Trainer.dev instead.", DeprecationWarning)
         super().__init__(endpoint="local://", **kwargs)
         self._tasks = tasks.copy()
         if len(self._tasks) == 0:
@@ -298,10 +303,10 @@ class DevTaskLoader(AgentLightningClient):
             self._resources_update = ResourcesUpdate(resources_id="local", resources=resources)
 
         # Store rollouts posted back to the loader for easy debugging of local runs
-        self._rollouts: List[Rollout] = []
+        self._rollouts: List[RolloutLegacy] = []
 
     @property
-    def rollouts(self) -> List[Rollout]:
+    def rollouts(self) -> List[RolloutLegacy]:
         """Return rollouts that have been posted back to the loader."""
         return self._rollouts
 
@@ -347,7 +352,7 @@ class DevTaskLoader(AgentLightningClient):
         logger.debug("DevTaskLoader returning latest resources.")
         return self._resources_update
 
-    def post_rollout(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
+    def post_rollout(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
         logger.debug(f"DevTaskLoader received rollout for task: {rollout.rollout_id}")
         self._rollouts.append(rollout)
         return {"status": "received", "rollout_id": rollout.rollout_id}
@@ -361,7 +366,7 @@ class DevTaskLoader(AgentLightningClient):
     async def get_latest_resources_async(self) -> Optional[ResourcesUpdate]:
         return self.get_latest_resources()
 
-    async def post_rollout_async(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
+    async def post_rollout_async(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
         return self.post_rollout(rollout)
 
     def __repr__(self):

agentlightning/config.py

@@ -31,6 +31,8 @@ CliConfigurable = Any
 
 logger = logging.getLogger(__name__)
 
+__all__ = ["lightning_cli"]
+
 # TypeVars for precise return type hinting with overloads
 _C = TypeVar("_C", bound=CliConfigurable)
 _C1 = TypeVar("_C1", bound=CliConfigurable)
@@ -307,6 +309,9 @@ def lightning_cli(cls1: Type[_C1], cls2: Type[_C2], cls3: Type[_C3], cls4: Type[
 def lightning_cli(*classes: Type[CliConfigurable]) -> Tuple[CliConfigurable, ...]: ...
 
 
+# FIXME: lightning_cli needs to be fixed to comply with the latest trainer implementation.
+
+
 def lightning_cli(*classes: Type[CliConfigurable]) -> CliConfigurable | Tuple[CliConfigurable, ...]:  # type: ignore
     """
     Parses command-line arguments to configure and instantiate provided CliConfigurable classes.

agentlightning/emitter/__init__.py

@@ -4,7 +4,6 @@ from .exception import emit_exception
 from .message import emit_message
 from .object import emit_object
 from .reward import (
-    RewardSpanData,
     emit_reward,
     find_final_reward,
     find_reward_spans,
@@ -20,7 +19,6 @@ __all__ = [
     "is_reward_span",
     "find_reward_spans",
     "find_final_reward",
-    "RewardSpanData",
     "emit_message",
     "emit_object",
     "emit_exception",

agentlightning/execution/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import ExecutionStrategy
+from .client_server import ClientServerExecutionStrategy
+from .events import ExecutionEvent, MultiprocessingEvent, ThreadingEvent
+from .shared_memory import SharedMemoryExecutionStrategy
+
+__all__ = [
+    "ExecutionStrategy",
+    "ClientServerExecutionStrategy",
+    "ExecutionEvent",
+    "ThreadingEvent",
+    "MultiprocessingEvent",
+    "SharedMemoryExecutionStrategy",
+]

agentlightning/execution/base.py

@@ -5,18 +5,18 @@ from typing import Protocol
 
 from agentlightning.store.base import LightningStore
 
-from .events import Event
+from .events import ExecutionEvent
 
 logger = logging.getLogger(__name__)
 
 
 class AlgorithmBundle(Protocol):
-    async def __call__(self, store: LightningStore, event: Event) -> None:
+    async def __call__(self, store: LightningStore, event: ExecutionEvent) -> None:
         """Initalization and execution logic."""
 
 
 class RunnerBundle(Protocol):
-    async def __call__(self, store: LightningStore, worker_id: int, event: Event) -> None:
+    async def __call__(self, store: LightningStore, worker_id: int, event: ExecutionEvent) -> None:
         """Initalization and execution logic."""
 
 

agentlightning/execution/client_server.py

@@ -13,7 +13,7 @@ from agentlightning.store.base import LightningStore
 from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
 
 from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
-from .events import Event, MultiprocessingEvent
+from .events import ExecutionEvent, MultiprocessingEvent
 
 logger = logging.getLogger(__name__)
 
@@ -127,7 +127,9 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
                 raise ValueError("main_process='runner' requires n_runners to be 1")
         self.main_process = main_process
 
-    async def _execute_algorithm(self, algorithm: AlgorithmBundle, store: LightningStore, stop_evt: Event) -> None:
+    async def _execute_algorithm(
+        self, algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent
+    ) -> None:
         logger.info("Starting LightningStore server on %s:%s", self.server_host, self.server_port)
         server_store = LightningStoreServer(store, host=self.server_host, port=self.server_port)
         server_started = False
@@ -155,7 +157,7 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
                 else:
                     logger.debug("LightningStore server shutdown completed")
 
-    async def _execute_runner(self, runner: RunnerBundle, worker_id: int, stop_evt: Event) -> None:
+    async def _execute_runner(self, runner: RunnerBundle, worker_id: int, stop_evt: ExecutionEvent) -> None:
         client_store = LightningStoreClient(f"http://{self.server_host}:{self.server_port}")
         try:
             logger.debug("Runner %s connecting to server at %s:%s", worker_id, self.server_host, self.server_port)
@@ -180,14 +182,14 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
     def _spawn_runners(
         self,
         runner: RunnerBundle,
-        stop_evt: Event,
+        stop_evt: ExecutionEvent,
         *,
         ctx: BaseContext,
     ) -> list[multiprocessing.Process]:
         """Used when `role == "runner"` or `role == "both"` and `n_runners > 1`."""
         processes: list[multiprocessing.Process] = []
 
-        def _runner_sync(runner: RunnerBundle, worker_id: int, stop_evt: Event) -> None:
+        def _runner_sync(runner: RunnerBundle, worker_id: int, stop_evt: ExecutionEvent) -> None:
             # Runners are executed in child processes; each process owns its own
             # event loop to keep the asyncio scheduler isolated.
             asyncio.run(self._execute_runner(runner, worker_id, stop_evt))
@@ -207,13 +209,13 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
         self,
         algorithm: AlgorithmBundle,
         store: LightningStore,
-        stop_evt: Event,
+        stop_evt: ExecutionEvent,
         *,
         ctx: BaseContext,
     ) -> multiprocessing.Process:
         """Used when `main_process == "runner"`."""
 
-        def _algorithm_sync(algorithm: AlgorithmBundle, store: LightningStore, stop_evt: Event) -> None:
+        def _algorithm_sync(algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent) -> None:
             asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
 
         process = cast(
@@ -257,7 +259,7 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
     def _shutdown_processes(
         self,
         processes: list[multiprocessing.Process],
-        stop_evt: Event,
+        stop_evt: ExecutionEvent,
     ) -> None:
         """4-step escalation shutdown of ``processes``."""
         if not processes:

agentlightning/execution/events.py

@@ -6,7 +6,7 @@ from multiprocessing.context import BaseContext
 from typing import Optional, Protocol
 
 
-class Event(Protocol):
+class ExecutionEvent(Protocol):
     """
     A minimal protocol similar to threading.Event.
 

agentlightning/execution/shared_memory.py

@@ -11,7 +11,7 @@ from agentlightning.store.base import LightningStore
 from agentlightning.store.threading import LightningStoreThreaded
 
 from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
-from .events import Event, ThreadingEvent
+from .events import ExecutionEvent, ThreadingEvent
 
 logger = logging.getLogger(__name__)
 
@@ -54,7 +54,7 @@ class SharedMemoryExecutionStrategy(ExecutionStrategy):
         self.graceful_delay = graceful_delay
         self.poll_interval = poll_interval
 
-    async def _run_until_completed_or_canceled(self, coro: Awaitable[Any], stop_evt: Event) -> Any:
+    async def _run_until_completed_or_canceled(self, coro: Awaitable[Any], stop_evt: ExecutionEvent) -> Any:
         """Run `coro` until it finishes or a cooperative stop is requested.
 
         Control flow:
@@ -149,7 +149,7 @@ class SharedMemoryExecutionStrategy(ExecutionStrategy):
         self,
         algorithm: AlgorithmBundle,
         store: LightningStore,
-        stop_evt: Event,
+        stop_evt: ExecutionEvent,
         thread_exceptions: Optional[SimpleQueue[BaseException]],
     ) -> None:
         try:
@@ -168,7 +168,7 @@ class SharedMemoryExecutionStrategy(ExecutionStrategy):
         runner: RunnerBundle,
         store: LightningStore,
         worker_id: int,
-        stop_evt: Event,
+        stop_evt: ExecutionEvent,
         thread_exceptions: Optional[SimpleQueue[BaseException]],
     ) -> None:
         try:

agentlightning/instrumentation/__init__.py

@@ -40,6 +40,7 @@ except ImportError:
 
 
 def instrument_all():
+    """Instrument all the instrumentation libraries."""
     if AGENTOPS_INSTALLED:
         from .agentops import instrument_agentops
 
@@ -70,6 +71,7 @@ def instrument_all():
 
 
 def uninstrument_all():
+    """Uninstrument all the instrumentation libraries."""
     if AGENTOPS_INSTALLED:
         try:
             from .agentops import uninstrument_agentops

agentlightning/instrumentation/agentops.py

@@ -14,6 +14,13 @@ import setproctitle
 
 logger = logging.getLogger(__name__)
 
+__all__ = [
+    "instrument_agentops",
+    "uninstrument_agentops",
+    "agentops_local_server",
+    "AgentOpsServerManager",
+]
+
 # Module-level storage for originals
 _original_handle_chat_attributes: Callable[..., Any] | None = None
 _original_handle_response: Callable[..., Any] | None = None
@@ -156,6 +163,7 @@ def instrument_agentops():
 
 
 def uninstrument_agentops():
+    """Uninstrument agentops to stop capturing token IDs."""
     try:
         _unpatch_new_agentops()
     except Exception:
@@ -197,6 +205,8 @@ def _run_server(**kwargs: Any):  # type: ignore
 
 
 class AgentOpsServerManager:
+    """Manages a AgentOps local server to bypass the online service of AgentOps."""
+
     def __init__(self, daemon: bool = True, port: int | None = None):
         self.server_process: multiprocessing.Process | None = None
         self.server_port = port

agentlightning/instrumentation/agentops_langchain.py

@@ -8,6 +8,11 @@ from agentops.integration.callbacks.langchain import LangchainCallbackHandler
 original_on_chain_start = LangchainCallbackHandler.on_chain_start
 langgraph_entry = None
 
+__all__ = [
+    "instrument_agentops_langchain",
+    "uninstrument_agentops_langchain",
+]
+
 
 def on_chain_start(self: Any, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
     if "name" in kwargs:
@@ -25,12 +30,14 @@ def on_chain_start(self: Any, serialized: Dict[str, Any], inputs: Dict[str, Any]
 
 
 def instrument_agentops_langchain():
+    """Bypass AgentOp's native support for Langchain."""
     global langgraph_entry
     langgraph_entry = instrumentation.AGENTIC_LIBRARIES.pop("langgraph", None)
     LangchainCallbackHandler.on_chain_start = on_chain_start
 
 
 def uninstrument_agentops_langchain():
+    """Restore AgentOp's native support for Langchain."""
     global langgraph_entry
     if langgraph_entry is not None:
         instrumentation.AGENTIC_LIBRARIES["langgraph"] = langgraph_entry

agentlightning/instrumentation/litellm.py

@@ -1,12 +1,20 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+"""LiteLLM instrumentations.
+
+It's unclear whether or not this file is useful.
+It seems that LiteLLM owns its own telemetry from their own entrance
+https://docs.litellm.ai/docs/observability/agentops_integration
+"""
+
 from typing import Any, Optional
 
 from litellm.integrations.opentelemetry import OpenTelemetry
 
-# It's unclear whether or not this file is useful
-# It seems that LiteLLM owns its own telemetry from their own entrance
-# https://docs.litellm.ai/docs/observability/agentops_integration
+__all__ = [
+    "instrument_litellm",
+    "uninstrument_litellm",
+]
 
 original_set_attributes = OpenTelemetry.set_attributes  # type: ignore
 
@@ -21,8 +29,10 @@ def patched_set_attributes(self: Any, span: Any, kwargs: Any, response_obj: Opti
 
 
 def instrument_litellm():
+    """Instrument litellm to capture token IDs."""
     OpenTelemetry.set_attributes = patched_set_attributes
 
 
 def uninstrument_litellm():
+    """Uninstrument litellm to stop capturing token IDs."""
     OpenTelemetry.set_attributes = original_set_attributes

agentlightning/instrumentation/vllm.py

@@ -9,6 +9,11 @@ import vllm.entrypoints.openai.protocol
 from vllm.entrypoints.openai.protocol import ChatCompletionResponse
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
 
+__all__ = [
+    "instrument_vllm",
+    "uninstrument_vllm",
+]
+
 
 class ChatCompletionResponsePatched(ChatCompletionResponse):
     prompt_token_ids: List[int] | None = None
@@ -59,6 +64,10 @@ async def chat_completion_full_generator(
 
 
 def instrument_vllm():
+    """Instrument vLLM to capture token IDs generated by engine.
+
+    This instrumentation has been merged to upstream vLLM since v0.10.2.
+    """
     if vllm.entrypoints.openai.protocol.ChatCompletionResponse is ChatCompletionResponsePatched:
         warnings.warn("vllm is already instrumented. Skip the instrumentation.")
         return
@@ -68,4 +77,5 @@ def instrument_vllm():
 
 
 def uninstrument_vllm():
+    """Uninstrument vLLM to stop capturing token IDs generated by engine."""
     OpenAIServingChat.chat_completion_full_generator = original_chat_completion_full_generator

agentlightning/litagent/__init__.py

@@ -5,7 +5,6 @@ from .litagent import *
 
 __all__ = [
     "LitAgent",
-    "is_v0_1_rollout_api",
     "llm_rollout",
     "prompt_rollout",
     "rollout",

agentlightning/litagent/decorator.py

@@ -13,8 +13,8 @@ from agentlightning.types import (
     NamedResources,
     PromptTemplate,
     ProxyLLM,
-    RolloutRawResultV2,
-    RolloutV2,
+    Rollout,
+    RolloutRawResult,
 )
 
 from .litagent import LitAgent
@@ -34,19 +34,19 @@ T_contra = TypeVar("T_contra", contravariant=True)
 
 
 class LlmRolloutFuncSync2(Protocol[T_contra]):
-    def __call__(self, task: T_contra, llm: LLM) -> RolloutRawResultV2: ...
+    def __call__(self, task: T_contra, llm: LLM) -> RolloutRawResult: ...
 
 
 class LlmRolloutFuncSync3(Protocol[T_contra]):
-    def __call__(self, task: T_contra, llm: LLM, rollout: RolloutV2) -> RolloutRawResultV2: ...
+    def __call__(self, task: T_contra, llm: LLM, rollout: Rollout) -> RolloutRawResult: ...
 
 
 class LlmRolloutFuncAsync2(Protocol[T_contra]):
-    def __call__(self, task: T_contra, llm: LLM) -> Awaitable[RolloutRawResultV2]: ...
+    def __call__(self, task: T_contra, llm: LLM) -> Awaitable[RolloutRawResult]: ...
 
 
 class LlmRolloutFuncAsync3(Protocol[T_contra]):
-    def __call__(self, task: T_contra, llm: LLM, rollout: RolloutV2) -> Awaitable[RolloutRawResultV2]: ...
+    def __call__(self, task: T_contra, llm: LLM, rollout: Rollout) -> Awaitable[RolloutRawResult]: ...
 
 
 LlmRolloutFunc = Union[
@@ -58,21 +58,21 @@ LlmRolloutFunc = Union[
 
 
 class PromptRolloutFuncSync2(Protocol[T_contra]):
-    def __call__(self, task: T_contra, prompt_template: PromptTemplate) -> RolloutRawResultV2: ...
+    def __call__(self, task: T_contra, prompt_template: PromptTemplate) -> RolloutRawResult: ...
 
 
 class PromptRolloutFuncAsync2(Protocol[T_contra]):
-    def __call__(self, task: T_contra, prompt_template: PromptTemplate) -> Awaitable[RolloutRawResultV2]: ...
+    def __call__(self, task: T_contra, prompt_template: PromptTemplate) -> Awaitable[RolloutRawResult]: ...
 
 
 class PromptRolloutFuncSync3(Protocol[T_contra]):
-    def __call__(self, task: T_contra, prompt_template: PromptTemplate, rollout: RolloutV2) -> RolloutRawResultV2: ...
+    def __call__(self, task: T_contra, prompt_template: PromptTemplate, rollout: Rollout) -> RolloutRawResult: ...
 
 
 class PromptRolloutFuncAsync3(Protocol[T_contra]):
     def __call__(
-        self, task: T_contra, prompt_template: PromptTemplate, rollout: RolloutV2
-    ) -> Awaitable[RolloutRawResultV2]: ...
+        self, task: T_contra, prompt_template: PromptTemplate, rollout: Rollout
+    ) -> Awaitable[RolloutRawResult]: ...
 
 
 PromptRolloutFunc = Union[
@@ -86,7 +86,7 @@ PromptRolloutFunc = Union[
 class FunctionalLitAgentFunc(Protocol[T_contra]):
     def __call__(
         self, task: T_contra, *args: Any, **kwargs: Any
-    ) -> Union[RolloutRawResultV2, Awaitable[RolloutRawResultV2]]: ...
+    ) -> Union[RolloutRawResult, Awaitable[RolloutRawResult]]: ...
 
 
 class FunctionalLitAgent(LitAgent[T]):
@@ -134,7 +134,7 @@ class FunctionalLitAgent(LitAgent[T]):
     def is_async(self) -> bool:
         return self._is_async
 
-    def rollout(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
+    def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
         """Execute a synchronous rollout using the wrapped function.
 
         Args:
@@ -151,7 +151,7 @@ class FunctionalLitAgent(LitAgent[T]):
         kwargs = self._get_kwargs(resources, rollout)
         return self._rollout_func(task, **kwargs)  # type: ignore
 
-    async def rollout_async(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
+    async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
         """Execute an asynchronous rollout using the wrapped function.
 
         Args:
@@ -168,7 +168,7 @@ class FunctionalLitAgent(LitAgent[T]):
         kwargs = self._get_kwargs(resources, rollout)
         return await self._rollout_func(task, **kwargs)  # type: ignore
 
-    def _get_kwargs(self, resources: NamedResources, rollout: RolloutV2) -> Dict[str, Any]:
+    def _get_kwargs(self, resources: NamedResources, rollout: Rollout) -> Dict[str, Any]:
         """Extract the kwargs needed for the rollout function based on its signature.
 
         Dynamically builds the kwargs dictionary by inspecting the function signature and
@@ -193,7 +193,7 @@ class FunctionalLitAgent(LitAgent[T]):
 
         return kwargs
 
-    def _get_llm_resource(self, resources: NamedResources, rollout: RolloutV2) -> LLM:
+    def _get_llm_resource(self, resources: NamedResources, rollout: Rollout) -> LLM:
         """Extract the first LLM resource from the resources dictionary.
 
         Strip the ProxyLLM resource into a LLM resource if needed.
@@ -224,7 +224,7 @@ class FunctionalLitAgent(LitAgent[T]):
 
         return resource_found
 
-    def _get_prompt_template_resource(self, resources: NamedResources, rollout: RolloutV2) -> PromptTemplate:
+    def _get_prompt_template_resource(self, resources: NamedResources, rollout: Rollout) -> PromptTemplate:
         """Extract the first PromptTemplate resource from the resources dictionary.
 
         Args:
@@ -252,7 +252,7 @@ class FunctionalLitAgent(LitAgent[T]):
 
         return resource_found
 
-    def _strip_proxy_helper(self, proxy_llm: LLM, rollout: RolloutV2) -> LLM:
+    def _strip_proxy_helper(self, proxy_llm: LLM, rollout: Rollout) -> LLM:
         """Strip the ProxyLLM resource into a concrete LLM resource.
 
         This method resolves ProxyLLM instances to their concrete LLM implementation
@@ -274,7 +274,7 @@ class FunctionalLitAgent(LitAgent[T]):
             # Not a ProxyLLM, nothing to strip here.
             return proxy_llm
 
-        # Rollout is still a RolloutV2 here because API is not stabilized yet.
+        # Rollout is still a Rollout here because API is not stabilized yet.
         # In practice, it must be an AttemptedRollout.
         if not isinstance(rollout, AttemptedRollout):
             raise ValueError("Rollout is not an AttemptedRollout.")

agentlightning/litagent/litagent.py

@@ -8,7 +8,7 @@ import warnings
 import weakref
 from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar
 
-from agentlightning.types import NamedResources, RolloutRawResultV2, RolloutV2, Task
+from agentlightning.types import NamedResources, Rollout, RolloutRawResult, Task
 
 if TYPE_CHECKING:
     from agentlightning.runner import BaseRunner
@@ -22,7 +22,6 @@ T = TypeVar("T")
 
 __all__ = [
     "LitAgent",
-    "is_v0_1_rollout_api",
 ]
 
 
@@ -120,7 +119,10 @@ class LitAgent(Generic[T]):
         Returns:
             The BaseTracer instance associated with this agent.
         """
-        return self.trainer.tracer
+        if hasattr(self.runner, "tracer"):
+            return self.runner.tracer  # type: ignore
+        else:
+            return self.trainer.tracer
 
     @property
     def tracer(self) -> BaseTracer:
@@ -170,7 +172,7 @@ class LitAgent(Generic[T]):
         no-op.
         """
 
-    def on_rollout_end(self, task: Task, rollout: RolloutV2, runner: BaseRunner[T], tracer: BaseTracer) -> None:
+    def on_rollout_end(self, task: Task, rollout: Rollout, runner: BaseRunner[T], tracer: BaseTracer) -> None:
         """Hook called after a rollout completes.
 
         Deprecated in favor of `on_rollout_end` in the `Hook` interface.
@@ -185,7 +187,7 @@ class LitAgent(Generic[T]):
         logging. By default, this is a no-op.
         """
 
-    def rollout(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
+    def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
         """Main entry point for executing a rollout.
 
         This method determines whether to call the synchronous or
@@ -214,7 +216,7 @@ class LitAgent(Generic[T]):
         """
         raise NotImplementedError("Agents must implement the `rollout` method.")
 
-    async def rollout_async(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
+    async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
         """Asynchronous version of the main rollout method.
 
         This method determines whether to call the synchronous or
@@ -240,7 +242,7 @@ class LitAgent(Generic[T]):
         """
         raise NotImplementedError("Agents must implement the `rollout_async` method for async operations.")
 
-    def training_rollout(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
+    def training_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
         """Defines the agent's behavior for a single training task.
 
         This method should contain the logic for how the agent processes an
@@ -256,7 +258,7 @@ class LitAgent(Generic[T]):
         """
         return self.rollout(task, resources, rollout)
 
-    def validation_rollout(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
+    def validation_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
         """Defines the agent's behavior for a single validation task.
 
         By default, this method redirects to `training_rollout`. Override it
@@ -274,9 +276,7 @@ class LitAgent(Generic[T]):
         """
         return self.rollout(task, resources, rollout)
 
-    async def training_rollout_async(
-        self, task: T, resources: NamedResources, rollout: RolloutV2
-    ) -> RolloutRawResultV2:
+    async def training_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
         """Asynchronous version of `training_rollout`.
 
         This method should be implemented by agents that perform asynchronous
@@ -293,9 +293,7 @@ class LitAgent(Generic[T]):
         """
         return await self.rollout_async(task, resources, rollout)
 
-    async def validation_rollout_async(
-        self, task: T, resources: NamedResources, rollout: RolloutV2
-    ) -> RolloutRawResultV2:
+    async def validation_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
         """Asynchronous version of `validation_rollout`.
 
         By default, this method redirects to `training_rollout_async`.

agentlightning/llm_proxy.py

@@ -30,6 +30,10 @@ from .store.base import LightningStore
 
 logger = logging.getLogger(__name__)
 
+__all__ = [
+    "LLMProxy",
+]
+
 
 class ModelConfig(TypedDict):
     """LiteLLM model registration entry.

agentlightning/logging.py

@@ -2,6 +2,8 @@
 
 import logging
 
+__all__ = ["configure_logger"]
+
 
 def configure_logger(level: int = logging.INFO, name: str = "agentlightning") -> logging.Logger:
     logger = logging.getLogger(name)

agentlightning/runner/__init__.py

@@ -1,11 +1,11 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .agent import AgentRunnerV2
+from .agent import LitAgentRunner
 from .base import BaseRunner
-from .legacy import AgentRunner
+from .legacy import LegacyAgentRunner
 
 __all__ = [
     "BaseRunner",
-    "AgentRunner",
-    "AgentRunnerV2",
+    "LegacyAgentRunner",
+    "LitAgentRunner",
 ]

agentlightning/runner/agent.py

@@ -25,14 +25,14 @@ from agentlightning.types import (
     AttemptedRollout,
     Hook,
     NamedResources,
+    Rollout,
     RolloutMode,
-    RolloutRawResultV2,
-    RolloutV2,
+    RolloutRawResult,
     Span,
 )
 
 if TYPE_CHECKING:
-    from agentlightning.execution.events import Event
+    from agentlightning.execution.events import ExecutionEvent
 
 from .base import BaseRunner
 
@@ -41,7 +41,7 @@ T_task = TypeVar("T_task")
 logger = logging.getLogger(__name__)
 
 
-class AgentRunnerV2(BaseRunner[T_task]):
+class LitAgentRunner(BaseRunner[T_task]):
     """Runner implementation for executing agent tasks with distributed support.
 
     This runner manages the complete lifecycle of agent rollout execution,
@@ -139,6 +139,15 @@ class AgentRunnerV2(BaseRunner[T_task]):
 
         self._tracer.teardown_worker(worker_id)
 
+    @property
+    def tracer(self) -> BaseTracer:
+        """Get the tracer instance.
+
+        Returns:
+            The BaseTracer instance used by this runner.
+        """
+        return self._tracer
+
     def get_agent(self) -> LitAgent[T_task]:
         """Get the agent instance.
 
@@ -221,7 +230,7 @@ class AgentRunnerV2(BaseRunner[T_task]):
                 logger.exception(f"{self._log_prefix()} Exception during {hook_type} hook {hook}.")
 
     async def _post_process_rollout_result(
-        self, rollout: AttemptedRollout, raw_result: RolloutRawResultV2
+        self, rollout: AttemptedRollout, raw_result: RolloutRawResult
     ) -> List[ReadableSpan] | List[Span]:
         """Standardizes the agent's return value and report what's needed to report to the store.
 
@@ -296,14 +305,14 @@ class AgentRunnerV2(BaseRunner[T_task]):
 
         return trace_spans
 
-    async def _sleep_until_next_poll(self, event: Optional[Event] = None) -> None:
+    async def _sleep_until_next_poll(self, event: Optional[ExecutionEvent] = None) -> None:
         """Sleep until the next poll interval, with optional event-based interruption.
 
         If an event is provided, the method will check it periodically (every 0.1s)
         and return early if the event is set.
 
         Args:
-            event: Optional Event object that can be used to interrupt the sleep.
+            event: Optional ExecutionEvent object that can be used to interrupt the sleep.
                 If set during the sleep period, the method returns immediately.
         """
         if event is None:
@@ -316,7 +325,7 @@ class AgentRunnerV2(BaseRunner[T_task]):
             if event.is_set():
                 return
 
-    async def _step_impl(self, next_rollout: AttemptedRollout, raise_on_exception: bool = False) -> None:
+    async def _step_impl(self, next_rollout: AttemptedRollout, raise_on_exception: bool = False) -> str:
         """Execute a single rollout implementation.
 
         This is the core method that handles the execution of a single rollout,
@@ -346,7 +355,7 @@ class AgentRunnerV2(BaseRunner[T_task]):
                 raise RuntimeError(f"{self._log_prefix(rollout_id)} Failed to fetch resources")
             else:
                 logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
-                return
+                return rollout_id
 
         trace_spans: List[ReadableSpan] | List[Span] = []
         has_exception: bool = False
@@ -420,7 +429,9 @@ class AgentRunnerV2(BaseRunner[T_task]):
                     f"{self._log_prefix(rollout_id)} Exception during update_attempt. Giving up the update."
                 )
 
-    async def iter(self, *, event: Optional[Event] = None) -> None:
+        return rollout_id
+
+    async def iter(self, *, event: Optional[ExecutionEvent] = None) -> None:
         """Run the runner, continuously iterating over tasks in the store.
 
         This method polls the store for new rollouts and executes them until:
@@ -432,7 +443,7 @@ class AgentRunnerV2(BaseRunner[T_task]):
         propagated, allowing the runner to continue processing subsequent tasks.
 
         Args:
-            event: Optional Event object to signal the runner to stop. The runner
+            event: Optional ExecutionEvent object to signal the runner to stop. The runner
                 will check this event periodically and stop gracefully when set.
         """
         num_tasks_processed = 0
@@ -443,7 +454,7 @@ class AgentRunnerV2(BaseRunner[T_task]):
             self._max_rollouts is None or num_tasks_processed < self._max_rollouts
         ):
             # Retrieve the next rollout
-            next_rollout: Optional[RolloutV2] = None
+            next_rollout: Optional[Rollout] = None
             while not (event is not None and event.is_set()):
                 logger.debug(f"{self._log_prefix()} Try to poll for next rollout.")
                 next_rollout = await store.dequeue_rollout()
@@ -481,8 +492,8 @@ class AgentRunnerV2(BaseRunner[T_task]):
         *,
         resources: Optional[NamedResources] = None,
         mode: Optional[RolloutMode] = None,
-        event: Optional[Event] = None,
-    ) -> None:
+        event: Optional[ExecutionEvent] = None,
+    ) -> Rollout:
         """Execute a single task directly, bypassing the task queue.
 
         This method creates a new rollout for the given input and executes it
@@ -495,9 +506,12 @@ class AgentRunnerV2(BaseRunner[T_task]):
                 If not provided, the latest resources from the store will be used.
             mode: Optional rollout mode ("train" or "validation"). If not provided,
                 the agent's default mode will be used.
-            event: Optional Event object to signal interruption (currently unused
+            event: Optional ExecutionEvent object to signal interruption (currently unused
                 but included for interface consistency).
 
+        Returns:
+            The completed rollout.
+
         Raises:
             Exception: Any exception that occurs during rollout execution will be
                 re-raised to the caller.
@@ -511,4 +525,9 @@ class AgentRunnerV2(BaseRunner[T_task]):
             resources_id = None
 
         attempted_rollout = await self.get_store().start_rollout(input=input, mode=mode, resources_id=resources_id)
-        await self._step_impl(attempted_rollout, raise_on_exception=True)
+        rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True)
+
+        completed_rollout = await store.get_rollout_by_id(rollout_id)
+        if completed_rollout is None:
+            raise RuntimeError(f"{self._log_prefix()} Failed to fetch completed rollout by id after step: {rollout_id}")
+        return completed_rollout

agentlightning/runner/base.py

@@ -13,12 +13,13 @@ import logging
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, Generic, Iterator, Optional, Sequence, TypeVar
 
+from agentlightning.execution.events import ExecutionEvent
 from agentlightning.litagent import LitAgent
 from agentlightning.store.base import LightningStore
-from agentlightning.types import Hook, NamedResources, ParallelWorkerBase, RolloutMode
+from agentlightning.types import Hook, NamedResources, ParallelWorkerBase, Rollout, RolloutMode
 
 if TYPE_CHECKING:
-    from agentlightning.execution.events import Event
+    from agentlightning.execution.events import ExecutionEvent
 
 
 T_task = TypeVar("T_task")
@@ -155,14 +156,14 @@ class BaseRunner(ParallelWorkerBase, Generic[T_task]):
             except Exception:
                 logger.error("Error during runner teardown", exc_info=True)
 
-    async def iter(self, *, event: Optional[Event] = None) -> None:
+    async def iter(self, *, event: Optional[ExecutionEvent] = None) -> None:
         """Run the runner, continuously iterating over tasks in the store.
 
         This method runs in a loop, polling the store for new tasks and executing
         them until interrupted by the event or when no more tasks are available.
 
         Args:
-            event: Optional Event object that can be used to signal the runner
+            event: Optional ExecutionEvent object that can be used to signal the runner
                 to stop gracefully. When set, the runner should finish its current
                 task and exit the iteration loop.
 
@@ -177,8 +178,8 @@ class BaseRunner(ParallelWorkerBase, Generic[T_task]):
         *,
         resources: Optional[NamedResources] = None,
         mode: Optional[RolloutMode] = None,
-        event: Optional[Event] = None,
-    ) -> None:
+        event: Optional[ExecutionEvent] = None,
+    ) -> Rollout:
         """Execute a single task with the given input.
 
         This method provides fine-grained control for executing individual tasks
@@ -190,9 +191,12 @@ class BaseRunner(ParallelWorkerBase, Generic[T_task]):
                 If not provided, the latest resources from the store will be used.
             mode: Optional rollout mode (e.g., "train", "test"). If not provided,
                 the default mode will be used.
-            event: Optional Event object to signal interruption. When set, the
+            event: Optional ExecutionEvent object to signal interruption. When set, the
                 runner may abort the current execution.
 
+        Returns:
+            The completed rollout.
+
         Raises:
             NotImplementedError: Must be implemented by subclasses.
         """

agentlightning/runner/legacy.py

@@ -1,7 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-# type: ignore
-
 import json
 import logging
 import time
@@ -9,22 +7,23 @@ from typing import Any, Dict, List, Optional, cast
 
 from opentelemetry.sdk.trace import ReadableSpan
 
-from agentlightning.adapter import TraceTripletAdapter
+from agentlightning.adapter import TracerTraceToTriplet
 from agentlightning.client import AgentLightningClient
-from agentlightning.litagent import LitAgent, is_v0_1_rollout_api
+from agentlightning.litagent import LitAgent
+from agentlightning.litagent.litagent import is_v0_1_rollout_api
 from agentlightning.tracer.base import BaseTracer
-from agentlightning.types import Rollout, RolloutRawResult, Triplet
+from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Triplet
 
 from .base import BaseRunner
 
 logger = logging.getLogger(__name__)
 
 __all__ = [
-    "AgentRunner",
+    "LegacyAgentRunner",
 ]
 
 
-class AgentRunner(BaseRunner[Any]):
+class LegacyAgentRunner(BaseRunner[Any]):
     """Manages the agent's execution loop and integrates with AgentOps.
 
     This class orchestrates the interaction between the agent (`LitAgent`) and
@@ -45,7 +44,7 @@ class AgentRunner(BaseRunner[Any]):
         agent: LitAgent[Any],
         client: AgentLightningClient,
         tracer: BaseTracer,
-        triplet_exporter: TraceTripletAdapter,
+        triplet_exporter: TracerTraceToTriplet,
         worker_id: Optional[int] = None,
         max_tasks: Optional[int] = None,
     ):
@@ -76,26 +75,26 @@ class AgentRunner(BaseRunner[Any]):
         """Generates a standardized log prefix for the current worker."""
         if self.worker_id is not None:
             if rollout_id:
-                return f"[Worker {self.worker_id} | Rollout {rollout_id}]"
+                return f"[Worker {self.worker_id} | RolloutLegacy {rollout_id}]"
             else:
                 return f"[Worker {self.worker_id}]"
         if rollout_id:
-            return f"[Rollout {rollout_id}]"
+            return f"[RolloutLegacy {rollout_id}]"
         return "[Default Worker]"
 
     def _to_rollout_object(
         self,
-        result: RolloutRawResult,
+        result: RolloutRawResultLegacy,
         rollout_id: str,
-    ) -> Rollout:
-        """Standardizes the agent's return value into a Rollout object.
+    ) -> RolloutLegacy:
+        """Standardizes the agent's return value into a RolloutLegacy object.
 
         Args:
             result: The output from the agent's rollout method.
             rollout_id: The unique identifier for the current task.
 
         Returns:
-            A standardized `Rollout` object for reporting to the server.
+            A standardized `RolloutLegacy` object for reporting to the server.
         """
         trace: Any = None
         final_reward: Optional[float] = None
@@ -116,8 +115,8 @@ class AgentRunner(BaseRunner[Any]):
         # Case 4: result is a list of dict (trace JSON)
         if isinstance(result, list) and all(isinstance(t, dict) for t in result):
             trace = result
-        # Case 5: result is a Rollout object
-        if isinstance(result, Rollout):
+        # Case 5: result is a RolloutLegacy object
+        if isinstance(result, RolloutLegacy):
             final_reward = result.final_reward
             triplets = result.triplets
             trace = result.trace
@@ -129,15 +128,15 @@ class AgentRunner(BaseRunner[Any]):
                 trace = [json.loads(readable_span.to_json()) for readable_span in spans]
                 trace_spans = spans
 
-        # Always extract triplets from the trace using TraceTripletAdapter
+        # Always extract triplets from the trace using TracerTraceToTriplet
         if trace_spans:
-            triplets = self.triplet_exporter(trace_spans)
+            triplets = self.triplet_exporter(trace_spans)  # type: ignore
 
         # If the agent has triplets, use the last one for final reward if not set
         if triplets and triplets[-1].reward is not None and final_reward is None:
             final_reward = triplets[-1].reward
 
-        # Create the Rollout object with standardized fields
+        # Create the RolloutLegacy object with standardized fields
         result_dict: Dict[str, Any] = {
             "rollout_id": rollout_id,
         }
@@ -148,9 +147,9 @@ class AgentRunner(BaseRunner[Any]):
         if trace is not None:
             result_dict["trace"] = trace
 
-        if isinstance(result, Rollout):
+        if isinstance(result, RolloutLegacy):
             return result.model_copy(update=result_dict)
-        return Rollout(**result_dict)
+        return RolloutLegacy(**result_dict)
 
     def run(self) -> bool:  # type: ignore
         """Poll the task and rollout once synchronously."""
@@ -173,7 +172,7 @@ class AgentRunner(BaseRunner[Any]):
             logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
             return False
 
-        rollout_obj = Rollout(rollout_id=task.rollout_id, task=task)  # Default empty rollout
+        rollout_obj = RolloutLegacy(rollout_id=task.rollout_id, task=task)  # Default empty rollout
 
         try:
             try:
@@ -187,14 +186,14 @@ class AgentRunner(BaseRunner[Any]):
                 # Pass the task input, not the whole task object
                 if is_v0_1_rollout_api(rollout_method):
                     result = cast(
-                        RolloutRawResult,
+                        RolloutRawResultLegacy,
                         rollout_method(
                             task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
                         ),
                     )  # type: ignore
                 else:
-                    result = rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)
-                rollout_obj = self._to_rollout_object(result, task.rollout_id)
+                    result = rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)  # type: ignore
+                rollout_obj = self._to_rollout_object(result, task.rollout_id)  # type: ignore
                 end_time = time.time()
                 logger.info(
                     f"{self._log_prefix(rollout_id)} Completed in "
@@ -207,7 +206,7 @@ class AgentRunner(BaseRunner[Any]):
             logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
         finally:
             try:
-                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)
+                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)  # type: ignore
             except Exception:
                 logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
             self.client.post_rollout(rollout_obj)
@@ -250,7 +249,7 @@ class AgentRunner(BaseRunner[Any]):
             logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
             return False
 
-        rollout_obj = Rollout(rollout_id=task.rollout_id, task=task)  # Default empty rollout
+        rollout_obj = RolloutLegacy(rollout_id=task.rollout_id, task=task)  # Default empty rollout
 
         try:
             try:
@@ -266,14 +265,14 @@ class AgentRunner(BaseRunner[Any]):
                 # Pass the task input, not the whole task object
                 if is_v0_1_rollout_api(rollout_method):
                     result = cast(
-                        RolloutRawResult,
+                        RolloutRawResultLegacy,
                         await rollout_method(
                             task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
                         ),
                     )  # type: ignore
                 else:
-                    result = await rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)
-                rollout_obj = self._to_rollout_object(result, task.rollout_id)
+                    result = await rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)  # type: ignore
+                rollout_obj = self._to_rollout_object(result, task.rollout_id)  # type: ignore
                 end_time = time.time()
                 logger.info(
                     f"{self._log_prefix(rollout_id)} Completed in "
@@ -285,7 +284,7 @@ class AgentRunner(BaseRunner[Any]):
             logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
         finally:
             try:
-                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)
+                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)  # type: ignore
             except Exception:
                 logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
             await self.client.post_rollout_async(rollout_obj)

agentlightning/server.py

@@ -9,6 +9,7 @@ import logging
 import threading
 import time
 import uuid
+import warnings
 from contextlib import asynccontextmanager
 from typing import Any, Dict, List, Literal, Optional
 
@@ -19,7 +20,7 @@ from .types import (
     GenericResponse,
     NamedResources,
     ResourcesUpdate,
-    Rollout,
+    RolloutLegacy,
     Task,
     TaskIfAny,
 )
@@ -36,7 +37,7 @@ class ServerDataStore:
     def __init__(self):
         self._task_queue: asyncio.Queue[Task] = asyncio.Queue()
         self._processing_tasks: Dict[str, Task] = {}  # Currently processing tasks
-        self._completed_rollouts: Dict[str, Rollout] = {}
+        self._completed_rollouts: Dict[str, RolloutLegacy] = {}
 
         # Store for versioned resources
         self._resource_versions: Dict[str, NamedResources] = {}
@@ -121,7 +122,7 @@ class ServerDataStore:
             return await self.get_resources_by_id(self._latest_resources_id)
         return None
 
-    async def store_rollout(self, rollout: Rollout):
+    async def store_rollout(self, rollout: RolloutLegacy):
         """
         Safely stores a completed rollout from a client.
         """
@@ -130,14 +131,14 @@ class ServerDataStore:
             self._completed_rollouts[rollout.rollout_id] = rollout
             logger.info(f"Rollout received and stored: {rollout.rollout_id}")
 
-    async def retrieve_rollout(self, rollout_id: str) -> Optional[Rollout]:
+    async def retrieve_rollout(self, rollout_id: str) -> Optional[RolloutLegacy]:
         """
         Safely retrieves a single rollout by its ID, removing it from the store.
         """
         async with self._results_lock:
             return self._completed_rollouts.pop(rollout_id, None)
 
-    async def retrieve_completed_rollouts(self) -> List[Rollout]:
+    async def retrieve_completed_rollouts(self) -> List[RolloutLegacy]:
         """
         Retrieves all completed rollouts and clears the store.
         """
@@ -176,6 +177,9 @@ class AgentLightningServer:
             port: The port to bind the server to.
             task_timeout_seconds: Time in seconds after which a claimed task is considered stale and requeued.
         """
+        warnings.warn(
+            "AgentLightningServer is deprecated. Please use LightningStoreServer instead.", DeprecationWarning
+        )
         self.host = host
         self.port = port
         self.endpoint = f"http://{host}:{port}"
@@ -272,7 +276,7 @@ class AgentLightningServer:
             return resources_update
 
         @self._app.post("/rollout", response_model=GenericResponse)
-        async def post_rollout(payload: Rollout) -> GenericResponse:  # type: ignore
+        async def post_rollout(payload: RolloutLegacy) -> GenericResponse:  # type: ignore
             """Endpoint for clients to report a completed rollout."""
             if not self._store:
                 raise HTTPException(status_code=503, detail="Server not fully initialized.")
@@ -328,7 +332,7 @@ class AgentLightningServer:
         await self._store.update_resources(update)
         return resources_id
 
-    async def get_completed_rollout(self, rollout_id: str) -> Optional[Rollout]:
+    async def get_completed_rollout(self, rollout_id: str) -> Optional[RolloutLegacy]:
         """
         Retrieves a specific completed rollout by its ID.
         """
@@ -336,7 +340,7 @@ class AgentLightningServer:
             raise RuntimeError("Store not initialized. The server may not be running.")
         return await self._store.retrieve_rollout(rollout_id)
 
-    async def poll_completed_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[Rollout]:
+    async def poll_completed_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[RolloutLegacy]:
         """
         Polls for a completed rollout by its ID, waiting up to `timeout` seconds.
         """
@@ -349,7 +353,7 @@ class AgentLightningServer:
                 return None
             await asyncio.sleep(1)
 
-    async def retrieve_completed_rollouts(self) -> List[Rollout]:
+    async def retrieve_completed_rollouts(self) -> List[RolloutLegacy]:
         """
         Retrieves all available completed trajectories and clears the internal store.
         """

agentlightning/store/__init__.py

@@ -1 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
+
+from .base import LightningStore
+from .client_server import LightningStoreClient, LightningStoreServer
+from .memory import InMemoryLightningStore
+from .threading import LightningStoreThreaded
+
+__all__ = [
+    "LightningStore",
+    "LightningStoreClient",
+    "LightningStoreServer",
+    "InMemoryLightningStore",
+    "LightningStoreThreaded",
+]

agentlightning/store/base.py

@@ -12,23 +12,23 @@ from agentlightning.types import (
     AttemptStatus,
     NamedResources,
     ResourcesUpdate,
+    Rollout,
     RolloutConfig,
     RolloutStatus,
-    RolloutV2,
     Span,
     TaskInput,
 )
 
 
-def is_queuing(rollout: RolloutV2) -> bool:
+def is_queuing(rollout: Rollout) -> bool:
     return rollout.status == "queuing" or rollout.status == "requeuing"
 
 
-def is_running(rollout: RolloutV2) -> bool:
+def is_running(rollout: Rollout) -> bool:
     return rollout.status == "preparing" or rollout.status == "running"
 
 
-def is_finished(rollout: RolloutV2) -> bool:
+def is_finished(rollout: Rollout) -> bool:
     return rollout.status == "failed" or rollout.status == "succeeded" or rollout.status == "cancelled"
 
 
@@ -88,7 +88,7 @@ class LightningStore:
         mode: Literal["train", "val", "test"] | None = None,
         resources_id: str | None = None,
         metadata: Dict[str, Any] | None = None,
-    ) -> RolloutV2:
+    ) -> Rollout:
         """
         Adds a new task to the queue with specific metadata and
         returns the rollout object with its unique ID.
@@ -134,7 +134,7 @@ class LightningStore:
 
     async def query_rollouts(
         self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
-    ) -> List[RolloutV2]:
+    ) -> List[Rollout]:
         """
         Query and retrieve rollouts filtered by their status.
         If no status is provided, returns all rollouts.
@@ -148,7 +148,7 @@ class LightningStore:
         """
         raise NotImplementedError()
 
-    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
         """
         Safely retrieves a specific rollout by its ID.
         """
@@ -181,7 +181,7 @@ class LightningStore:
         """
         raise NotImplementedError()
 
-    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
         """
         Wait for specified rollouts to complete with a timeout.
         Returns the completed rollouts, potentially incomplete if timeout is reached.
@@ -219,7 +219,7 @@ class LightningStore:
         status: RolloutStatus | Unset = UNSET,
         config: RolloutConfig | Unset = UNSET,
         metadata: Optional[Dict[str, Any]] | Unset = UNSET,
-    ) -> RolloutV2:
+    ) -> Rollout:
         """
         Update the rollout status and related metadata.
 

agentlightning/store/client_server.py

@@ -24,9 +24,9 @@ from agentlightning.types import (
     AttemptStatus,
     NamedResources,
     ResourcesUpdate,
+    Rollout,
     RolloutConfig,
     RolloutStatus,
-    RolloutV2,
     Span,
     TaskInput,
 )
@@ -284,7 +284,7 @@ class LightningStoreServer(LightningStore):
                 metadata=request.metadata,
             )
 
-        @self.app.post("/enqueue_rollout", response_model=RolloutV2)
+        @self.app.post("/enqueue_rollout", response_model=Rollout)
         async def enqueue_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
             return await self.store.enqueue_rollout(
                 input=request.input,
@@ -301,7 +301,7 @@ class LightningStoreServer(LightningStore):
         async def start_attempt(request: RolloutId):  # pyright: ignore[reportUnusedFunction]
             return await self.store.start_attempt(request.rollout_id)
 
-        @self.app.post("/query_rollouts", response_model=List[RolloutV2])
+        @self.app.post("/query_rollouts", response_model=List[Rollout])
         async def query_rollouts(request: QueryRolloutsRequest):  # pyright: ignore[reportUnusedFunction]
             return await self.store.query_rollouts(status=request.status)
 
@@ -313,7 +313,7 @@ class LightningStoreServer(LightningStore):
         async def get_latest_attempt(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.store.get_latest_attempt(rollout_id)
 
-        @self.app.get("/get_rollout_by_id/{rollout_id}", response_model=Optional[RolloutV2])
+        @self.app.get("/get_rollout_by_id/{rollout_id}", response_model=Optional[Rollout])
         async def get_rollout_by_id(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.store.get_rollout_by_id(rollout_id)
 
@@ -335,13 +335,15 @@ class LightningStoreServer(LightningStore):
 
         @self.app.post("/add_span", response_model=Span)
         async def add_span(span: Span):  # pyright: ignore[reportUnusedFunction]
+            print("!!!!! add_span received")
             return await self.store.add_span(span)
 
         @self.app.get("/get_next_span_sequence_id/{rollout_id}/{attempt_id}", response_model=int)
         async def get_next_span_sequence_id(rollout_id: str, attempt_id: str):  # pyright: ignore[reportUnusedFunction]
+            print("!!!!! get_next_span_sequence_id received")
             return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)
 
-        @self.app.post("/wait_for_rollouts", response_model=List[RolloutV2])
+        @self.app.post("/wait_for_rollouts", response_model=List[Rollout])
         async def wait_for_rollouts(request: WaitForRolloutsRequest):  # pyright: ignore[reportUnusedFunction]
             return await self.store.wait_for_rollouts(rollout_ids=request.rollout_ids, timeout=request.timeout)
 
@@ -351,7 +353,7 @@ class LightningStoreServer(LightningStore):
         ):
             return await self.store.query_spans(rollout_id, attempt_id)
 
-        @self.app.post("/update_rollout", response_model=RolloutV2)
+        @self.app.post("/update_rollout", response_model=Rollout)
         async def update_rollout(request: UpdateRolloutRequest):  # pyright: ignore[reportUnusedFunction]
             return await self.store.update_rollout(
                 rollout_id=request.rollout_id,
@@ -392,7 +394,7 @@ class LightningStoreServer(LightningStore):
         mode: Literal["train", "val", "test"] | None = None,
         resources_id: str | None = None,
         metadata: Dict[str, Any] | None = None,
-    ) -> RolloutV2:
+    ) -> Rollout:
         return await self._backend().enqueue_rollout(input, mode, resources_id, metadata)
 
     async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
@@ -403,7 +405,7 @@ class LightningStoreServer(LightningStore):
 
     async def query_rollouts(
         self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
-    ) -> List[RolloutV2]:
+    ) -> List[Rollout]:
         return await self._backend().query_rollouts(status=status, rollout_ids=rollout_ids)
 
     async def query_attempts(self, rollout_id: str) -> List[Attempt]:
@@ -412,7 +414,7 @@ class LightningStoreServer(LightningStore):
     async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
         return await self._backend().get_latest_attempt(rollout_id)
 
-    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
         return await self._backend().get_rollout_by_id(rollout_id)
 
     async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
@@ -442,7 +444,7 @@ class LightningStoreServer(LightningStore):
     ) -> Span:
         return await self._backend().add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)
 
-    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
         return await self._backend().wait_for_rollouts(rollout_ids=rollout_ids, timeout=timeout)
 
     async def query_spans(
@@ -461,7 +463,7 @@ class LightningStoreServer(LightningStore):
         status: RolloutStatus | Unset = UNSET,
         config: RolloutConfig | Unset = UNSET,
         metadata: Optional[Dict[str, Any]] | Unset = UNSET,
-    ) -> RolloutV2:
+    ) -> Rollout:
         return await self._backend().update_rollout(
             rollout_id=rollout_id,
             input=input,
@@ -491,6 +493,18 @@ class LightningStoreServer(LightningStore):
         )
 
 
+# def _make_trace():
+#     tc = aiohttp.TraceConfig()
+#     # async def log_evt(session, context, params):
+#     #     print(f"[TRACE] {context}: {params}")
+#     tc.on_dns_resolvehost_start.append(lambda *a, **k: print("[TRACE] dns_start", a[-1].host))
+#     tc.on_connection_create_start.append(lambda *a, **k: print("[TRACE] conn_start"))
+#     tc.on_request_start.append(lambda *a, **k: print("[TRACE] req_start", a[-1].method, a[-1].url))
+#     tc.on_request_end.append(lambda *a, **k: print("[TRACE] req_end", a[-1].method, a[-1].url))
+#     tc.on_request_exception.append(lambda *a, **k: print("[TRACE] req_exc", a[-1].method, a[-1].url))
+#     return tc
+
+
 class LightningStoreClient(LightningStore):
     """HTTP client that talks to a remote LightningStoreServer.
 
@@ -541,11 +555,18 @@ class LightningStoreClient(LightningStore):
 
         loop = asyncio.get_running_loop()
         key = id(loop)
+        print("!!!!! _get_session received %s", key)
         with self._lock:
+            print("!!!!! _get_session with lock")
             sess = self._sessions.get(key)
             if sess is None or sess.closed:
+                # connector = aiohttp.TCPConnector(
+                #     limit=64, limit_per_host=16, ttl_dns_cache=300, enable_cleanup_closed=True
+                # )
+                # sess = aiohttp.ClientSession(trace_configs=[_make_trace()])
                 sess = aiohttp.ClientSession()
                 self._sessions[key] = sess
+        print(self._sessions)
         return sess
 
     async def _wait_until_healthy(self, session: aiohttp.ClientSession) -> bool:
@@ -591,6 +612,7 @@ class LightningStoreClient(LightningStore):
         """
         session = await self._get_session()
         url = f"{self.server_address}{path if path.startswith('/') else '/'+path}"
+        print("$$$$$$ session acquired", url)
 
         # attempt 0 is immediate, then follow retry schedule
         attempts = (0.0,) + self._retry_delays
@@ -602,13 +624,16 @@ class LightningStoreClient(LightningStore):
                 await asyncio.sleep(delay)
             try:
                 http_call = getattr(session, method)
-                async with http_call(url, json=json) as resp:
+                print("$$$$$$ http_call", http_call)
+                timeout = aiohttp.ClientTimeout(total=3.5, connect=1.0, sock_connect=1.0, sock_read=2.5)
+                async with http_call(url, json=json, timeout=timeout) as resp:
+                    print("$$$$$ resp", resp)
                     resp.raise_for_status()
                     return await resp.json()
             except aiohttp.ClientResponseError as cre:
                 # Respect app-level 4xx as final (server marks app faults as 400)
                 # 4xx => application issue; do not retry (except 408 which is transient)
-                logger.exception(f"ClientResponseError: {cre.status} {cre.message}")
+                logger.debug(f"ClientResponseError: {cre.status} {cre.message}", exc_info=True)
                 if 400 <= cre.status < 500 and cre.status != 408:
                     raise
                 # 5xx and others will be retried below if they raise
@@ -624,7 +649,7 @@ class LightningStoreClient(LightningStore):
                 asyncio.TimeoutError,
             ) as net_exc:
                 # Network/session issue: probe health before retrying
-                logger.exception(f"Network/session issue: {net_exc}")
+                logger.debug(f"Network/session issue: {net_exc}", exc_info=True)
                 last_exc = net_exc
                 logger.info(f"Network/session issue will be retried. Retrying the request {method}: {path}")
                 if not await self._wait_until_healthy(session):
@@ -675,13 +700,13 @@ class LightningStoreClient(LightningStore):
         mode: Literal["train", "val", "test"] | None = None,
         resources_id: str | None = None,
         metadata: Dict[str, Any] | None = None,
-    ) -> RolloutV2:
+    ) -> Rollout:
         data = await self._request_json(
             "post",
             "/enqueue_rollout",
             json=RolloutRequest(input=input, mode=mode, resources_id=resources_id, metadata=metadata).model_dump(),
         )
-        return RolloutV2.model_validate(data)
+        return Rollout.model_validate(data)
 
     async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
         """
@@ -720,7 +745,7 @@ class LightningStoreClient(LightningStore):
 
     async def query_rollouts(
         self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
-    ) -> List[RolloutV2]:
+    ) -> List[Rollout]:
         data = await self._request_json(
             "post",
             "/query_rollouts",
@@ -729,7 +754,7 @@ class LightningStoreClient(LightningStore):
                 rollout_ids=list(rollout_ids) if rollout_ids else None,
             ).model_dump(),
         )
-        return [RolloutV2.model_validate(item) for item in data]
+        return [Rollout.model_validate(item) for item in data]
 
     async def query_attempts(self, rollout_id: str) -> List[Attempt]:
         data = await self._request_json("get", f"/query_attempts/{rollout_id}")
@@ -756,7 +781,7 @@ class LightningStoreClient(LightningStore):
             logger.error(f"get_latest_attempt failed after all retries for rollout_id={rollout_id}: {e}", exc_info=True)
             return None
 
-    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
         """
         Get a rollout by its ID.
 
@@ -764,7 +789,7 @@ class LightningStoreClient(LightningStore):
             rollout_id: ID of the rollout to retrieve.
 
         Returns:
-            RolloutV2 if found, None if not found or if all retries are exhausted.
+            Rollout if found, None if not found or if all retries are exhausted.
 
         Note:
             This method retries on transient failures (network errors, 5xx status codes).
@@ -772,7 +797,7 @@ class LightningStoreClient(LightningStore):
         """
         try:
             data = await self._request_json("get", f"/get_rollout_by_id/{rollout_id}")
-            return RolloutV2.model_validate(data) if data else None
+            return Rollout.model_validate(data) if data else None
         except Exception as e:
             logger.error(f"get_rollout_by_id failed after all retries for rollout_id={rollout_id}: {e}", exc_info=True)
             return None
@@ -832,6 +857,7 @@ class LightningStoreClient(LightningStore):
             return None
 
     async def add_span(self, span: Span) -> Span:
+        print("$$$$$$ add_span received")
         data = await self._request_json("post", "/add_span", json=span.model_dump(mode="json"))
         return Span.model_validate(data)
 
@@ -848,6 +874,7 @@ class LightningStoreClient(LightningStore):
         sequence_id: int | None = None,
     ) -> Span:
         # unchanged logic, now benefits from retries inside add_span/get_next_span_sequence_id
+        print("$$$$$$ add_otel_span received")
         if sequence_id is None:
             sequence_id = await self.get_next_span_sequence_id(rollout_id, attempt_id)
         span = Span.from_opentelemetry(
@@ -856,10 +883,11 @@ class LightningStoreClient(LightningStore):
             attempt_id=attempt_id,
             sequence_id=sequence_id,
         )
+        print("$$$$$$ span created")
         await self.add_span(span)
         return span
 
-    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
         if timeout is not None and timeout > 0.1:
             raise ValueError(
                 "Timeout must be less than 0.1 seconds in LightningStoreClient to avoid blocking the event loop"
@@ -869,7 +897,7 @@ class LightningStoreClient(LightningStore):
             "/wait_for_rollouts",
             json=WaitForRolloutsRequest(rollout_ids=rollout_ids, timeout=timeout).model_dump(),
         )
-        return [RolloutV2.model_validate(item) for item in data]
+        return [Rollout.model_validate(item) for item in data]
 
     async def query_spans(
         self,
@@ -891,7 +919,7 @@ class LightningStoreClient(LightningStore):
         status: RolloutStatus | Unset = UNSET,
         config: RolloutConfig | Unset = UNSET,
         metadata: Optional[Dict[str, Any]] | Unset = UNSET,
-    ) -> RolloutV2:
+    ) -> Rollout:
         payload: Dict[str, Any] = {"rollout_id": rollout_id}
         if not isinstance(input, Unset):
             payload["input"] = input
@@ -907,7 +935,7 @@ class LightningStoreClient(LightningStore):
             payload["metadata"] = metadata
 
         data = await self._request_json("post", "/update_rollout", json=payload)
-        return RolloutV2.model_validate(data)
+        return Rollout.model_validate(data)
 
     async def update_attempt(
         self,

agentlightning/store/memory.py

@@ -20,9 +20,9 @@ from agentlightning.types import (
     AttemptStatus,
     NamedResources,
     ResourcesUpdate,
+    Rollout,
     RolloutConfig,
     RolloutStatus,
-    RolloutV2,
     Span,
     TaskInput,
 )
@@ -95,8 +95,8 @@ class InMemoryLightningStore(LightningStore):
         self._lock = asyncio.Lock()
 
         # Task queue and rollouts storage
-        self._task_queue: deque[RolloutV2] = deque()
-        self._rollouts: Dict[str, RolloutV2] = {}
+        self._task_queue: deque[Rollout] = deque()
+        self._rollouts: Dict[str, Rollout] = {}
 
         # Resources storage (similar to legacy server.py)
         self._resources: Dict[str, ResourcesUpdate] = {}
@@ -127,7 +127,7 @@ class InMemoryLightningStore(LightningStore):
             rollout_id = _generate_rollout_id()
             current_time = time.time()
 
-            rollout = RolloutV2(
+            rollout = Rollout(
                 rollout_id=rollout_id,
                 input=input,
                 mode=mode,
@@ -162,7 +162,7 @@ class InMemoryLightningStore(LightningStore):
         mode: Literal["train", "val", "test"] | None = None,
         resources_id: str | None = None,
         metadata: Dict[str, Any] | None = None,
-    ) -> RolloutV2:
+    ) -> Rollout:
         """
         Adds a new task to the queue with specific metadata and returns its unique ID.
         """
@@ -170,7 +170,7 @@ class InMemoryLightningStore(LightningStore):
             rollout_id = _generate_rollout_id()
             current_time = time.time()
 
-            rollout = RolloutV2(
+            rollout = Rollout(
                 rollout_id=rollout_id,
                 input=input,
                 mode=mode,
@@ -276,7 +276,7 @@ class InMemoryLightningStore(LightningStore):
     @_healthcheck_wrapper
     async def query_rollouts(
         self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
-    ) -> List[RolloutV2]:
+    ) -> List[Rollout]:
         """
         Query and retrieve rollouts filtered by their status and rollout ids.
         If no status is provided, returns all rollouts.
@@ -297,7 +297,7 @@ class InMemoryLightningStore(LightningStore):
             return rollouts
 
     @_healthcheck_wrapper
-    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
         """
         Safely retrieves a specific rollout by its ID.
         """
@@ -440,14 +440,14 @@ class InMemoryLightningStore(LightningStore):
         return span
 
     @_healthcheck_wrapper
-    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
         """
         Wait for specified rollouts to complete with a timeout.
         Returns the completed rollouts, potentially incomplete if timeout is reached.
 
         This method does not change the state of the store.
         """
-        completed_rollouts: List[RolloutV2] = []
+        completed_rollouts: List[Rollout] = []
 
         async def wait_for_rollout(rollout_id: str):
             # First check if already completed
@@ -522,7 +522,7 @@ class InMemoryLightningStore(LightningStore):
         status: RolloutStatus | Unset = UNSET,
         config: RolloutConfig | Unset = UNSET,
         metadata: Optional[Dict[str, Any]] | Unset = UNSET,
-    ) -> RolloutV2:
+    ) -> Rollout:
         """
         Update the rollout status and related metadata.
         """
@@ -571,7 +571,7 @@ class InMemoryLightningStore(LightningStore):
         status: RolloutStatus | Unset = UNSET,
         config: RolloutConfig | Unset = UNSET,
         metadata: Optional[Dict[str, Any]] | Unset = UNSET,
-    ) -> RolloutV2:
+    ) -> Rollout:
         # No lock inside this one.
         rollout = self._rollouts.get(rollout_id)
         if not rollout:
@@ -614,7 +614,7 @@ class InMemoryLightningStore(LightningStore):
                 )
 
         # Re-validate the rollout to ensure legality
-        RolloutV2.model_validate(rollout.model_dump())
+        Rollout.model_validate(rollout.model_dump())
 
         return rollout
 
@@ -664,7 +664,7 @@ class InMemoryLightningStore(LightningStore):
 
         if attempt == latest_attempt:
 
-            async def _update_status(rollout_id: str, status: RolloutStatus) -> RolloutV2:
+            async def _update_status(rollout_id: str, status: RolloutStatus) -> Rollout:
                 return await self._update_rollout_unlocked(rollout_id, status=status)
 
             # Propagate the status to the rollout
@@ -693,7 +693,7 @@ class InMemoryLightningStore(LightningStore):
             async def _update_attempt_status(rollout_id: str, attempt_id: str, status: AttemptStatus) -> Attempt:
                 return await self._update_attempt_unlocked(rollout_id, attempt_id, status=status)
 
-            async def _update_rollout_status(rollout_id: str, status: RolloutStatus) -> RolloutV2:
+            async def _update_rollout_status(rollout_id: str, status: RolloutStatus) -> Rollout:
                 return await self._update_rollout_unlocked(rollout_id, status=status)
 
             await healthcheck(

agentlightning/store/sqlite.py

@@ -1 +1,3 @@
 # Copyright (c) Microsoft. All rights reserved.
+
+# TODO: Implement this

agentlightning/store/threading.py

@@ -13,9 +13,9 @@ from agentlightning.types import (
     AttemptStatus,
     NamedResources,
     ResourcesUpdate,
+    Rollout,
     RolloutConfig,
     RolloutStatus,
-    RolloutV2,
     Span,
     TaskInput,
 )
@@ -51,7 +51,7 @@ class LightningStoreThreaded(LightningStore):
         mode: Literal["train", "val", "test"] | None = None,
         resources_id: str | None = None,
         metadata: Dict[str, Any] | None = None,
-    ) -> RolloutV2:
+    ) -> Rollout:
         with self._lock:
             return await self.store.enqueue_rollout(input, mode, resources_id, metadata)
 
@@ -68,7 +68,7 @@ class LightningStoreThreaded(LightningStore):
         *,
         status: Optional[Sequence[RolloutStatus]] = None,
         rollout_ids: Optional[Sequence[str]] = None,
-    ) -> List[RolloutV2]:
+    ) -> List[Rollout]:
         with self._lock:
             return await self.store.query_rollouts(status=status, rollout_ids=rollout_ids)
 
@@ -76,7 +76,7 @@ class LightningStoreThreaded(LightningStore):
         with self._lock:
             return await self.store.query_attempts(rollout_id)
 
-    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
         with self._lock:
             return await self.store.get_rollout_by_id(rollout_id)
 
@@ -114,7 +114,7 @@ class LightningStoreThreaded(LightningStore):
         with self._lock:
             return await self.store.add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)
 
-    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
         # This method does not change the state of the store, and it's not thread-safe.
         return await self.store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=timeout)
 
@@ -139,7 +139,7 @@ class LightningStoreThreaded(LightningStore):
         status: RolloutStatus | Unset = UNSET,
         config: RolloutConfig | Unset = UNSET,
         metadata: Optional[Dict[str, Any]] | Unset = UNSET,
-    ) -> RolloutV2:
+    ) -> Rollout:
         with self._lock:
             return await self.store.update_rollout(
                 rollout_id=rollout_id,

agentlightning/store/utils.py

@@ -3,9 +3,9 @@
 import time
 from typing import Awaitable, Callable, List, cast
 
-from agentlightning.types import Attempt, AttemptedRollout, AttemptStatus, RolloutConfig, RolloutStatus, RolloutV2
+from agentlightning.types import Attempt, AttemptedRollout, AttemptStatus, Rollout, RolloutConfig, RolloutStatus
 
-UpdateRolloutStatus = Callable[[str, RolloutStatus], Awaitable[RolloutV2]]
+UpdateRolloutStatus = Callable[[str, RolloutStatus], Awaitable[Rollout]]
 UpdateAttemptStatus = Callable[[str, str, AttemptStatus], Awaitable[Attempt]]
 
 
@@ -13,7 +13,7 @@ async def propagate_status(
     update_rollout_status: UpdateRolloutStatus,  # this should be unlocked
     attempt: Attempt,
     config: RolloutConfig,
-) -> RolloutV2:
+) -> Rollout:
     """
     Propagate the status of an attempt to the rollout.
 

agentlightning/tracer/agentops.py

@@ -197,7 +197,7 @@ class AgentOpsTracer(BaseTracer):
             raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
         return self._lightning_span_processor.spans()
 
-    def get_langchain_callback_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
+    def get_langchain_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
         """
         Get the Langchain callback handler for integrating with Langchain.
 
@@ -221,6 +221,43 @@ class AgentOpsTracer(BaseTracer):
             )
         return LangchainCallbackHandler(api_key=api_key, tags=tags)
 
+    get_langchain_callback_handler = get_langchain_handler  # alias
+
+
+async def heartbeat(name="exporter-loop", period=0.5):
+    import asyncio
+    import time
+
+    last = time.perf_counter()
+    while True:
+        await asyncio.sleep(period)
+        now = time.perf_counter()
+        dt = now - last
+        last = now
+        if dt > period * 4:  # e.g., >2s if period=0.5s
+            print("!!!!!!! [%s] loop stall detected: slept %.3fs (expected %.3fs)" % (name, dt, period))
+
+
+import asyncio
+import logging
+
+# logging.basicConfig(level=logging.DEBUG)
+# asyncio.get_event_loop().set_debug(True)
+import time
+
+
+def debug_dump(loop):
+    while True:
+        try:
+            print("=== Pending tasks ===")
+            for t in asyncio.all_tasks(loop):
+                if not t.done():
+                    print(t, "awaiting", t.get_coro())
+                    t.print_stack()
+        except Exception:
+            pass
+        time.sleep(5)
+
 
 class LightningSpanProcessor(SpanProcessor):
     def __init__(self):
@@ -242,8 +279,13 @@ class LightningSpanProcessor(SpanProcessor):
     def _loop_runner(self):
         loop = asyncio.new_event_loop()
         self._loop = loop
+        self._loop.set_debug(True)
         asyncio.set_event_loop(loop)
         self._loop_ready.set()
+
+        thread = threading.Thread(target=debug_dump, args=(loop,), daemon=True)
+        thread.start()
+        # asyncio.create_task(heartbeat())
         loop.run_forever()
         loop.close()
 
@@ -261,6 +303,32 @@ class LightningSpanProcessor(SpanProcessor):
         # submit to the dedicated loop and wait synchronously
         if self._loop is None:
             raise RuntimeError("Loop is not initialized. This should not happen.")
+
+        # If already on the exporter loop thread, schedule and return immediately.
+        # ---------------------------------------------------------------------------
+        # WHY THIS CONDITIONAL EXISTS:
+        # In rare cases, span.end() is triggered from a LangchainCallbackHandler.__del__
+        # (or another finalizer) while the Python garbage collector is running on the
+        # *same thread* that owns our exporter event loop ("otel-loop").
+        #
+        # When that happens, on_end() executes on the exporter loop thread itself.
+        # If we were to call `asyncio.run_coroutine_threadsafe(...).result()` here,
+        # it would deadlock immediately — because the loop cannot both wait on and run
+        # the same coroutine. The Future stays pending forever and the loop stops
+        # processing scheduled callbacks.
+        #
+        # To avoid that self-deadlock, we detect when on_end() runs on the exporter
+        # loop thread. If so, we *schedule* the coroutine on the loop (fire-and-forget)
+        # instead of blocking with .result().
+        #
+        # This situation can occur because Python calls __del__ in whatever thread
+        # releases the last reference, which can easily be our loop thread if the
+        # object is dereferenced during loop._run_once().
+        # ---------------------------------------------------------------------------
+        if threading.current_thread() is self._loop_thread:
+            self._loop.call_soon_threadsafe(asyncio.create_task, coro)  # type: ignore
+            return None
+
         fut = asyncio.run_coroutine_threadsafe(coro, self._loop)  # type: ignore
         return fut.result(timeout=timeout)  # raises on error  # type: ignore
 
@@ -306,6 +374,10 @@ class LightningSpanProcessor(SpanProcessor):
         Args:
             span: The span that has ended.
         """
+        import traceback
+
+        # print("ON_END")
+        # print(traceback.format_stack())
         # Skip if span is not sampled
         if not span.context or not span.context.trace_flags.sampled:
             return
@@ -313,10 +385,27 @@ class LightningSpanProcessor(SpanProcessor):
         if self._store and self._rollout_id and self._attempt_id:
             try:
                 # Submit add_otel_span to the event loop and wait for it to complete
-                self._await_in_loop(
-                    self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
-                    timeout=5.0,
-                )
+                print("!!! before,")
+                print("Ready callbacks:", self._loop._ready)
+                print("Scheduled callbacks:", len(self._loop._scheduled))
+                if self._loop._scheduled:
+                    print("First in the queue:", self._loop._scheduled[0])
+                print("..... Current thread: ", threading.current_thread())
+                print("..... Loop thread: ", self._loop_thread)
+                if self._loop_thread.ident == threading.current_thread().ident:
+                    traceback.print_stack()
+                    print("Span content: ", span.attributes)
+                from opentelemetry.instrumentation.utils import suppress_instrumentation
+
+                with suppress_instrumentation():
+                    self._await_in_loop(
+                        self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
+                        timeout=30.0,
+                    )
+                print("!!! after,")
+                print("All tasks")
+                print("Ready callbacks:", self._loop._ready)
+                print("Scheduled callbacks:", self._loop._scheduled)
             except Exception:
                 # log; on_end MUST NOT raise
                 logger.exception(f"Error adding span to store: {span.name}")

agentlightning/tracer/base.py

@@ -1,14 +1,19 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from __future__ import annotations
+
 import logging
 from contextlib import contextmanager
-from typing import Any, Awaitable, Callable, Iterator, List, Optional
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterator, List, Optional
 
 from opentelemetry.sdk.trace import ReadableSpan
 
 from agentlightning.store.base import LightningStore
 from agentlightning.types import ParallelWorkerBase
 
+if TYPE_CHECKING:
+    from langchain.callbacks.base import BaseCallbackHandler
+
 logger = logging.getLogger(__name__)
 
 
@@ -42,7 +47,7 @@ class BaseTracer(ParallelWorkerBase):
 
     # Process the trace data
     if trace_tree:
-        rl_triplets = TraceTripletAdapter().adapt(spans)
+        rl_triplets = TracerTraceToTriplet().adapt(spans)
         # ... do something with the triplets
     ```
     """
@@ -112,3 +117,11 @@ class BaseTracer(ParallelWorkerBase):
         """
         with self.trace_context(name=func.__name__):
             return await func(*args, **kwargs)
+
+    def get_langchain_handler(self) -> Optional[BaseCallbackHandler]:
+        """Get a handler to install in langchain agent callback.
+
+        Agents are expected to use this handler in their agents to enable tracing.
+        """
+        logger.warning(f"{self.__class__.__name__} does not provide a LangChain callback handler.")
+        return None

agentlightning/trainer/__init__.py

@@ -1,5 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from .init_utils import build_component
 from .trainer import Trainer
 
-__all__ = ["Trainer"]
+__all__ = ["Trainer", "build_component"]

agentlightning/trainer/legacy.py

@@ -0,0 +1,367 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import multiprocessing
+import signal
+import time
+import warnings
+from typing import Any, List, Optional, TypeVar, Union
+
+from agentlightning.adapter import TraceAdapter, TracerTraceToTriplet
+from agentlightning.algorithm import BaseAlgorithm
+from agentlightning.client import AgentLightningClient
+from agentlightning.litagent import LitAgent
+from agentlightning.runner import LegacyAgentRunner
+from agentlightning.tracer.base import BaseTracer
+from agentlightning.types import Dataset, ParallelWorkerBase
+
+logger = logging.getLogger(__name__)
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+class TrainerLegacy(ParallelWorkerBase):
+    """Trainer for legacy mode for v0.1 compatibility."""
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        """Initialize the TrainerLegacy.
+
+        This method is mainly to make type checker happy.
+        It won't be used in practice.
+        """
+        self._dev = kwargs.pop("dev", False)
+        self.algorithm: Optional[BaseAlgorithm] = kwargs.pop("algorithm", None)
+        self.tracer: BaseTracer = kwargs.pop("tracer", None)
+        self.n_workers: int = kwargs.pop("n_workers", None)
+        self.max_tasks: Optional[int] = kwargs.pop("max_tasks", None)
+        self.daemon: bool = kwargs.pop("daemon", True)
+        self.triplet_exporter: TraceAdapter[Any] = kwargs.pop("triplet_exporter", None)
+
+    def _extract_client_from_data(
+        self, data: Union[str, AgentLightningClient, Dataset[Any]]
+    ) -> Optional[AgentLightningClient]:
+        """Extract client from data if it's a string URL or AgentLightningClient."""
+        if isinstance(data, str):
+            if not data.startswith("http://") and not data.startswith("https://"):
+                raise ValueError("String data must be a valid URL starting with http:// or https://")
+            return AgentLightningClient(endpoint=data)
+        elif isinstance(data, AgentLightningClient):
+            return data
+        return None
+
+    def _extract_dataset_from_data(
+        self, data: Union[str, AgentLightningClient, Dataset[Any]]
+    ) -> Optional[Dataset[Any]]:
+        """Extract dataset from data if it's a Dataset."""
+        if isinstance(data, str) or isinstance(data, AgentLightningClient):
+            return None
+        return data
+
+    def _determine_backend(
+        self,
+        train_data: Union[str, AgentLightningClient, Dataset[Any]],
+        dev_data: Union[str, AgentLightningClient, Dataset[Any], None] = None,
+    ) -> Union[str, AgentLightningClient]:
+        """Determine which backend to use for initialization."""
+        if self._dev:
+            if dev_data is None:
+                raise ValueError("dev_data must be provided when dev=True.")
+            client = self._extract_client_from_data(dev_data)
+            if client is None:
+                raise ValueError("dev_data must be a string URL or AgentLightningClient when dev=True.")
+            return client
+        else:
+            client = self._extract_client_from_data(train_data)
+            if client is None and self.algorithm is None:
+                raise ValueError(
+                    "train_data must be a string URL or AgentLightningClient when no algorithm is provided."
+                )
+            elif client is None and self.algorithm is not None:
+                # Algorithm will be responsible for creating the client
+                client = self.algorithm.get_client()
+                logger.info(f"Algorithm created client: {client}")
+                return client
+            if client is None:
+                raise ValueError(
+                    "train_data must be a string URL or AgentLightningClient when no algorithm is provided."
+                )
+            return client
+
+    def init(self, backend: Union[str, AgentLightningClient]) -> None:
+        logger.info(f"Initializing Trainer...")
+
+        self._init_client(backend)
+
+        self.tracer.init()
+
+        logger.info(f"Trainer main initialization complete.")
+
+    def teardown(self) -> None:
+        logger.info(f"Cleaning up Trainer...")
+        self.tracer.teardown()
+
+        self._client = None
+        logger.info(f"Trainer main cleanup complete.")
+
+    def client(self) -> AgentLightningClient:
+        """Returns the AgentLightningClient instance."""
+        if self._client is None:
+            raise RuntimeError("AgentLightningClient has not been initialized. Call `init` first.")
+        return self._client
+
+    def _init_client(self, backend: Union[str, AgentLightningClient]) -> AgentLightningClient:
+        if self._client is None:
+            if isinstance(backend, AgentLightningClient):
+                logger.info("Using provided AgentLightningClient instance.")
+                self._client = backend
+            else:
+                logger.info(f"Initializing AgentLightningClient with endpoint: {backend}")
+                if not isinstance(backend, str):  # type: ignore
+                    raise ValueError("backend must be a string URL or an AgentLightningClient instance.")
+                if not backend.startswith("http://") and not backend.startswith("https://"):
+                    raise ValueError("backend must be a valid URL starting with http:// or https://")
+                # Initialize the client with the provided backend URL
+                self._client = AgentLightningClient(endpoint=backend)
+        else:
+            logger.warning("AgentLightningClient already initialized. Returning existing instance.")
+        return self._client
+
+    def _worker_main_loop(self, agent: LitAgent[Any], worker_id: int, is_async: bool):
+        """The main function for each worker process.
+
+        This function initializes the client and the loop, then starts the
+        execution. It also configures process-specific settings like the
+        process title and signal handling.
+
+        Args:
+            agent: The `LitAgent` instance to run.
+            worker_id: The unique ID for this worker.
+            is_async: A boolean indicating if the async loop should be run.
+        """
+        if self.n_workers > 1:
+            import setproctitle
+
+            # Ignore Ctrl+C in worker processes; the main process handles it
+            signal.signal(signal.SIGINT, signal.SIG_IGN)
+            setproctitle.setproctitle(multiprocessing.current_process().name)
+
+        # Now we are in child processes, so we can safely set up the environment.
+        agent.set_trainer(self)  # type: ignore
+        if not isinstance(self.triplet_exporter, TracerTraceToTriplet):  # type: ignore
+            raise ValueError("triplet_exporter must be a TracerTraceToTriplet for the legacy trainer.")
+        # TODO: this should be set elsewhere
+        if agent.trained_agents:
+            self.triplet_exporter.agent_match = agent.trained_agents
+        self._initialize_worker_env(worker_id)
+
+        mode = "Async" if is_async else "Sync"
+        logger.info(f"[Worker {worker_id}] {mode} worker process started.")
+
+        num_processed = 0
+
+        try:
+            client = self.client()
+            loop = LegacyAgentRunner(
+                agent=agent,
+                client=client,
+                tracer=self.tracer,
+                triplet_exporter=self.triplet_exporter,
+                max_tasks=self.max_tasks,
+                worker_id=worker_id,
+            )
+            loop.init_worker(worker_id)  # type: ignore
+            if is_async:
+                num_processed = asyncio.run(loop.iter_async())
+            else:
+                num_processed = loop.iter()
+        except Exception:
+            logger.exception(f"[Worker {worker_id}] Unhandled exception in worker loop.")
+        finally:
+            self._teardown_worker_env(worker_id)
+
+        return num_processed
+
+    def _initialize_worker_env(self, worker_id: int):
+        logger.info(f"[Worker {worker_id}] Setting up trainer environment...")  # worker_id included in process name
+        self.tracer.init_worker(worker_id)
+
+    def _teardown_worker_env(self, worker_id: int):
+        logger.info(f"[Worker {worker_id}] Cleaning up trainer environment...")
+        self.tracer.teardown_worker(worker_id)
+        logger.info(f"[Worker {worker_id}] Environment cleanup complete.")
+
+    @staticmethod
+    def kill_orphaned_processes() -> None:
+        """
+        Kill any orphaned processes that may have been left behind by previous runs.
+        This is useful for cleaning up after crashes or unexpected exits.
+        """
+        import psutil
+
+        for proc in psutil.process_iter():  # type: ignore
+            # check whether the process name matches
+            if proc.name().startswith("AgentLightning-"):
+                proc.kill()
+
+    def _terminate_processes(self, processes: List[multiprocessing.Process]) -> None:
+        if self.n_workers > 1 and len(processes) > 0:
+            for i, p in enumerate(processes):
+                if p.is_alive():
+                    logger.info(f"Terminating worker {i} (name: {p.name}, PID: {p.pid})...")
+                    p.terminate()
+                else:
+                    logger.info(f"Worker {i} (name: {p.name}, PID: {p.pid}) is not alive or has already terminated.")
+            for i, p in enumerate(processes):
+                if p.is_alive():
+                    p.join(timeout=10)  # Give some time to terminate
+                if p.is_alive():  # If still alive, kill
+                    logger.warning(
+                        f"Worker {i} (name: {p.name}, PID: {p.pid}) did not terminate gracefully, killing..."
+                    )
+                    p.kill()
+                    p.join(timeout=10)  # Ensure it's reaped
+
+    def fit_v0(
+        self,
+        agent: LitAgent[T_co],
+        train_data: Union[str, AgentLightningClient, Dataset[T_co]],
+        *,
+        val_data: Union[str, AgentLightningClient, Dataset[T_co], None] = None,
+        dev_data: Union[str, AgentLightningClient, Dataset[T_co], None] = None,
+        dev_backend: Union[str, AgentLightningClient, None] = None,
+    ):
+        """Train the agent using the provided data.
+
+        Each data argument can be a string URL connecting to a agent-lightning server,
+        or an AgentLightningClient instance connecting to a server (or mock server), or a dataset.
+        If no algorithm is provided when instantiating the trainer, the data must be
+        provided to connecting a server. Otherwise, dataset is also allowed and will be
+        passed to the algorithm.
+
+        If the algorithm is instantiated and there is no URL/client provided,
+        the algorithm will be responsible for creating a client that will connect to itself.
+        It can also create a mock client if the algorithm does not require a server.
+        """
+
+        if dev_backend is not None:
+            warnings.warn("dev_backend is deprecated. Use dev_data instead.")
+            if dev_data is not None:
+                raise ValueError("dev_data and dev_backend cannot be provided at the same time.")
+            dev_data = dev_backend
+
+        # Extract datasets for algorithm if available
+        train_dataset = self._extract_dataset_from_data(train_data)
+        val_dataset = self._extract_dataset_from_data(val_data) if val_data else None
+
+        # Initialize the algorithm with trainer if provided
+        if self.algorithm is not None:
+            self.algorithm.set_trainer(self)  # type: ignore
+            # DO NOT RUN TRAINING HERE. Need to spawn the worker first.
+
+        # Determine the backend to use for client-server mode
+        backend = self._determine_backend(train_data, dev_data)
+
+        if self._dev:
+            logger.warning(f"Running in dev mode. Using dev backend: {backend}")
+        else:
+            logger.debug(f"Running in non-dev mode. Using backend: {backend}")
+
+        self.init(backend)
+
+        processes: List[multiprocessing.Process] = []
+
+        # Determine if the agent is asynchronous
+
+        mode = "asynchronous" if agent.is_async() else "synchronous"
+
+        try:
+            if self.n_workers == 1:
+                logger.info(f"Running with n_workers=1 ({mode} in main process).")
+
+                # Warn if algorithm is set with single worker mode
+                if self.algorithm is not None:
+                    logger.warning(
+                        "Algorithm is set but using single worker mode. Algorithm will never get the chance to run."
+                    )
+                    # Ideally the single worker should be run in a separate thread or process.
+
+                num_tasks = self._worker_main_loop(agent, 0, agent.is_async())
+                logger.info(f"Single worker mode finished. Tasks processed: {num_tasks}")
+
+                # If algorithm is provided and we have datasets, run algorithm after worker completes
+                if self.algorithm is not None and train_dataset is not None:
+                    logger.info("Running algorithm training after worker completion.")
+                    self.algorithm.run(
+                        train_dataset=train_dataset,
+                        val_dataset=val_dataset,
+                    )
+            else:
+                logger.info(f"Running with n_workers={self.n_workers} ({mode} multiprocessing).")
+                for i in range(self.n_workers):
+                    process_name = f"AgentLightning-Worker-{i}"
+                    p = multiprocessing.Process(
+                        target=self._worker_main_loop,
+                        args=(agent, i, agent.is_async()),
+                        daemon=self.daemon,
+                        name=process_name,
+                    )
+                    processes.append(p)
+                    logger.info(f"Starting worker process {i} (name: {process_name})...")
+                    p.start()
+
+                if self.daemon:
+                    # If algorithm is provided and we have datasets, pass them to the algorithm
+                    if self.algorithm is not None:
+                        logger.info("All workers have been spawned. Running algorithm training with provided datasets.")
+                        self.algorithm.run(
+                            train_dataset=train_dataset,
+                            val_dataset=val_dataset,
+                        )
+                        logger.info("Algorithm exits. Killing the workers.")
+                        self._terminate_processes(processes)
+
+                    for i, p in enumerate(processes):
+                        p.join()  # Wait for the process to complete
+                        logger.info(
+                            f"Worker process {i} (name: {p.name}, PID: {p.pid}) joined with exit code {p.exitcode}."
+                        )
+                        if p.exitcode != 0:
+                            logger.warning(
+                                f"Worker process {i} (name: {p.name}, PID: {p.pid}) exited with non-zero code: {p.exitcode}."
+                            )
+
+                    logger.info(f"All {self.n_workers} worker processes have completed.")
+                else:
+                    logger.info("All worker processes started. Main process will not wait.")
+
+                    # A hack to stop the main process from waiting for child processes to finish.
+                    time.sleep(1)  # Give workers time to start
+                    import multiprocessing.process as multiprocessing_process
+
+                    multiprocessing_process._children.clear()  # type: ignore
+
+                    if self.algorithm is not None:
+                        logger.info("Main process continues to run algorithm.")
+                        self.algorithm.run(
+                            train_dataset=train_dataset,
+                            val_dataset=val_dataset,
+                        )
+                        logger.info("Algorithm exits. Killing the workers.")
+                        self._terminate_processes(processes)
+
+        except KeyboardInterrupt:
+            logger.info("KeyboardInterrupt received. Killing the workers.")
+            self._terminate_processes(processes)
+            logger.info(f"Workers terminated or single worker interrupted.")
+            raise
+        except Exception:
+            logger.exception(f"Unhandled exception in fit method.")
+            self._terminate_processes(processes)
+            logger.info(f"Workers terminated or single worker interrupted.")
+            raise
+        finally:
+            if self.daemon:
+                self.teardown()
+            else:
+                logger.info("Main process exiting. Please use Trainer.kill_orphaned_processes() for cleanup.")

agentlightning/trainer/registry.py

@@ -1,5 +1,10 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+"""Put components in this file to make them available to the Trainer.
+
+Currently only used for ExecutionStrategy.
+"""
+
 ExecutionStrategyRegistry = {
     "shm": "agentlightning.execution.shared_memory.SharedMemoryExecutionStrategy",
     # "ipc": "agentlightning.execution.inter_process.InterProcessExecutionStrategy",

agentlightning/trainer/trainer.py

@@ -3,29 +3,26 @@
 import asyncio
 import functools
 import logging
-import multiprocessing
-import signal
-import time
 import warnings
-from typing import Any, Callable, Dict, List, Optional, Sequence, TypeVar, Union
+from typing import Any, Callable, Dict, Optional, Sequence, TypeVar, Union
 
-from agentlightning.adapter import TraceAdapter, TraceTripletAdapter
-from agentlightning.algorithm.base import BaseAlgorithm, FastAlgorithm
-from agentlightning.algorithm.mock import MockAlgorithm
+from agentlightning.adapter import TraceAdapter, TracerTraceToTriplet
+from agentlightning.algorithm import BaseAlgorithm, Baseline, FastAlgorithm
 from agentlightning.client import AgentLightningClient
 from agentlightning.execution.base import ExecutionStrategy
 from agentlightning.execution.client_server import ClientServerExecutionStrategy
-from agentlightning.execution.events import Event
+from agentlightning.execution.events import ExecutionEvent
 from agentlightning.litagent import LitAgent
 from agentlightning.llm_proxy import LLMProxy
-from agentlightning.runner import AgentRunner, AgentRunnerV2, BaseRunner
+from agentlightning.runner import BaseRunner, LitAgentRunner
 from agentlightning.store.base import LightningStore
 from agentlightning.store.memory import InMemoryLightningStore
 from agentlightning.tracer.agentops import AgentOpsTracer
 from agentlightning.tracer.base import BaseTracer
-from agentlightning.types import Dataset, Hook, NamedResources, ParallelWorkerBase
+from agentlightning.types import Dataset, Hook, NamedResources
 
 from .init_utils import build_component, instantiate_component
+from .legacy import TrainerLegacy
 from .registry import ExecutionStrategyRegistry
 
 logger = logging.getLogger(__name__)
@@ -36,7 +33,7 @@ T = TypeVar("T")
 ComponentSpec = Union[T, type[T], Callable[[], T], str, Dict[str, Any], None]
 
 
-class Trainer(ParallelWorkerBase):
+class Trainer(TrainerLegacy):
     """Orchestrates the distributed execution of agent rollouts.
 
     The Trainer is responsible for launching one or more worker processes
@@ -60,7 +57,7 @@ class Trainer(ParallelWorkerBase):
                 If None, a default `AgentOpsTracer` will be created with the current settings.
         hooks: A sequence of `Hook` instances to be called at various lifecycle stages (e.g., on_trace_start,
                on_trace_end, on_rollout_start, on_rollout_end).
-        adapter: An instance of `TraceTripletAdapter` to export data consumble by algorithms from traces.
+        adapter: An instance of `TracerTraceToTriplet` to export data consumble by algorithms from traces.
         llm_proxy: An instance of `LLMProxy` to use for intercepting the LLM calls.
                    If not provided, algorithm will create one on its own.
         n_workers: Number of agent workers to run in parallel. Deprecated in favor of `n_runners`.
@@ -68,7 +65,7 @@ class Trainer(ParallelWorkerBase):
         daemon: Whether worker processes should be daemons. Daemon processes
                 are terminated automatically when the main process exits. Deprecated.
                 Only have effect with `fit_v0`.
-        triplet_exporter: An instance of `TraceTripletAdapter` to export triplets from traces,
+        triplet_exporter: An instance of `TracerTraceToTriplet` to export triplets from traces,
                           or a dictionary with the initialization parameters for the exporter.
                           Deprecated. Use `adapter` instead.
         dev: If True, rollouts are run against the dev endpoint provided in `fit`.
@@ -92,10 +89,13 @@ class Trainer(ParallelWorkerBase):
         n_workers: Optional[int] = None,
         max_tasks: Optional[int] = None,
         daemon: bool = True,
-        triplet_exporter: ComponentSpec[TraceTripletAdapter] = None,
+        triplet_exporter: ComponentSpec[TracerTraceToTriplet] = None,
         hooks: Optional[Union[Hook, Sequence[Hook]]] = None,
     ):
-        super().__init__()
+        # Do not call super().__init__() here.
+        # super().__init__() will call TrainerLegacy's initialization, which is not intended.
+        self.worker_id: Optional[int] = None
+
         self._dev = dev
         self.daemon = daemon
         self._client: AgentLightningClient | None = None  # Will be initialized in fit or fit_v0
@@ -212,9 +212,9 @@ class Trainer(ParallelWorkerBase):
             adapter,
             expected_type=TraceAdapter,
             spec_name="adapter",
-            default_factory=TraceTripletAdapter,
+            default_factory=TracerTraceToTriplet,
             dict_requires_type=False,
-            dict_default_cls=TraceTripletAdapter,
+            dict_default_cls=TracerTraceToTriplet,
             invalid_spec_error_fmt="Invalid adapter type: {actual_type}. Expected TraceAdapter, dict, or None.",
             type_error_fmt="Adapter factory returned {type_name}, which is not a TraceAdapter subclass.",
         )
@@ -283,7 +283,7 @@ class Trainer(ParallelWorkerBase):
             optional_defaults["max_rollouts"] = lambda: self.max_rollouts
 
         def default_runner_factory() -> BaseRunner[Any]:
-            return instantiate_component(AgentRunnerV2, optional_defaults=optional_defaults)
+            return instantiate_component(LitAgentRunner, optional_defaults=optional_defaults)
 
         return build_component(
             runner,
@@ -302,7 +302,7 @@ class Trainer(ParallelWorkerBase):
             return (hooks,)
         return tuple(hooks)
 
-    def fit_v2(
+    def fit(
         self,
         agent: LitAgent[T_co],
         train_dataset: Optional[Dataset[T_co]] = None,
@@ -349,7 +349,7 @@ class Trainer(ParallelWorkerBase):
 
         # Sanity check
         if self.algorithm is None:
-            algorithm = MockAlgorithm()
+            algorithm = Baseline()
         else:
             algorithm = self.algorithm
 
@@ -371,7 +371,7 @@ class Trainer(ParallelWorkerBase):
     async def _algorithm_bundle(
         self,
         store: LightningStore,
-        event: Event,
+        event: ExecutionEvent,
         train_dataset: Optional[Dataset[T_co]],
         val_dataset: Optional[Dataset[T_co]],
         algorithm: Optional[BaseAlgorithm],
@@ -407,7 +407,9 @@ class Trainer(ParallelWorkerBase):
             logger.exception("Algorithm bundle encountered an error.")
             raise
 
-    async def _runner_bundle(self, store: LightningStore, worker_id: int, event: Event, agent: LitAgent[T_co]) -> None:
+    async def _runner_bundle(
+        self, store: LightningStore, worker_id: int, event: ExecutionEvent, agent: LitAgent[T_co]
+    ) -> None:
         runner_instance: BaseRunner[Any] | None = None
         runner_initialized = False
         worker_initialized = False
@@ -434,331 +436,3 @@ class Trainer(ParallelWorkerBase):
                         runner_instance.teardown()
                     except Exception:
                         logger.exception("Error during runner teardown (worker_id=%s).", worker_id)
-
-    def _extract_client_from_data(
-        self, data: Union[str, AgentLightningClient, Dataset[Any]]
-    ) -> Optional[AgentLightningClient]:
-        """Extract client from data if it's a string URL or AgentLightningClient."""
-        if isinstance(data, str):
-            if not data.startswith("http://") and not data.startswith("https://"):
-                raise ValueError("String data must be a valid URL starting with http:// or https://")
-            return AgentLightningClient(endpoint=data)
-        elif isinstance(data, AgentLightningClient):
-            return data
-        return None
-
-    def _extract_dataset_from_data(
-        self, data: Union[str, AgentLightningClient, Dataset[Any]]
-    ) -> Optional[Dataset[Any]]:
-        """Extract dataset from data if it's a Dataset."""
-        if isinstance(data, str) or isinstance(data, AgentLightningClient):
-            return None
-        return data
-
-    def _determine_backend(
-        self,
-        train_data: Union[str, AgentLightningClient, Dataset[Any]],
-        dev_data: Union[str, AgentLightningClient, Dataset[Any], None] = None,
-    ) -> Union[str, AgentLightningClient]:
-        """Determine which backend to use for initialization."""
-        if self._dev:
-            if dev_data is None:
-                raise ValueError("dev_data must be provided when dev=True.")
-            client = self._extract_client_from_data(dev_data)
-            if client is None:
-                raise ValueError("dev_data must be a string URL or AgentLightningClient when dev=True.")
-            return client
-        else:
-            client = self._extract_client_from_data(train_data)
-            if client is None and self.algorithm is None:
-                raise ValueError(
-                    "train_data must be a string URL or AgentLightningClient when no algorithm is provided."
-                )
-            elif client is None and self.algorithm is not None:
-                # Algorithm will be responsible for creating the client
-                client = self.algorithm.get_client()
-                logger.info(f"Algorithm created client: {client}")
-                return client
-            if client is None:
-                raise ValueError(
-                    "train_data must be a string URL or AgentLightningClient when no algorithm is provided."
-                )
-            return client
-
-    def init(self, backend: Union[str, AgentLightningClient]) -> None:
-        logger.info(f"Initializing Trainer...")
-
-        self._init_client(backend)
-
-        self.tracer.init()
-
-        logger.info(f"Trainer main initialization complete.")
-
-    def teardown(self) -> None:
-        logger.info(f"Cleaning up Trainer...")
-        self.tracer.teardown()
-
-        self._client = None
-        logger.info(f"Trainer main cleanup complete.")
-
-    def client(self) -> AgentLightningClient:
-        """Returns the AgentLightningClient instance."""
-        if self._client is None:
-            raise RuntimeError("AgentLightningClient has not been initialized. Call `init` first.")
-        return self._client
-
-    def _init_client(self, backend: Union[str, AgentLightningClient]) -> AgentLightningClient:
-        if self._client is None:
-            if isinstance(backend, AgentLightningClient):
-                logger.info("Using provided AgentLightningClient instance.")
-                self._client = backend
-            else:
-                logger.info(f"Initializing AgentLightningClient with endpoint: {backend}")
-                if not isinstance(backend, str):  # type: ignore
-                    raise ValueError("backend must be a string URL or an AgentLightningClient instance.")
-                if not backend.startswith("http://") and not backend.startswith("https://"):
-                    raise ValueError("backend must be a valid URL starting with http:// or https://")
-                # Initialize the client with the provided backend URL
-                self._client = AgentLightningClient(endpoint=backend)
-        else:
-            logger.warning("AgentLightningClient already initialized. Returning existing instance.")
-        return self._client
-
-    def _worker_main_loop(self, agent: LitAgent[Any], worker_id: int, is_async: bool):
-        """The main function for each worker process.
-
-        This function initializes the client and the loop, then starts the
-        execution. It also configures process-specific settings like the
-        process title and signal handling.
-
-        Args:
-            agent: The `LitAgent` instance to run.
-            worker_id: The unique ID for this worker.
-            is_async: A boolean indicating if the async loop should be run.
-        """
-        if self.n_workers > 1:
-            import setproctitle
-
-            # Ignore Ctrl+C in worker processes; the main process handles it
-            signal.signal(signal.SIGINT, signal.SIG_IGN)
-            setproctitle.setproctitle(multiprocessing.current_process().name)
-
-        # Now we are in child processes, so we can safely set up the environment.
-        agent.set_trainer(self)
-        if not isinstance(self.triplet_exporter, TraceTripletAdapter):
-            raise ValueError("triplet_exporter must be a TraceTripletAdapter for the legacy trainer.")
-        # TODO: this should be set elsewhere
-        if agent.trained_agents:
-            self.triplet_exporter.agent_match = agent.trained_agents
-        self._initialize_worker_env(worker_id)
-
-        mode = "Async" if is_async else "Sync"
-        logger.info(f"[Worker {worker_id}] {mode} worker process started.")
-
-        num_processed = 0
-
-        try:
-            client = self.client()
-            loop = AgentRunner(
-                agent=agent,
-                client=client,
-                tracer=self.tracer,
-                triplet_exporter=self.triplet_exporter,
-                max_tasks=self.max_tasks,
-                worker_id=worker_id,
-            )
-            loop.init_worker(worker_id)  # type: ignore
-            if is_async:
-                num_processed = asyncio.run(loop.iter_async())
-            else:
-                num_processed = loop.iter()
-        except Exception:
-            logger.exception(f"[Worker {worker_id}] Unhandled exception in worker loop.")
-        finally:
-            self._teardown_worker_env(worker_id)
-
-        return num_processed
-
-    def _initialize_worker_env(self, worker_id: int):
-        logger.info(f"[Worker {worker_id}] Setting up trainer environment...")  # worker_id included in process name
-        self.tracer.init_worker(worker_id)
-
-    def _teardown_worker_env(self, worker_id: int):
-        logger.info(f"[Worker {worker_id}] Cleaning up trainer environment...")
-        self.tracer.teardown_worker(worker_id)
-        logger.info(f"[Worker {worker_id}] Environment cleanup complete.")
-
-    @staticmethod
-    def kill_orphaned_processes() -> None:
-        """
-        Kill any orphaned processes that may have been left behind by previous runs.
-        This is useful for cleaning up after crashes or unexpected exits.
-        """
-        import psutil
-
-        for proc in psutil.process_iter():  # type: ignore
-            # check whether the process name matches
-            if proc.name().startswith("AgentLightning-"):
-                proc.kill()
-
-    def _terminate_processes(self, processes: List[multiprocessing.Process]) -> None:
-        if self.n_workers > 1 and len(processes) > 0:
-            for i, p in enumerate(processes):
-                if p.is_alive():
-                    logger.info(f"Terminating worker {i} (name: {p.name}, PID: {p.pid})...")
-                    p.terminate()
-                else:
-                    logger.info(f"Worker {i} (name: {p.name}, PID: {p.pid}) is not alive or has already terminated.")
-            for i, p in enumerate(processes):
-                if p.is_alive():
-                    p.join(timeout=10)  # Give some time to terminate
-                if p.is_alive():  # If still alive, kill
-                    logger.warning(
-                        f"Worker {i} (name: {p.name}, PID: {p.pid}) did not terminate gracefully, killing..."
-                    )
-                    p.kill()
-                    p.join(timeout=10)  # Ensure it's reaped
-
-    def fit(
-        self,
-        agent: LitAgent[T_co],
-        train_data: Union[str, AgentLightningClient, Dataset[T_co]],
-        *,
-        val_data: Union[str, AgentLightningClient, Dataset[T_co], None] = None,
-        dev_data: Union[str, AgentLightningClient, Dataset[T_co], None] = None,
-        dev_backend: Union[str, AgentLightningClient, None] = None,
-    ):
-        """Train the agent using the provided data.
-
-        Each data argument can be a string URL connecting to a agent-lightning server,
-        or an AgentLightningClient instance connecting to a server (or mock server), or a dataset.
-        If no algorithm is provided when instantiating the trainer, the data must be
-        provided to connecting a server. Otherwise, dataset is also allowed and will be
-        passed to the algorithm.
-
-        If the algorithm is instantiated and there is no URL/client provided,
-        the algorithm will be responsible for creating a client that will connect to itself.
-        It can also create a mock client if the algorithm does not require a server.
-        """
-
-        if dev_backend is not None:
-            warnings.warn("dev_backend is deprecated. Use dev_data instead.")
-            if dev_data is not None:
-                raise ValueError("dev_data and dev_backend cannot be provided at the same time.")
-            dev_data = dev_backend
-
-        # Extract datasets for algorithm if available
-        train_dataset = self._extract_dataset_from_data(train_data)
-        val_dataset = self._extract_dataset_from_data(val_data) if val_data else None
-
-        # Initialize the algorithm with trainer if provided
-        if self.algorithm is not None:
-            self.algorithm.set_trainer(self)
-            # DO NOT RUN TRAINING HERE. Need to spawn the worker first.
-
-        # Determine the backend to use for client-server mode
-        backend = self._determine_backend(train_data, dev_data)
-
-        if self._dev:
-            logger.warning(f"Running in dev mode. Using dev backend: {backend}")
-        else:
-            logger.debug(f"Running in non-dev mode. Using backend: {backend}")
-
-        self.init(backend)
-
-        processes: List[multiprocessing.Process] = []
-
-        # Determine if the agent is asynchronous
-
-        mode = "asynchronous" if agent.is_async() else "synchronous"
-
-        try:
-            if self.n_workers == 1:
-                logger.info(f"Running with n_workers=1 ({mode} in main process).")
-
-                # Warn if algorithm is set with single worker mode
-                if self.algorithm is not None:
-                    logger.warning(
-                        "Algorithm is set but using single worker mode. Algorithm will never get the chance to run."
-                    )
-                    # Ideally the single worker should be run in a separate thread or process.
-
-                num_tasks = self._worker_main_loop(agent, 0, agent.is_async())
-                logger.info(f"Single worker mode finished. Tasks processed: {num_tasks}")
-
-                # If algorithm is provided and we have datasets, run algorithm after worker completes
-                if self.algorithm is not None and train_dataset is not None:
-                    logger.info("Running algorithm training after worker completion.")
-                    self.algorithm.run(
-                        train_dataset=train_dataset,
-                        val_dataset=val_dataset,
-                    )
-            else:
-                logger.info(f"Running with n_workers={self.n_workers} ({mode} multiprocessing).")
-                for i in range(self.n_workers):
-                    process_name = f"AgentLightning-Worker-{i}"
-                    p = multiprocessing.Process(
-                        target=self._worker_main_loop,
-                        args=(agent, i, agent.is_async()),
-                        daemon=self.daemon,
-                        name=process_name,
-                    )
-                    processes.append(p)
-                    logger.info(f"Starting worker process {i} (name: {process_name})...")
-                    p.start()
-
-                if self.daemon:
-                    # If algorithm is provided and we have datasets, pass them to the algorithm
-                    if self.algorithm is not None:
-                        logger.info("All workers have been spawned. Running algorithm training with provided datasets.")
-                        self.algorithm.run(
-                            train_dataset=train_dataset,
-                            val_dataset=val_dataset,
-                        )
-                        logger.info("Algorithm exits. Killing the workers.")
-                        self._terminate_processes(processes)
-
-                    for i, p in enumerate(processes):
-                        p.join()  # Wait for the process to complete
-                        logger.info(
-                            f"Worker process {i} (name: {p.name}, PID: {p.pid}) joined with exit code {p.exitcode}."
-                        )
-                        if p.exitcode != 0:
-                            logger.warning(
-                                f"Worker process {i} (name: {p.name}, PID: {p.pid}) exited with non-zero code: {p.exitcode}."
-                            )
-
-                    logger.info(f"All {self.n_workers} worker processes have completed.")
-                else:
-                    logger.info("All worker processes started. Main process will not wait.")
-
-                    # A hack to stop the main process from waiting for child processes to finish.
-                    time.sleep(1)  # Give workers time to start
-                    import multiprocessing.process as multiprocessing_process
-
-                    multiprocessing_process._children.clear()  # type: ignore
-
-                    if self.algorithm is not None:
-                        logger.info("Main process continues to run algorithm.")
-                        self.algorithm.run(
-                            train_dataset=train_dataset,
-                            val_dataset=val_dataset,
-                        )
-                        logger.info("Algorithm exits. Killing the workers.")
-                        self._terminate_processes(processes)
-
-        except KeyboardInterrupt:
-            logger.info("KeyboardInterrupt received. Killing the workers.")
-            self._terminate_processes(processes)
-            logger.info(f"Workers terminated or single worker interrupted.")
-            raise
-        except Exception:
-            logger.exception(f"Unhandled exception in fit method.")
-            self._terminate_processes(processes)
-            logger.info(f"Workers terminated or single worker interrupted.")
-            raise
-        finally:
-            if self.daemon:
-                self.teardown()
-            else:
-                logger.info("Main process exiting. Please use Trainer.kill_orphaned_processes() for cleanup.")

agentlightning/types/core.py

@@ -12,6 +12,7 @@ from typing import (
     Literal,
     Optional,
     Protocol,
+    SupportsIndex,
     TypeVar,
     Union,
     cast,
@@ -29,12 +30,12 @@ if TYPE_CHECKING:
 
 __all__ = [
     "Triplet",
-    "Rollout",
+    "RolloutLegacy",
     "Task",
     "TaskInput",
     "TaskIfAny",
+    "RolloutRawResultLegacy",
     "RolloutRawResult",
-    "RolloutRawResultV2",
     "RolloutMode",
     "GenericResponse",
     "ParallelWorkerBase",
@@ -42,7 +43,7 @@ __all__ = [
     "AttemptStatus",
     "RolloutStatus",
     "RolloutConfig",
-    "RolloutV2",
+    "Rollout",
     "Attempt",
     "AttemptedRollout",
     "Hook",
@@ -60,7 +61,7 @@ class Triplet(BaseModel):
     metadata: Dict[str, Any] = Field(default_factory=dict)
 
 
-class Rollout(BaseModel):
+class RolloutLegacy(BaseModel):
     """The standard reporting object from client to server."""
 
     rollout_id: str
@@ -141,7 +142,7 @@ class RolloutConfig(BaseModel):
     )  # list of statuses that should trigger a retry
 
 
-class RolloutV2(BaseModel):
+class Rollout(BaseModel):
     rollout_id: str
 
     # Inputs
@@ -163,7 +164,7 @@ class RolloutV2(BaseModel):
     metadata: Optional[Dict[str, Any]] = None
 
 
-class AttemptedRollout(RolloutV2):
+class AttemptedRollout(Rollout):
     """A rollout along with its active attempt."""
 
     attempt: Attempt
@@ -176,10 +177,11 @@ class AttemptedRollout(RolloutV2):
 
 
 TaskInput = Any
+"""Task input type. Can be any type."""
 
 
 class Task(BaseModel):
-    """A task (rollout request) to be processed by the client agent."""
+    """A task (rollout request) to be processed by the client agent. Deprecated."""
 
     rollout_id: str
     input: TaskInput
@@ -201,9 +203,9 @@ class TaskIfAny(BaseModel):
     task: Optional[Task] = None
 
 
-RolloutRawResult = Union[None, float, List[Triplet], List[Dict[str, Any]], List[ReadableSpan], Rollout]
+RolloutRawResultLegacy = Union[None, float, List[Triplet], List[Dict[str, Any]], List[ReadableSpan], RolloutLegacy]
 
-RolloutRawResultV2 = Union[
+RolloutRawResult = Union[
     None,  # nothing (relies on tracer)
     float,  # only final reward
     List[ReadableSpan],  # constructed OTEL spans by user
@@ -264,7 +266,7 @@ class Dataset(Protocol, Generic[T_co]):
     You don't have to inherit from this class; you can use a simple list if you want to.
     """
 
-    def __getitem__(self, index: int) -> T_co: ...
+    def __getitem__(self, index: SupportsIndex, /) -> T_co: ...
 
     def __len__(self) -> int: ...
 
@@ -273,7 +275,7 @@ class Hook(ParallelWorkerBase):
     """Base class for defining hooks in the agent runner's lifecycle."""
 
     async def on_trace_start(
-        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: RolloutV2
+        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: Rollout
     ) -> None:
         """Hook called immediately after the tracer enters the trace context but before the rollout begins.
 
@@ -281,14 +283,14 @@ class Hook(ParallelWorkerBase):
             agent: The :class:`LitAgent` instance associated with the runner.
             runner: The :class:`BaseRunner` managing the rollout.
             tracer: The :class:`BaseTracer` instance associated with the runner.
-            rollout: The :class:`RolloutV2` object that will be processed.
+            rollout: The :class:`Rollout` object that will be processed.
 
         Subclasses can override this method to implement custom logic such as logging,
         metric collection, or resource setup. By default, this is a no-op.
         """
 
     async def on_trace_end(
-        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: RolloutV2
+        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: Rollout
     ) -> None:
         """Hook called immediately after the rollout completes but before the tracer exits the trace context.
 
@@ -296,19 +298,19 @@ class Hook(ParallelWorkerBase):
             agent: The :class:`LitAgent` instance associated with the runner.
             runner: The :class:`BaseRunner` managing the rollout.
             tracer: The :class:`BaseTracer` instance associated with the runner.
-            rollout: The :class:`RolloutV2` object that has been processed.
+            rollout: The :class:`Rollout` object that has been processed.
 
         Subclasses can override this method to implement custom logic such as logging,
         metric collection, or resource cleanup. By default, this is a no-op.
         """
 
-    async def on_rollout_start(self, *, agent: LitAgent[Any], runner: BaseRunner[Any], rollout: RolloutV2) -> None:
+    async def on_rollout_start(self, *, agent: LitAgent[Any], runner: BaseRunner[Any], rollout: Rollout) -> None:
         """Hook called immediately before a rollout *attempt* begins.
 
         Args:
             agent: The :class:`LitAgent` instance associated with the runner.
             runner: The :class:`BaseRunner` managing the rollout.
-            rollout: The :class:`RolloutV2` object that will be processed.
+            rollout: The :class:`Rollout` object that will be processed.
 
         Subclasses can override this method to implement custom logic such as
         logging, metric collection, or resource setup. By default, this is a
@@ -320,7 +322,7 @@ class Hook(ParallelWorkerBase):
         *,
         agent: LitAgent[Any],
         runner: BaseRunner[Any],
-        rollout: RolloutV2,
+        rollout: Rollout,
         spans: Union[List[ReadableSpan], List[Span]],
     ) -> None:
         """Hook called after a rollout *attempt* completes.
@@ -328,7 +330,7 @@ class Hook(ParallelWorkerBase):
         Args:
             agent: The :class:`LitAgent` instance associated with the runner.
             runner: The :class:`BaseRunner` managing the rollout.
-            rollout: The :class:`RolloutV2` object that has been processed.
+            rollout: The :class:`Rollout` object that has been processed.
             spans: The spans that have been added to the store.
 
         Subclasses can override this method for cleanup or additional

agentlightning/verl/__init__.py

@@ -1,5 +1,8 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+"""This package contains a *hacky* integration of VERL with Agent Lightning."""
+
 from .daemon import *
 from .dataset import *
+from .entrypoint import *
 from .trainer import *

agentlightning/verl/daemon.py

@@ -17,14 +17,20 @@ from flask import Flask, Response, abort, request
 from tensordict import TensorDict
 from verl import DataProto
 
-from agentlightning import LLM, AgentLightningServer, NamedResources, Rollout, configure_logger
-from agentlightning.adapter.triplet import BaseTraceTripletAdapter, TraceTripletAdapter
+from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLegacy, configure_logger
+from agentlightning.adapter.triplet import TracerTraceToTriplet, TraceToTripletBase
 from agentlightning.llm_proxy import LLMProxy, ModelConfig
 from agentlightning.store.base import LightningStore
-from agentlightning.types import RolloutConfig, RolloutV2, Task
+from agentlightning.types import Rollout, RolloutConfig, Task
 
 configure_logger()
 
+__all__ = [
+    "AgentModeDaemon",
+    "get_left_padded_ids_and_attention_mask",
+    "get_right_padded_ids_and_attention_mask",
+]
+
 
 def get_left_padded_ids_and_attention_mask(
     ids: List[int], max_length: int, pad_token_id: int
@@ -138,7 +144,7 @@ class AgentModeDaemon:
         mode: Literal["v0", "v1"] = "v1",
         llm_proxy: LLMProxy | None = None,
         store: LightningStore | None = None,
-        adapter: BaseTraceTripletAdapter | None = None,
+        adapter: TraceToTripletBase | None = None,
     ):
         self.mode = mode
         self.llm_timeout_seconds = llm_timeout_seconds
@@ -164,7 +170,7 @@ class AgentModeDaemon:
                 # Reuse the existing LLM proxy (probably configured by user)
                 self.llm_proxy = llm_proxy
             if adapter is None:
-                self.adapter = TraceTripletAdapter()
+                self.adapter = TracerTraceToTriplet()
             else:
                 # Reuse the one from trainer
                 self.adapter = adapter
@@ -183,7 +189,7 @@ class AgentModeDaemon:
         # Internal State
         self.backend_llm_server_addresses: List[str] = []
         self._total_tasks_queued = 0
-        self._completed_rollouts_v0: Dict[str, Rollout] = {}
+        self._completed_rollouts_v0: Dict[str, RolloutLegacy] = {}
         self._task_id_to_original_sample: Dict[str, Dict[str, Any]] = {}
         self._server_thread: Optional[threading.Thread] = None
         self._proxy_thread: Optional[threading.Thread] = None
@@ -434,7 +440,7 @@ class AgentModeDaemon:
             print(f"Failed to set up data on server: {e}")
             raise
 
-    def _validate_data(self, rollout: Rollout):
+    def _validate_data(self, rollout: RolloutLegacy):
         if rollout.final_reward is None:
             print(
                 f"Warning: Reward is None for rollout {rollout.rollout_id}, will be auto-set to {self.reward_fillna_value}."
@@ -448,10 +454,10 @@ class AgentModeDaemon:
         elif any(not r.prompt.get("token_ids", []) for r in rollout.triplets):
             print(f"Warning: Rollout {rollout.rollout_id} contains empty prompt: {rollout.triplets}")
 
-    async def _validate_data_v1(self, rollout: RolloutV2) -> Rollout:
-        """Convert RolloutV2 to Rollout and validate.
+    async def _validate_data_v1(self, rollout: Rollout) -> RolloutLegacy:
+        """Convert Rollout to RolloutLegacy and validate.
 
-        1. Task: construct from RolloutV2
+        1. Task: construct from Rollout
         2. Triplets: obtained by querying spans and feeding into the adapter
         3. Final reward: extracted from last triplet's reward, searching backwards if not found
         """
@@ -474,7 +480,7 @@ class AgentModeDaemon:
                     final_reward = triplet.reward
                     break
 
-        # Construct the Task object from RolloutV2
+        # Construct the Task object from Rollout
         task = Task(
             rollout_id=rollout.rollout_id,
             input=rollout.input,
@@ -484,7 +490,7 @@ class AgentModeDaemon:
         )
 
         # Create the Rollout object (without trace and logs as per user's note)
-        result_rollout = Rollout(
+        result_rollout = RolloutLegacy(
             rollout_id=rollout.rollout_id,
             task=task,
             final_reward=final_reward,
@@ -510,7 +516,7 @@ class AgentModeDaemon:
                 if rollout.rollout_id in self._completed_rollouts_v0:
                     # Already processed, skip
                     continue
-                if isinstance(rollout, RolloutV2):
+                if isinstance(rollout, Rollout):
                     rollout = await self._validate_data_v1(rollout)
                 else:
                     self._validate_data(rollout)
@@ -736,7 +742,7 @@ class AgentModeDaemon:
         # This implementation assumes that `set_up_data_and_server` is called
         # for each new run, effectively starting a fresh batch.
 
-    def _fillna_reward(self, rollout: Rollout):
+    def _fillna_reward(self, rollout: RolloutLegacy):
         if rollout.final_reward is None:
             if self.reward_fillna_value is not None:  # type: ignore
                 final_reward = self.reward_fillna_value

agentlightning/verl/dataset.py

@@ -9,6 +9,11 @@ from verl.utils.dataset.rl_dataset import RLHFDataset
 
 from agentlightning.types import Dataset
 
+__all__ = [
+    "AgentDataset",
+    "LoadedDataset",
+]
+
 
 class AgentDataset(RLHFDataset):
 

agentlightning/verl/entrypoint.py

@@ -17,6 +17,12 @@ from agentlightning.types import Dataset
 from .dataset import AgentDataset, LoadedDataset
 from .trainer import AgentLightningTrainer
 
+__all__ = [
+    "main",
+    "run_ppo",
+    "TaskRunner",
+]
+
 
 @hydra.main(config_path="pkg://agentlightning/verl", config_name="config", version_base=None)
 def main(config):

agentlightning/verl/trainer.py

@@ -33,12 +33,16 @@ from verl.trainer.ppo.ray_trainer import (
 from verl.utils.metric import reduce_metrics
 from verl.utils.tracking import Tracking
 
-from agentlightning.adapter import BaseTraceTripletAdapter, TraceAdapter
+from agentlightning.adapter import TraceAdapter, TraceToTripletBase
 from agentlightning.llm_proxy import LLMProxy
 from agentlightning.store.base import LightningStore
 
 from .daemon import AgentModeDaemon
 
+__all__ = [
+    "AgentLightningTrainer",
+]
+
 
 @contextmanager
 def _timer(name: str, timing_raw: Dict[str, float]):
@@ -291,8 +295,8 @@ class AgentLightningTrainer(RayPPOTrainer):
         self._load_checkpoint()
 
         assert self.async_rollout_mode, "If agent mode is enabled, async server must be enabled"
-        if self.adapter is not None and not isinstance(self.adapter, BaseTraceTripletAdapter):
-            raise ValueError("Adapter must be a BaseTraceTripletAdapter for currently VERL implementation.")
+        if self.adapter is not None and not isinstance(self.adapter, TraceToTripletBase):
+            raise ValueError("Adapter must be a TraceToTripletBase for currently VERL implementation.")
         self.agent_mode_daemon = AgentModeDaemon(
             self.config.agentlightning.port,
             self.config.actor_rollout_ref.rollout.n,

docs/algorithm-zoo/apo.md

@@ -0,0 +1,25 @@
+# APO
+
+!!! tip "Shortcut"
+
+    You can use the shortcut `agl.APO(...)` to create an APO instance.
+
+    ```python
+    import agentlightning as agl
+
+    agl.APO(...)
+    ```
+
+## Installation
+
+```bash
+pip install agentlightning[apo]
+```
+
+## Tutorials Using APO
+
+TBD
+
+## References
+
+::: agentlightning.algorithm.apo

docs/algorithm-zoo/index.md

@@ -0,0 +1,10 @@
+# Algorithm Zoo
+
+AgentLightning includes several popular and frequently requested algorithms in its built-in library, allowing agent developers to use them directly. These algorithms are designed to be compatible with most agent scenarios.
+
+For customizing algorithms, see [Algorithm-side References](../reference/algorithm.md).
+
+| Algorithm | Optimizing Resources | Description |
+| --------- | ------------------- | ----------- |
+| [APO](./apo.md) | [PromptTemplate][agentlightning.PromptTemplate] | Automatic Prompt Optimization (APO) algorithm using textual gradients and beam search. |
+| [VERL](./verl.md) | [LLM][agentlightning.LLM] | Reinforcement Learning with [VERL framework](https://github.com/volcengine/verl). |

docs/algorithm-zoo/verl.md

@@ -0,0 +1,37 @@
+# VERL
+
+!!! tip "Shortcut"
+
+    You can use the shortcut `agl.VERL(...)` to create a VERL instance.
+
+    ```python
+    import agentlightning as agl
+
+    agl.VERL(...)
+    ```
+
+!!! warning "Customization note"
+
+    Customization of VERL is not supported as of current version. We recommend copying the source code from VERL and modifying it as needed to suit your requirements.
+
+## Installation
+
+```bash
+pip install agentlightning[verl]
+```
+
+!!! warning
+
+    For best results, follow the steps in the [installation guide](../quickstart/installation.md) to set up VERL and its dependencies. Installing VERL directly with `pip install agentlightning[verl]` can cause issues unless you already have a compatible version of PyTorch installed.
+
+## Tutorials Using VERL
+
+TBD
+
+## References - Entrypoint
+
+::: agentlightning.algorithm.verl
+
+## References - Implementation
+
+::: agentlightning.verl

docs/deep-dive/birds-eye-view.md

@@ -156,7 +156,7 @@ sequenceDiagram
 
 The Adapter is a component used by the Algorithm to transform raw data from the Store into a format suitable for learning. Runners stream raw spans into the Store during execution. Later, the Algorithm queries these spans and uses an Adapter to convert them into structured data, like training examples for a reinforcement learning model.
 
-For instance, the `TraceTripletAdapter` processes OpenTelemetry spans to create `(prompt, response, reward)` triplets, which are the fundamental data structure for many RL fine-tuning algorithms.
+For instance, the `TracerTraceToTriplet` processes OpenTelemetry spans to create `(prompt, response, reward)` triplets, which are the fundamental data structure for many RL fine-tuning algorithms.
 
 ```mermaid
 flowchart LR
@@ -237,7 +237,7 @@ flowchart TD
     %% === Left side: Algorithm domain ===
     subgraph L["Algorithm Side"]
         Algorithm["Algorithm<br>(no default)"]
-        Adapter["Adapter<br>(TraceTripletAdapter*)"]
+        Adapter["Adapter<br>(TracerTraceToTriplet*)"]
         LLMProxy["LLM Proxy<br>(no default)"]
         Algorithm -.injects.-> Adapter
         Algorithm -.injects.-> LLMProxy
@@ -257,7 +257,7 @@ flowchart TD
 
     %% === Right side: Runner side ===
     subgraph R["Runner Side"]
-        Runner["Runner<br>(AgentRunnerV2* default)"]
+        Runner["Runner<br>(LitAgentRunner* default)"]
         Tracer["Tracer<br>(AgentOpsTracer*)"]
         Hooks["Hooks (empty default)"]
         Agent["Agent<br>(LitAgent*)"]
@@ -305,7 +305,7 @@ In Agent-lightning, the environment is implicit in the agent’s workflow, which
 3. Querying the spans generated, extracting triplets, and converting them into a format that the underlying RL library can consume;
 4. Updating the language model based on the learning signals.
 
-In the VERL integration, the algorithm launches a chat completion endpoint using `vLLM` and wraps training with `FSDP` for distributed optimization. It enqueues tasks from the dataset. After rollouts finish, it queries spans and converts them to triplets with `TraceTripletAdapter`. VERL’s native training loop then consumes these triplets to update model weights. The workflow can be summarized in the following diagram.
+In the VERL integration, the algorithm launches a chat completion endpoint using `vLLM` and wraps training with `FSDP` for distributed optimization. It enqueues tasks from the dataset. After rollouts finish, it queries spans and converts them to triplets with `TracerTraceToTriplet`. VERL’s native training loop then consumes these triplets to update model weights. The workflow can be summarized in the following diagram.
 
 ```mermaid
 sequenceDiagram
@@ -313,7 +313,7 @@ sequenceDiagram
     participant vLLM as vLLM Chat<br>Completion Endpoint
     participant FSDP as FSDP / Megatron<br>Weights Optimizer
     participant Algo as Algorithm<br>Main Controller<br>(Main Process)
-    participant Adapter as TraceTripletAdapter
+    participant Adapter as TracerTraceToTriplet
     participant LLMProxy as LLM Proxy
     participant Store as LightningStore
     participant Runner as Runner + Agent

docs/index.md

@@ -14,7 +14,6 @@ Agent Lightning is the absolute trainer to light up AI agents.
 - [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
 - [Join our Discord community](https://discord.gg/RYk7CdvDR7) - Connect with other users and contributors
 
 

docs/quickstart/getting-started.md

@@ -66,7 +66,7 @@ from agentlightning.trainer import Trainer
 
 agent = SimpleAgent()
 trainer = Trainer(n_workers=2)  # Create 2 parallel workers
-trainer.fit(agent, "http://127.0.0.1:9997")
+trainer.fit_v0(agent, "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.

docs/reference/agent.md

@@ -0,0 +1,45 @@
+# Agent Developer APIs
+
+## Customizing Agents - Decorators
+
+!!! tip
+
+    These are convenient helpers for creating agents from functions. First-time users are recommended to use these decorators to create agents.
+
+::: agentlightning.rollout
+
+!!! warning
+
+    The following two decorators are implementations of [`agentlightning.rollout`][agentlightning.rollout]. They are not recommended for new users.
+
+::: agentlightning.llm_rollout
+
+::: agentlightning.prompt_rollout
+
+## Class-based Agents
+
+::: agentlightning.LitAgent
+
+## Emitter
+
+::: agentlightning.emit_reward
+
+::: agentlightning.emit_message
+
+::: agentlightning.emit_object
+
+::: agentlightning.emit_exception
+
+## Reward Helpers
+
+::: agentlightning.find_final_reward
+
+::: agentlightning.find_reward_spans
+
+::: agentlightning.get_reward_value
+
+::: agentlightning.is_reward_span
+
+## Legacy Emitter Decorators
+
+::: agentlightning.reward.reward

docs/reference/algorithm.md

@@ -0,0 +1,42 @@
+## Algorithm-side References
+
+!!! note
+
+    This reference covers APIs that are designed to be used at "Algorithm Side".
+    For built-in algorithms, see [Algorithm Zoo](../algorithm-zoo/index.md).
+
+## Base Class and Decorators
+
+::: agentlightning.BaseAlgorithm
+
+::: agentlightning.algo
+
+## Fast Algorithms (for Debugging)
+
+::: agentlightning.FastAlgorithm
+
+::: agentlightning.Baseline
+
+## Adapter
+
+::: agentlightning.TraceAdapter
+
+::: agentlightning.Adapter
+
+::: agentlightning.TraceToTripletBase
+
+::: agentlightning.TracerTraceToTriplet
+
+::: agentlightning.LlmProxyTraceToTriplet
+
+::: agentlightning.TraceToMessages
+
+## LLM Proxy
+
+::: agentlightning.LLMProxy
+
+::: agentlightning.llm_proxy.ModelConfig
+
+::: agentlightning.llm_proxy.LightningSpanExporter
+
+::: agentlightning.llm_proxy.AddReturnTokenIds

docs/reference/cli.md

@@ -0,0 +1,89 @@
+# Command Line Interface
+
+<!-- TODO: This document should be auto-generated. -->
+
+!!! warning
+
+    This document is a work in progress and might not be updated with the latest changes.
+    Try to use `agl -h` to get the latest help message.
+
+!!! tip
+
+    Agent-lightning also provides utilities to help you build your own CLI for [LitAgent][agentlightning.LitAgent] and [Trainer][agentlightning.Trainer]. See [Trainer](./trainer.md) for references.
+
+## agl
+
+```text
+usage: agl [-h] {vllm,store,agentops}
+
+Agent Lightning CLI entry point.
+
+Available subcommands:
+  vllm      Run the vLLM CLI with Agent Lightning instrumentation.
+  store     Run a LightningStore server.
+  agentops  Start the AgentOps server manager.
+
+positional arguments:
+  {vllm,store,agentops}
+                        Subcommand to run.
+
+options:
+  -h, --help            show this help message and exit
+```
+
+## agl vllm
+
+Agent-lightning's instrumented vLLM CLI.
+
+```text
+usage: agl vllm [-h] [-v] {chat,complete,serve,bench,collect-env,run-batch} ...
+
+vLLM CLI
+
+positional arguments:
+  {chat,complete,serve,bench,collect-env,run-batch}
+    chat                Generate chat completions via the running API server.
+    complete            Generate text completions based on the given prompt via the running API server.
+    collect-env         Start collecting environment information.
+    run-batch           Run batch prompts and write results to file.
+
+options:
+  -h, --help            show this help message and exit
+  -v, --version         show program's version number and exit
+
+For full list:            vllm [subcommand] --help=all
+For a section:            vllm [subcommand] --help=ModelConfig    (case-insensitive)
+For a flag:               vllm [subcommand] --help=max-model-len  (_ or - accepted)
+Documentation:            https://docs.vllm.ai
+```
+
+## agl store
+
+Agent-lightning's LightningStore CLI. Use it to start an independent LightningStore server.
+
+Currently the store data are stored in memory and will be lost when the server is stopped.
+
+```text
+usage: agl store [-h] [--port PORT]
+
+Run a LightningStore server
+
+options:
+  -h, --help   show this help message and exit
+  --port PORT  Port to run the server on
+```
+
+## agl agentops
+
+Start a mock AgentOps server to bypass the online service of AgentOps.
+
+```text
+usage: agl agentops [-h] [--daemon] [--port PORT]
+
+Start AgentOps server
+
+options:
+  -h, --help   show this help message and exit
+  --daemon     Run server as a daemon
+  --port PORT  Port to run the server on
+```

docs/reference/core.md

@@ -1,51 +0,0 @@
-# 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/instrumentation.md

@@ -0,0 +1,21 @@
+# Instrumentation API
+
+::: agentlightning.instrumentation.instrument_all
+
+::: agentlightning.instrumentation.uninstrument_all
+
+## AgentOps LangChain
+
+::: agentlightning.instrumentation.agentops_langchain
+
+## AgentOps
+
+::: agentlightning.instrumentation.agentops
+
+## LiteLLM
+
+::: agentlightning.instrumentation.litellm
+
+## vLLM
+
+::: agentlightning.instrumentation.vllm

docs/reference/rl.md

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

docs/reference/runner.md

@@ -0,0 +1,19 @@
+# Runner-side References
+
+!!! note
+
+    This reference covers APIs that are designed to be used at "Runner Side".
+
+## Runners
+
+::: agentlightning.LitAgentRunner
+
+::: agentlightning.BaseRunner
+
+## Tracer
+
+::: agentlightning.AgentOpsTracer
+
+::: agentlightning.OtelTracer
+
+::: agentlightning.BaseTracer

docs/reference/store.md

@@ -0,0 +1,15 @@
+# Store References
+
+::: agentlightning.LightningStore
+
+## Store Implementations
+
+::: agentlightning.InMemoryLightningStore
+
+## Client-Server and Thread-safe Wrappers
+
+::: agentlightning.LightningStoreServer
+
+::: agentlightning.LightningStoreClient
+
+::: agentlightning.LightningStoreThreaded

docs/reference/trainer.md

@@ -0,0 +1,25 @@
+# Agent-lightning Trainer
+
+::: agentlightning.Trainer
+
+::: agentlightning.build_component
+
+## Execution Strategy
+
+::: agentlightning.ExecutionStrategy
+
+::: agentlightning.ClientServerExecutionStrategy
+
+::: agentlightning.SharedMemoryExecutionStrategy
+
+## Events
+
+::: agentlightning.ExecutionEvent
+
+::: agentlightning.ThreadingEvent
+
+::: agentlightning.MultiprocessingEvent
+
+## CLI Builder
+
+::: agentlightning.lightning_cli

docs/reference/types.md

@@ -0,0 +1,73 @@
+# Type References
+
+## Core Types
+
+::: agentlightning.Triplet
+
+::: agentlightning.TaskInput
+
+::: agentlightning.RolloutRawResult
+
+::: agentlightning.RolloutMode
+
+::: agentlightning.GenericResponse
+
+::: agentlightning.ParallelWorkerBase
+
+::: agentlightning.Dataset
+
+::: agentlightning.AttemptStatus
+
+::: agentlightning.RolloutStatus
+
+::: agentlightning.RolloutConfig
+
+::: agentlightning.Rollout
+
+::: agentlightning.Attempt
+
+::: agentlightning.AttemptedRollout
+
+::: agentlightning.Hook
+
+## Resources
+
+::: agentlightning.Resource
+
+::: agentlightning.LLM
+
+::: agentlightning.ProxyLLM
+
+::: agentlightning.PromptTemplate
+
+::: agentlightning.ResourceUnion
+
+::: agentlightning.NamedResources
+
+::: agentlightning.ResourcesUpdate
+
+## Traces
+
+::: agentlightning.AttributeValue
+
+::: agentlightning.Attributes
+
+::: agentlightning.TraceState
+
+::: agentlightning.SpanContext
+
+::: agentlightning.TraceStatus
+
+::: agentlightning.Event
+
+::: agentlightning.Link
+
+::: agentlightning.Resource
+
+::: agentlightning.Span
+
+::: agentlightning.SpanNames
+
+::: agentlightning.SpanAttributeNames
+
+::: agentlightning.SpanLike

examples/.gitignore

@@ -8,3 +8,4 @@ agentops.log
 unsloth/models/
 unsloth/unsloth_compiled_cache/
 unsloth/unsloth_training_checkpoints/
+apo/pomltrace/

examples/apo/apo.py

@@ -1,9 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-"""This is the APO example written in the legacy client-server style (agent-lightning v0.1).
-
-New users should refer to the `examples/apo/apo.py` for the modern APO example.
-"""
+"""This is the APO sample with both rollout and algo in one file."""
 
 from typing import List, Optional
 
@@ -11,10 +8,10 @@ from openai import AsyncOpenAI
 from rich.console import Console
 
 from agentlightning import Trainer, configure_logger
-from agentlightning.algorithm.base import algo
+from agentlightning.algorithm import algo
 from agentlightning.litagent.decorator import rollout
 from agentlightning.reward import find_final_reward
-from agentlightning.store.base import LightningStore
+from agentlightning.store import LightningStore
 from agentlightning.types import NamedResources, PromptTemplate, Span
 
 console = Console()
@@ -130,4 +127,4 @@ Return only a number between 0 and 1. No text, punctuation, or explanation."""
 if __name__ == "__main__":
     configure_logger()
     trainer = Trainer(n_workers=1, algorithm=apo_algorithm)
-    trainer.fit_v2(apo_rollout)
+    trainer.fit(apo_rollout)

examples/apo/apo_debug.py

@@ -9,8 +9,8 @@ from typing import cast
 from apo import apo_rollout
 
 from agentlightning import Trainer, configure_logger
-from agentlightning.runner import AgentRunnerV2
-from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.runner import LitAgentRunner
+from agentlightning.store import InMemoryLightningStore
 from agentlightning.tracer import OtelTracer
 from agentlightning.types import Dataset, PromptTemplate
 
@@ -26,7 +26,7 @@ async def debug_with_runner():
     # Tracer is used to record the events (spans) in background during the agent's execution.
     # If you don't need any tracing functionality yet, you can use a dummy OtelTracer.
     tracer = OtelTracer()
-    runner = AgentRunnerV2[str](tracer)
+    runner = LitAgentRunner[str](tracer)
 
     # You also need a store here to store the data collected.
     store = InMemoryLightningStore()

examples/apo/legacy_apo_client.py

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

examples/apo/room_selector.py

@@ -0,0 +1,362 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import json
+import traceback
+from typing import List, Optional, Tuple, TypedDict, cast
+
+from openai import OpenAI
+from openai.types.chat import (
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionMessageFunctionToolCallParam,
+    ChatCompletionMessageParam,
+    ChatCompletionToolMessageParam,
+    ChatCompletionToolParam,
+)
+from pydantic import BaseModel, Field
+from rich.console import Console
+
+from agentlightning.adapter import TraceToMessages
+from agentlightning.litagent import rollout
+from agentlightning.reward import find_final_reward
+from agentlightning.runner import LitAgentRunner
+from agentlightning.store import InMemoryLightningStore
+from agentlightning.tracer.agentops import AgentOpsTracer
+from agentlightning.types import Dataset, PromptTemplate
+
+console = Console()
+
+
+class JudgeResponse(BaseModel):
+    reason: str = Field(description="The reason for the score. No more than 100 characters.")
+    score: float = Field(description="The score for the match on a 0-1 scale. Be critical.")
+
+
+class Room(TypedDict):
+    id: str
+    capacity: int
+    equipment: List[str]
+    accessible: bool
+    distance_m: int
+    booked: List[Tuple[str, str, int]]
+
+
+class RoomStatus(Room):
+    free: bool
+
+
+class AvailableRooms(TypedDict):
+    rooms: List[RoomStatus]
+
+
+class RoomRequirement(TypedDict):
+    date: str
+    time: str
+    duration_min: int
+    attendees: int
+    needs: List[str]
+    accessible_required: bool
+
+
+class RoomSelectionTask(TypedDict):
+    id: str
+    task_input: RoomRequirement
+    expected_choice: str
+
+
+TOOL_DEFINITIONS: List[ChatCompletionToolParam] = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_rooms_and_availability",
+            "description": "Return meeting rooms with capacity, equipment, accessibility, distance, and booked time slots.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "date": {"type": "string", "description": "YYYY-MM-DD"},
+                    "time": {"type": "string", "description": "HH:MM 24h local"},
+                    "duration_min": {"type": "integer", "description": "Meeting duration minutes"},
+                },
+                "required": ["date", "time", "duration_min"],
+            },
+        },
+    },
+]
+
+
+def prompt_template_baseline() -> PromptTemplate:
+    return PromptTemplate(
+        template="Find a room on {date} at {time} for {duration_min} minutes, {attendees} attendees. Needs: {needs}. Accessible required: {accessible_required}",
+        engine="f-string",
+    )
+
+
+def room_selection_grader(client: OpenAI, final_message: Optional[str], expected_choice: str) -> float:
+    judge_prompt = (
+        f"You are a strict grader of exact room choice."
+        f"Task output:\n{final_message}\n\n"
+        f"Task expected answer:\n{expected_choice}\n\n"
+        f"Score the match on a 0-1 scale. Be critical.\n"
+        f"Bear in mind that the score can be partially correct (between 0 and 1)."
+    )
+    judge = client.chat.completions.parse(
+        model="gpt-4.1-mini",
+        messages=[
+            {"role": "user", "content": judge_prompt},
+        ],
+        response_format=JudgeResponse,
+        temperature=0.0,
+    )
+
+    judge_result = judge.choices[0].message.content
+    console.print(f"[bold yellow]=== Judge ===[/bold yellow]")
+    console.print(judge_result)
+
+    judge_result_parsed = JudgeResponse.model_validate_json(judge_result)  # type: ignore
+
+    console.print(f"[bold yellow]=== Judge Score ===[/bold yellow]")
+    console.print(judge_result_parsed.score)
+    return judge_result_parsed.score
+
+
+@rollout
+def room_selector(task: RoomSelectionTask, prompt_template: PromptTemplate) -> float:
+    """An agent to select a room based on the given requirements.
+
+    Oracle System Prompt (works with 100% accuracy with gpt-5 mini low reasoning effort):
+
+        You are a scheduling assistant.
+        Hard constraints: free for slot, capacity >= attendees, includes all required equipment,
+        accessible==True if requested.
+        Tie-break scoring (lower is better):
+        1) capacity_slack = capacity - attendees (minimize)
+        2) extra_equipment = provided_equipment_count - required_equipment_count (minimize)
+        3) distance_m (minimize)
+        4) fewer total booked blocks that day (minimize)
+        Return No Room if no room is found that satisfies the constraints.
+        Return strictly:
+        final_choice: <ROOM_ID>
+        reason: <one line stating the decisive criteria>
+
+    Oracle User Prompt Template:
+
+        Find a room on {task_input['date']} at {task_input['time']} for {task_input['duration_min']} minutes,
+        {task_input['attendees']} attendees. Needs: {', '.join(task_input['needs']) or 'none'}.
+        Accessible required: {task_input['accessible_required']}
+
+    The current implementation greatly simply the oracle prompt and prompt template is provided by a parameter.
+    The prompt template should be tuned by Agent-lightning's APO algorithm.
+    It also should work with a very small model like gpt-4.1-nano.
+    """
+
+    client = OpenAI()
+    model = "gpt-4.1-nano"
+
+    user_message = prompt_template.format(**task["task_input"])
+
+    messages: List[ChatCompletionMessageParam] = [
+        {"role": "system", "content": "You are a scheduling assistant."},
+        {
+            "role": "user",
+            "content": user_message,
+        },
+    ]
+
+    console.print(f"[bold yellow]=== User Message ===[/bold yellow]")
+    console.print(user_message)
+
+    resp = client.chat.completions.create(
+        model=model,
+        messages=messages,
+        tools=TOOL_DEFINITIONS,
+        tool_choice="auto",
+        # Minimize the randomness
+        temperature=0.0,
+        # Uncomment for gpt-5
+        # reasoning_effort="low",
+    )
+
+    console.print(f"[bold yellow]=== Assistant Message ===[/bold yellow]")
+    console.print(resp.choices[0].message)
+
+    # Parse and process the tool calls
+    tool_calls = resp.choices[0].message.tool_calls
+    if tool_calls:
+
+        tool_call_params: List[ChatCompletionMessageFunctionToolCallParam] = []
+        tool_results: List[ChatCompletionToolMessageParam] = []
+        for tc in tool_calls:
+            if tc.type != "function":
+                raise ValueError(f"Tool call is not a function: {tc}")
+            if tc.function.name != "get_rooms_and_availability":
+                raise ValueError(f"Tool call is not get_rooms_and_availability: {tc}")
+            tool_call_params.append(
+                ChatCompletionMessageFunctionToolCallParam(
+                    id=tc.id,
+                    type="function",
+                    function={"name": tc.function.name, "arguments": tc.function.arguments},
+                )
+            )
+            args = json.loads(tc.function.arguments)
+            try:
+                tool_output = get_rooms_and_availability(args["date"], args["time"], args["duration_min"])
+            except Exception as e:
+                tool_output = {
+                    "error": str(e),
+                    "traceback": traceback.format_exc(),
+                }
+            console.print(f"[bold yellow]=== Tool Message ===[/bold yellow]")
+            console.print(tool_output)
+            tool_results.append(
+                ChatCompletionToolMessageParam(
+                    role="tool",
+                    tool_call_id=tc.id,
+                    content=json.dumps(tool_output),
+                )
+            )
+
+        # Update the messages for hte next call
+        messages.append(
+            ChatCompletionAssistantMessageParam(
+                role="assistant",
+                content=resp.choices[0].message.content,
+                tool_calls=tool_call_params,
+            )
+        )
+        messages.extend(tool_results)
+
+        next_resp = client.chat.completions.create(
+            model=model,
+            messages=messages,
+            # Minimize the randomness
+            temperature=0.0,
+        )
+        console.print(f"[bold yellow]=== Final Assistant Message ===[/bold yellow]")
+        console.print(next_resp.choices[0].message.content)
+        final_message = next_resp.choices[0].message.content
+
+    else:
+        final_message = resp.choices[0].message.content
+
+    return room_selection_grader(client, final_message, task["expected_choice"])
+
+
+# Local tool database (there might be multiple plausible fits)
+ROOMS: List[Room] = [
+    {
+        "id": "Orion",
+        "capacity": 4,
+        "equipment": ["tv", "whiteboard"],
+        "accessible": True,
+        "distance_m": 12,
+        "booked": [("2025-10-13", "10:00", 60), ("2025-10-13", "15:00", 30)],
+    },
+    {
+        "id": "Lyra",
+        "capacity": 10,
+        "equipment": ["projector", "whiteboard", "confphone"],
+        "accessible": True,
+        "distance_m": 30,
+        "booked": [("2025-10-13", "09:30", 30), ("2025-10-13", "11:00", 60)],
+    },
+    {
+        "id": "Vega",
+        "capacity": 6,
+        "equipment": ["tv"],
+        "accessible": False,
+        "distance_m": 22,
+        "booked": [("2025-10-13", "14:00", 60)],
+    },
+    {
+        "id": "Nova",
+        "capacity": 12,
+        "equipment": ["ledwall", "whiteboard", "confphone"],
+        "accessible": True,
+        "distance_m": 45,
+        "booked": [],
+    },
+    {
+        "id": "Quark",
+        "capacity": 8,
+        "equipment": ["projector", "whiteboard"],
+        "accessible": False,
+        "distance_m": 18,
+        "booked": [("2025-10-13", "10:30", 30)],
+    },
+    # Two extra to create harder ties
+    {
+        "id": "Atlas",
+        "capacity": 6,
+        "equipment": ["projector", "whiteboard"],
+        "accessible": True,
+        "distance_m": 10,
+        "booked": [("2025-10-13", "09:00", 30), ("2025-10-13", "13:30", 30)],
+    },
+    {
+        "id": "Pulse",
+        "capacity": 8,
+        "equipment": ["tv", "whiteboard", "confphone"],
+        "accessible": True,
+        "distance_m": 8,
+        "booked": [("2025-10-13", "16:30", 30)],
+    },
+]
+
+
+def overlaps(start: str, dur: int, other_start: str, other_dur: int) -> bool:
+    def tmin(t: str):
+        return int(t[:2]) * 60 + int(t[3:])
+
+    a0, a1 = tmin(start), tmin(start) + dur
+    b0, b1 = tmin(other_start), tmin(other_start) + other_dur
+    return max(a0, b0) < min(a1, b1)
+
+
+def get_rooms_and_availability(date: str, time_str: str, duration_min: int) -> AvailableRooms:
+    avail: List[RoomStatus] = []
+    for r in ROOMS:
+        free = all(
+            not (b_date == date and overlaps(time_str, duration_min, b_time, b_dur))
+            for (b_date, b_time, b_dur) in r["booked"]
+        )
+        item: RoomStatus = {
+            **r,
+            "free": free,
+        }
+        avail.append(item)
+    return {"rooms": avail}
+
+
+def load_room_tasks() -> Dataset[RoomSelectionTask]:
+    tasks: List[RoomSelectionTask] = []
+    for line in open("room_tasks.jsonl"):
+        task = json.loads(line)
+        tasks.append(RoomSelectionTask(**task))
+    return cast(Dataset[RoomSelectionTask], tasks)
+
+
+async def debug_room_selector(limit: int = 1):
+    # Prepare all the components to run the agent
+    runner = LitAgentRunner[RoomSelectionTask](AgentOpsTracer())
+    store = InMemoryLightningStore()
+    prompt_template = prompt_template_baseline()
+    tasks = load_room_tasks()
+    with runner.run_context(agent=room_selector, store=store):
+        for task in tasks:
+            console.print("[bold green]=== Task ===[/bold green]", task, sep="\n")
+            # Run the agent
+            rollout = await runner.step(tasks[0], resources={"main_prompt": prompt_template})
+            # Get the spans and convert them to messages
+            # Useful for debugging and analysis
+            spans = await store.query_spans(rollout.rollout_id)
+            adapter = TraceToMessages()
+            messages = adapter.adapt(spans)
+            for message_idx, message in enumerate(messages):
+                console.print(f"[bold purple]=== Postmortem Message #{message_idx} ===[/bold purple]")
+                console.print(json.dumps(message))
+            reward = find_final_reward(spans)
+            console.print("[bold purple]=== Postmortem Reward ===[/bold purple]", reward, sep="\n")
+
+
+if __name__ == "__main__":
+    asyncio.run(debug_room_selector())

examples/apo/room_selector_apo.py

@@ -0,0 +1,69 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""This sample code demonstrates how to use an existing APO algorithm to tune the prompts."""
+
+import logging
+from typing import Tuple, cast
+
+from openai import AsyncOpenAI
+from room_selector import RoomSelectionTask, load_room_tasks, prompt_template_baseline, room_selector
+
+from agentlightning import Trainer, configure_logger
+from agentlightning.adapter import TraceToMessages
+from agentlightning.algorithm.apo import APO
+from agentlightning.types import Dataset
+
+
+def load_train_val_dataset() -> Tuple[Dataset[RoomSelectionTask], Dataset[RoomSelectionTask]]:
+    dataset_full = load_room_tasks()
+    train_split = len(dataset_full) // 2
+    dataset_train = [dataset_full[i] for i in range(train_split)]
+    dataset_val = [dataset_full[i] for i in range(train_split, len(dataset_full))]
+    return cast(Dataset[RoomSelectionTask], dataset_train), cast(Dataset[RoomSelectionTask], dataset_val)
+
+
+def setup_apo_logger(file_path: str = "apo.log") -> None:
+    """Dump a copy of all the logs produced by APO algorithm to a file."""
+
+    file_handler = logging.FileHandler(file_path)
+    file_handler.setLevel(logging.INFO)
+    formatter = logging.Formatter("%(asctime)s [%(levelname)s] (Process-%(process)d %(name)s)   %(message)s")
+    file_handler.setFormatter(formatter)
+    logging.getLogger("agentlightning.algorithm.apo").addHandler(file_handler)
+
+
+def main() -> None:
+    configure_logger()
+    setup_apo_logger()
+
+    openai_client = AsyncOpenAI()
+
+    algo = APO[RoomSelectionTask](
+        openai_client,
+        val_batch_size=10,
+        gradient_batch_size=4,
+        beam_width=2,
+        branch_factor=2,
+        beam_rounds=2,
+        _poml_trace=True,
+    )
+    trainer = Trainer(
+        algorithm=algo,
+        # Increase the number of runners to run more rollouts in parallel
+        n_runners=8,
+        # APO algorithm needs a baseline
+        # Set it either here or in the algo
+        initial_resources={
+            # The resource key can be arbitrary
+            "prompt_template": prompt_template_baseline()
+        },
+        # APO algorithm needs an adapter to process the traces produced by rollouts
+        # Use this adapter to convert spans to messages
+        adapter=TraceToMessages(),
+    )
+    dataset_train, dataset_val = load_train_val_dataset()
+    trainer.fit(agent=room_selector, train_dataset=dataset_train, val_dataset=dataset_val)
+
+
+if __name__ == "__main__":
+    main()

examples/apo/room_tasks.jsonl

@@ -0,0 +1,57 @@
+{"id": "s01", "task_input": {"date": "2025-10-13", "time": "16:30", "duration_min": 30, "attendees": 12, "needs": ["projector", "confphone"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s02", "task_input": {"date": "2025-10-13", "time": "14:30", "duration_min": 30, "attendees": 12, "needs": ["whiteboard", "confphone"], "accessible_required": true}, "expected_choice": "Nova"}
+{"id": "s03", "task_input": {"date": "2025-10-13", "time": "11:00", "duration_min": 60, "attendees": 10, "needs": ["projector", "whiteboard"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s04", "task_input": {"date": "2025-10-13", "time": "09:45", "duration_min": 30, "attendees": 8, "needs": ["projector"], "accessible_required": false}, "expected_choice": "Quark"}
+{"id": "s05", "task_input": {"date": "2025-10-13", "time": "12:15", "duration_min": 45, "attendees": 10, "needs": ["projector", "whiteboard"], "accessible_required": true}, "expected_choice": "Lyra"}
+{"id": "s06", "task_input": {"date": "2025-10-13", "time": "10:30", "duration_min": 20, "attendees": 4, "needs": ["whiteboard"], "accessible_required": true}, "expected_choice": "Atlas"}
+{"id": "s07", "task_input": {"date": "2025-10-13", "time": "16:30", "duration_min": 60, "attendees": 12, "needs": ["confphone", "whiteboard"], "accessible_required": true}, "expected_choice": "Nova"}
+{"id": "s08", "task_input": {"date": "2025-10-13", "time": "12:00", "duration_min": 30, "attendees": 8, "needs": ["confphone", "whiteboard"], "accessible_required": true}, "expected_choice": "Pulse"}
+{"id": "s09", "task_input": {"date": "2025-10-13", "time": "13:30", "duration_min": 30, "attendees": 10, "needs": ["projector", "whiteboard"], "accessible_required": true}, "expected_choice": "Lyra"}
+{"id": "s10", "task_input": {"date": "2025-10-13", "time": "12:30", "duration_min": 30, "attendees": 5, "needs": ["tv"], "accessible_required": true}, "expected_choice": "Pulse"}
+{"id": "s11", "task_input": {"date": "2025-10-13", "time": "10:30", "duration_min": 30, "attendees": 10, "needs": ["confphone", "whiteboard"], "accessible_required": true}, "expected_choice": "Lyra"}
+{"id": "s12", "task_input": {"date": "2025-10-13", "time": "11:30", "duration_min": 45, "attendees": 10, "needs": ["projector", "whiteboard"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s13", "task_input": {"date": "2025-10-13", "time": "15:00", "duration_min": 30, "attendees": 3, "needs": ["whiteboard"], "accessible_required": true}, "expected_choice": "Atlas"}
+{"id": "s14", "task_input": {"date": "2025-10-13", "time": "11:30", "duration_min": 30, "attendees": 4, "needs": ["whiteboard"], "accessible_required": false}, "expected_choice": "Orion"}
+{"id": "s15", "task_input": {"date": "2025-10-13", "time": "13:00", "duration_min": 30, "attendees": 3, "needs": [], "accessible_required": true}, "expected_choice": "Orion"}
+{"id": "s16", "task_input": {"date": "2025-10-13", "time": "13:00", "duration_min": 30, "attendees": 8, "needs": ["projector"], "accessible_required": false}, "expected_choice": "Quark"}
+{"id": "s17", "task_input": {"date": "2025-10-13", "time": "13:45", "duration_min": 30, "attendees": 5, "needs": ["tv", "whiteboard"], "accessible_required": true}, "expected_choice": "Pulse"}
+{"id": "s18", "task_input": {"date": "2025-10-13", "time": "12:45", "duration_min": 30, "attendees": 10, "needs": ["projector", "confphone"], "accessible_required": true}, "expected_choice": "Lyra"}
+{"id": "s19", "task_input": {"date": "2025-10-13", "time": "16:30", "duration_min": 30, "attendees": 6, "needs": ["whiteboard"], "accessible_required": true}, "expected_choice": "Atlas"}
+{"id": "s20", "task_input": {"date": "2025-10-13", "time": "10:30", "duration_min": 45, "attendees": 4, "needs": ["projector", "whiteboard", "confphone"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s21", "task_input": {"date": "2025-10-13", "time": "12:00", "duration_min": 30, "attendees": 3, "needs": ["tv"], "accessible_required": true}, "expected_choice": "Orion"}
+{"id": "s22", "task_input": {"date": "2025-10-13", "time": "16:00", "duration_min": 45, "attendees": 8, "needs": ["projector", "whiteboard"], "accessible_required": false}, "expected_choice": "Quark"}
+{"id": "s23", "task_input": {"date": "2025-10-13", "time": "11:45", "duration_min": 30, "attendees": 6, "needs": [], "accessible_required": true}, "expected_choice": "Atlas"}
+{"id": "s24", "task_input": {"date": "2025-10-13", "time": "12:15", "duration_min": 30, "attendees": 10, "needs": ["whiteboard"], "accessible_required": true}, "expected_choice": "Lyra"}
+{"id": "s25", "task_input": {"date": "2025-10-13", "time": "15:30", "duration_min": 30, "attendees": 10, "needs": ["projector", "confphone"], "accessible_required": true}, "expected_choice": "Lyra"}
+{"id": "s26", "task_input": {"date": "2025-10-13", "time": "14:30", "duration_min": 60, "attendees": 12, "needs": ["projector", "whiteboard", "confphone"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s27", "task_input": {"date": "2025-10-13", "time": "13:45", "duration_min": 30, "attendees": 12, "needs": ["projector", "whiteboard", "confphone"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s28", "task_input": {"date": "2025-10-13", "time": "14:00", "duration_min": 60, "attendees": 4, "needs": ["tv", "whiteboard"], "accessible_required": true}, "expected_choice": "Orion"}
+{"id": "s29", "task_input": {"date": "2025-10-13", "time": "14:30", "duration_min": 30, "attendees": 10, "needs": ["whiteboard", "confphone"], "accessible_required": true}, "expected_choice": "Lyra"}
+{"id": "s30", "task_input": {"date": "2025-10-13", "time": "12:00", "duration_min": 60, "attendees": 4, "needs": ["tv"], "accessible_required": false}, "expected_choice": "Orion"}
+{"id": "s31", "task_input": {"date": "2025-10-13", "time": "15:00", "duration_min": 30, "attendees": 9, "needs": ["tv", "whiteboard"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s32", "task_input": {"date": "2025-10-13", "time": "10:45", "duration_min": 30, "attendees": 8, "needs": ["projector"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s33", "task_input": {"date": "2025-10-13", "time": "10:00", "duration_min": 30, "attendees": 3, "needs": ["tv"], "accessible_required": true}, "expected_choice": "Pulse"}
+{"id": "s34", "task_input": {"date": "2025-10-13", "time": "09:00", "duration_min": 30, "attendees": 12, "needs": ["projector", "whiteboard"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s35", "task_input": {"date": "2025-10-13", "time": "09:15", "duration_min": 30, "attendees": 6, "needs": ["projector", "whiteboard"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s36", "task_input": {"date": "2025-10-13", "time": "12:00", "duration_min": 30, "attendees": 9, "needs": ["tv"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s37", "task_input": {"date": "2025-10-13", "time": "13:45", "duration_min": 30, "attendees": 4, "needs": ["tv", "whiteboard"], "accessible_required": true}, "expected_choice": "Orion"}
+{"id": "s38", "task_input": {"date": "2025-10-13", "time": "11:30", "duration_min": 30, "attendees": 6, "needs": ["whiteboard"], "accessible_required": true}, "expected_choice": "Atlas"}
+{"id": "s39", "task_input": {"date": "2025-10-13", "time": "16:00", "duration_min": 30, "attendees": 8, "needs": ["confphone", "whiteboard"], "accessible_required": true}, "expected_choice": "Pulse"}
+{"id": "s40", "task_input": {"date": "2025-10-13", "time": "14:15", "duration_min": 30, "attendees": 6, "needs": ["tv"], "accessible_required": false}, "expected_choice": "Pulse"}
+{"id": "s41", "task_input": {"date": "2025-10-13", "time": "12:30", "duration_min": 60, "attendees": 10, "needs": ["projector", "whiteboard"], "accessible_required": true}, "expected_choice": "Lyra"}
+{"id": "s42", "task_input": {"date": "2025-10-13", "time": "15:30", "duration_min": 30, "attendees": 4, "needs": ["whiteboard"], "accessible_required": true}, "expected_choice": "Orion"}
+{"id": "s43", "task_input": {"date": "2025-10-13", "time": "10:30", "duration_min": 30, "attendees": 12, "needs": ["confphone", "whiteboard"], "accessible_required": true}, "expected_choice": "Nova"}
+{"id": "s44", "task_input": {"date": "2025-10-13", "time": "13:30", "duration_min": 30, "attendees": 12, "needs": ["projector"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s45", "task_input": {"date": "2025-10-13", "time": "10:30", "duration_min": 45, "attendees": 6, "needs": ["whiteboard", "projector"], "accessible_required": true}, "expected_choice": "Atlas"}
+{"id": "s46", "task_input": {"date": "2025-10-13", "time": "09:30", "duration_min": 30, "attendees": 10, "needs": ["projector", "confphone"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s47", "task_input": {"date": "2025-10-13", "time": "11:30", "duration_min": 30, "attendees": 10, "needs": ["projector", "whiteboard"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s48", "task_input": {"date": "2025-10-13", "time": "09:30", "duration_min": 60, "attendees": 10, "needs": ["projector", "confphone"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s49", "task_input": {"date": "2025-10-13", "time": "12:00", "duration_min": 30, "attendees": 12, "needs": ["confphone", "whiteboard"], "accessible_required": true}, "expected_choice": "Nova"}
+{"id": "s50", "task_input": {"date": "2025-10-13", "time": "15:00", "duration_min": 30, "attendees": 3, "needs": ["whiteboard", "confphone"], "accessible_required": true}, "expected_choice": "Pulse"}
+{"id": "s51", "task_input": {"date": "2025-10-13", "time": "14:30", "duration_min": 30, "attendees": 6, "needs": ["tv"], "accessible_required": false}, "expected_choice": "Pulse"}
+{"id": "s52", "task_input": {"date": "2025-10-13", "time": "11:00", "duration_min": 30, "attendees": 6, "needs": ["projector"], "accessible_required": true}, "expected_choice": "Atlas"}
+{"id": "s53", "task_input": {"date": "2025-10-13", "time": "13:30", "duration_min": 45, "attendees": 6, "needs": ["projector"], "accessible_required": true}, "expected_choice": "Lyra"}
+{"id": "s54", "task_input": {"date": "2025-10-13", "time": "10:00", "duration_min": 30, "attendees": 12, "needs": ["projector", "whiteboard", "confphone"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s55", "task_input": {"date": "2025-10-13", "time": "16:30", "duration_min": 30, "attendees": 8, "needs": ["tv"], "accessible_required": true}, "expected_choice": "No Room"}
+{"id": "s56", "task_input": {"date": "2025-10-13", "time": "09:00", "duration_min": 30, "attendees": 4, "needs": ["whiteboard"], "accessible_required": true}, "expected_choice": "Orion"}
+{"id": "s57", "task_input": {"date": "2025-10-13", "time": "10:30", "duration_min": 45, "attendees": 8, "needs": ["projector"], "accessible_required": true}, "expected_choice": "No Room"}

examples/calc_x/calc_agent.py

@@ -12,7 +12,8 @@ from autogen_core.models import ModelFamily
 from autogen_ext.models.openai import OpenAIChatCompletionClient
 from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams
 
-from agentlightning import LLM, LitAgent, NamedResources, Trainer, configure_logger, reward
+from agentlightning import LLM, LitAgent, NamedResources, Trainer, configure_logger
+from agentlightning.reward import reward
 
 configure_logger()
 
@@ -140,4 +141,4 @@ class CalcAgent(LitAgent[Any]):
 
 
 if __name__ == "__main__":
-    Trainer(n_workers=10).fit(CalcAgent(), "http://localhost:9999/")
+    Trainer(n_workers=10).fit_v0(CalcAgent(), "http://localhost:9999/")

examples/calc_x/calc_agent_dev.py

@@ -32,4 +32,6 @@ def dev_task_loader() -> DevTaskLoader:
 
 
 if __name__ == "__main__":
-    Trainer(n_workers=1, dev=True, max_tasks=2).fit(CalcAgent(), "http://localhost:9999/", dev_data=dev_task_loader())
+    Trainer(n_workers=1, dev=True, max_tasks=2).fit_v0(
+        CalcAgent(), "http://localhost:9999/", dev_data=dev_task_loader()
+    )

examples/calc_x/calc_agent_v0_2.py

@@ -123,7 +123,7 @@ def main():
     print(val_dataset[:5])  # type: ignore
 
     trainer = Trainer(algorithm=verl_algorithm(), n_workers=4)
-    trainer.fit_v2(calc_agent, train_dataset, val_dataset=val_dataset)  # type: ignore
+    trainer.fit(calc_agent, train_dataset, val_dataset=val_dataset)  # type: ignore
 
 
 if __name__ == "__main__":

examples/calc_x/calc_agent_v0_2_llm_proxy.py

@@ -8,7 +8,7 @@ from calc_agent import eval, get_agent
 from datasets import Dataset
 
 from agentlightning import LLM, Trainer, rollout
-from agentlightning.adapter import LlmProxyTripletAdapter
+from agentlightning.adapter import LlmProxyTraceToTriplet
 from agentlightning.algorithm.verl import VERL
 from agentlightning.tracer import OtelTracer
 
@@ -114,9 +114,9 @@ def main():
     print(val_dataset[:5])  # type: ignore
 
     tracer = OtelTracer()
-    adapter = LlmProxyTripletAdapter()
+    adapter = LlmProxyTraceToTriplet()
     trainer = Trainer(algorithm=VERL(rl_training_config), n_workers=10, tracer=tracer, adapter=adapter)
-    trainer.fit_v2(calc_agent, train_dataset, val_dataset=val_dataset)  # type: ignore
+    trainer.fit(calc_agent, train_dataset, val_dataset=val_dataset)  # type: ignore
 
 
 if __name__ == "__main__":

examples/rag/rag_agent.py

@@ -77,4 +77,4 @@ class RAGAgent(LitAgent[Any]):
 
 
 if __name__ == "__main__":
-    Trainer(n_workers=12).fit(RAGAgent(), "http://localhost:9999/")
+    Trainer(n_workers=12).fit_v0(RAGAgent(), "http://localhost:9999/")

examples/spider/spider_eval/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

examples/spider/spider_eval/async_utils.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import queue
+import threading
+
+
+def run_sync_ephemeral(coro):
+    """
+    Run an async coroutine from sync code.
+    - If no loop in this thread: use asyncio.run() directly.
+    - If already in an event loop: spawn a worker thread that calls asyncio.run()
+      (which creates and closes a brand-new event loop per call).
+    """
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        # No running loop in this thread; safe to use asyncio.run
+        return asyncio.run(coro)
+
+    # Already in a running loop -> execute in a worker thread
+    q = queue.Queue()
+
+    def worker():
+        try:
+            result = asyncio.run(coro)  # creates & closes its own loop
+            q.put((True, result))
+        except BaseException as e:
+            q.put((False, e))
+
+    t = threading.Thread(target=worker, daemon=True)
+    t.start()
+    ok, payload = q.get()
+    t.join()
+    if ok:
+        return payload
+    raise payload

examples/spider/spider_eval/exec_eval.py

@@ -18,6 +18,7 @@ from typing import Any, List, Set, Tuple
 
 import tqdm
 
+from .async_utils import run_sync_ephemeral
 from .parse import get_all_preds_for_execution, remove_distinct
 
 threadLock = threading.Lock()
@@ -225,8 +226,8 @@ def eval_exec_match(
             ranger = db_paths
 
         for db_path in ranger:
-            g_flag, g_denotation = asyncio.run(exec_on_db(db_path, g_str))
-            p_flag, p_denotation = asyncio.run(exec_on_db(db_path, pred))
+            g_flag, g_denotation = run_sync_ephemeral(exec_on_db(db_path, g_str))
+            p_flag, p_denotation = run_sync_ephemeral(exec_on_db(db_path, pred))
 
             # we should expect the gold to be succesfully executed on the database
             assert g_flag != "exception", "gold query %s has error on database file %s" % (g_str, db_path)

examples/spider/sql_agent.py

@@ -1,6 +1,8 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-"""
+"""Sample code that demonstrates an SQL agent using LangGraph and LangChain,
+trainable with Agent-lightning.
+
 Adapted from https://python.langchain.com/docs/tutorials/sql_qa/
 as well as https://langchain-ai.github.io/langgraph/tutorials/sql-agent/
 """
@@ -12,9 +14,9 @@ import re
 import shutil
 import tempfile
 import time
-from typing import Any, Dict, Literal, Optional, cast
+from typing import Any, Dict, List, Literal, Optional, cast
 
-import dotenv
+import pandas as pd
 import termcolor
 from langchain.chat_models import init_chat_model
 from langchain_community.tools.sql_database.tool import QuerySQLDatabaseTool
@@ -25,11 +27,11 @@ from langgraph.graph import END, START, MessagesState, StateGraph
 from langgraph.graph.state import CompiledStateGraph
 from spider_eval.exec_eval import eval_exec_match
 
-import agentlightning
+import agentlightning as agl
 
-agentlightning.configure_logger()
+agl.configure_logger()
 
-logger = agentlightning.configure_logger(name=__name__)
+logger = agl.configure_logger(name=__name__)
 
 
 WRITE_QUERY_PROMPT = ChatPromptTemplate(
@@ -411,7 +413,7 @@ def evaluate_query(query: str, ground_truth: str, database: str, raise_on_error:
             return 0.0
 
 
-class LitSQLAgent(agentlightning.LitAgent[Any]):
+class LitSQLAgent(agl.LitAgent[Dict[str, Any]]):
 
     def __init__(
         self,
@@ -428,20 +430,21 @@ class LitSQLAgent(agentlightning.LitAgent[Any]):
         self.table_info_truncate = table_info_truncate
         self.execution_truncate = execution_truncate
 
-    def _execute_rollout(
-        self, sample: dict[str, Any], *, resources: agentlightning.NamedResources, rollout_id: str, is_training: bool
+    def rollout(
+        self,
+        task: Dict[str, Any],
+        resources: agl.NamedResources,
+        rollout: agl.Rollout,
     ) -> float | None:
-        question = sample["question"]
+        question = task["question"]
         start_time = time.time()
-        llm: agentlightning.LLM = cast(agentlightning.LLM, resources["main_llm"])
+        llm: agl.LLM = cast(agl.LLM, resources["main_llm"])
 
-        if is_training:
-            original_db_path = os.path.join(self.spider_dir, "database", sample["db_id"], sample["db_id"] + ".sqlite")
+        if rollout.mode == "train":
+            original_db_path = os.path.join(self.spider_dir, "database", task["db_id"], task["db_id"] + ".sqlite")
         else:
-            original_db_path = os.path.join(
-                self.spider_dir, "test_database", sample["db_id"], sample["db_id"] + ".sqlite"
-            )
-        ground_truth = sample["query"]
+            original_db_path = os.path.join(self.spider_dir, "test_database", task["db_id"], task["db_id"] + ".sqlite")
+        ground_truth = task["query"]
 
         if not os.path.exists(original_db_path):
             logger.error(f"Database {original_db_path} does not exist. Skipping.")
@@ -455,6 +458,8 @@ class LitSQLAgent(agentlightning.LitAgent[Any]):
             logger.error("Schema file not found: %s", schema_path)
             schema = "No schema available."
 
+        rollout_id = rollout.rollout_id
+
         with tempfile.TemporaryDirectory() as temp_dir:
             db_path = os.path.join(temp_dir, os.path.basename(original_db_path))
             shutil.copyfile(original_db_path, db_path)
@@ -472,7 +477,7 @@ class LitSQLAgent(agentlightning.LitAgent[Any]):
                 endpoint=llm.endpoint,
                 verl_replacement=(
                     {"model": llm.model, **llm.sampling_parameters}
-                    if is_training
+                    if rollout.mode == "train"
                     else {
                         "model": llm.model,
                         "temperature": (
@@ -484,9 +489,11 @@ class LitSQLAgent(agentlightning.LitAgent[Any]):
                 ),
             ).graph()
             try:
+                # Required to make the langchain tracing work
+                handler = self.tracer.get_langchain_handler()
                 result = agent.invoke(  # type: ignore
                     {"question": question},  # type: ignore
-                    {"callbacks": [self.tracer.get_langchain_callback_handler()], "recursion_limit": 100},  # type: ignore
+                    {"callbacks": [handler] if handler else [], "recursion_limit": 100},
                 )
             except Exception as e:
                 logger.exception(f"[Rollout {rollout_id}] Error during agent invocation: {e}")
@@ -512,42 +519,27 @@ class LitSQLAgent(agentlightning.LitAgent[Any]):
 
         return reward
 
-    def training_rollout(self, task: Any, rollout_id: str, resources: agentlightning.NamedResources) -> Any:  # type: ignore
-        return self._execute_rollout(task, resources=resources, rollout_id=rollout_id, is_training=True)
-
-    def validation_rollout(self, task: Any, rollout_id: str, resources: agentlightning.NamedResources) -> Any:  # type: ignore
-        return self._execute_rollout(task, resources=resources, rollout_id=rollout_id, is_training=False)
-
-
-def spider_dev_data():
-    # Read from dev.parquet
-    import pandas as pd
 
+def debug_sql_agent():
     spider_dev_data_path = os.path.join(os.environ.get("VERL_SPIDER_DATA_DIR", "data"), "dev.parquet")
     if not os.path.exists(spider_dev_data_path):
         raise FileNotFoundError(f"Spider dev data file {spider_dev_data_path} does not exist.")
-    df = pd.read_parquet(spider_dev_data_path)  # type: ignore
-    if "OPENAI_API_BASE" not in os.environ:
-        logger.warning(
-            "Environment variable OPENAI_API_BASE is not set. Using default value 'https://api.openai.com/v1'."
-        )
-        openai_api_base = "https://api.openai.com/v1"
-    else:
-        openai_api_base = os.environ["OPENAI_API_BASE"]
+    df = pd.read_parquet(spider_dev_data_path).head(10)  # type: ignore
+    df = cast(List[Dict[str, Any]], df.to_dict(orient="records"))  # type: ignore
+    print("Debug data:", df)
 
-    resource = {
-        "main_llm": agentlightning.LLM(
-            model="gpt-4.1-nano",
-            endpoint=openai_api_base,
-            sampling_parameters={
-                "temperature": 0.0,
-            },
-        )
-    }
-    return agentlightning.DevTaskLoader(df.head(10).to_dict(orient="records"), resource)  # type: ignore
+    trainer = agl.Trainer(
+        n_workers=1,
+        initial_resources={
+            "main_llm": agl.LLM(
+                endpoint=os.environ["OPENAI_API_BASE"],
+                model="gpt-4.1-nano",
+                sampling_parameters={"temperature": 0.7},
+            )
+        },
+    )
+    trainer.dev(LitSQLAgent(), df)
 
 
 if __name__ == "__main__":
-    dotenv.load_dotenv()
-    agent, trainer = agentlightning.lightning_cli(LitSQLAgent, agentlightning.Trainer)
-    trainer.fit(agent, os.environ["VERL_API_BASE"], dev_data=spider_dev_data())
+    debug_sql_agent()

examples/spider/train_sql_agent.py

@@ -0,0 +1,189 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Train an SQL agent on the Spider dataset using Agent-lightning.
+
+This module provides a training script for SQL agents using different model configurations.
+The script supports three different training configurations:
+
+1. 'fast' - A lightweight configuration optimized for CI testing with reduced epochs
+2. 'qwen' - Standard configuration using Qwen-2.5-Coder-1.5B-Instruct model
+3. 'llama' - Configuration using LLaMA-3.2-3B-Instruct model with JSON formatting
+
+Usage:
+    python train_sql_agent.py fast    # Fast training for CI/testing
+    python train_sql_agent.py qwen    # Standard Qwen model training
+    python train_sql_agent.py llama   # LLaMA model training
+
+The script uses reinforcement learning with VERL (Versatile Efficient RL) algorithm
+to train agents on the Spider dataset for text-to-SQL generation tasks.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from copy import deepcopy
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import pandas as pd
+from sql_agent import LitSQLAgent
+
+import agentlightning as agl
+
+RL_TRAINING_CONFIG: Dict[str, Any] = {
+    "algorithm": {
+        "adv_estimator": "grpo",
+        "use_kl_in_reward": False,
+    },
+    "data": {
+        "train_files": "data/train_spider.parquet",
+        "val_files": "data/test_dev_500.parquet",
+        "train_batch_size": 32,
+        "max_prompt_length": 4096,
+        "max_response_length": 2048,
+        "truncation": "error",
+    },
+    "actor_rollout_ref": {
+        "rollout": {
+            "tensor_model_parallel_size": 1,
+            "n": 4,
+            "log_prob_micro_batch_size_per_gpu": 4,
+            "multi_turn": {"format": "hermes"},
+            "name": "vllm",
+            "gpu_memory_utilization": 0.8,
+        },
+        "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": 8,
+            "fsdp_config": {"param_offload": True},
+        },
+        "model": {
+            "path": "Qwen/Qwen2.5-Coder-1.5B-Instruct",
+            "use_remove_padding": True,
+            "enable_gradient_checkpointing": True,
+        },
+    },
+    "trainer": {
+        "n_gpus_per_node": 1,
+        "val_before_train": True,
+        "critic_warmup": 0,
+        "logger": ["console", "wandb"],
+        "project_name": "AgentLightning",
+        "experiment_name": "spider",
+        "nnodes": 1,
+        "test_freq": 32,
+        "total_epochs": 2,
+    },
+}
+
+
+def config_train_fast() -> Dict[str, Any]:
+    """A fast training run for CI testing purposes."""
+
+    # `EXPERIMENT_NAME="spider_$(date +%Y%m%d%H%M%S)"`
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    EXPERIMENT_NAME = f"spider_{timestamp}"
+
+    # `PROJECT_NAME=AgentLightningCI`
+    PROJECT_NAME = "AgentLightningCI"
+
+    # Simulate writing to $GITHUB_OUTPUT if it’s set
+    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")
+
+    print("Set environment variables:")
+    print(f"PROJECT_NAME={PROJECT_NAME}")
+    print(f"EXPERIMENT_NAME={EXPERIMENT_NAME}")
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.6
+    config["model"]["path"] = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+    config["data"]["val_files"] = "data/test_dev.parquet"
+    config["trainer"]["total_epochs"] = 1
+    config["trainer"]["total_training_steps"] = 1
+    config["trainer"]["experiment_name"] = EXPERIMENT_NAME
+    config["trainer"]["project_name"] = PROJECT_NAME
+    config["trainer"]["test_freq"] = 1
+    return config
+
+
+def config_train_qwen() -> Dict[str, Any]:
+    """A configuration for training with Qwen-2.5B."""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    return config
+
+
+def config_train_llama() -> Dict[str, Any]:
+    """A configuration for training with LLaMA-3.2-3B-Instruct."""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["multi_turn"]["format"] = "llama3_json"
+    config["actor_rollout_ref"]["model"]["path"] = "meta-llama/Llama-3.2-3B-Instruct"
+    return config
+
+
+def train(config: Dict[str, Any], active_agent: Optional[str]) -> None:
+    """Train the SQL agent with the given configuration."""
+
+    agent = LitSQLAgent()
+    algorithm = agl.VERL(config)
+    trainer = agl.Trainer(n_workers=10, algorithm=algorithm, adapter={"agent_match": active_agent})
+    print("Adapter agent match acknowledged:", trainer.adapter.agent_match)  # type: ignore
+
+    train_data = pd.read_parquet(config["data"]["train_files"]).to_dict(orient="records")  # type: ignore
+    val_data = pd.read_parquet(config["data"]["val_files"]).to_dict(orient="records")  # type: ignore
+    trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)  # type: ignore
+
+
+def main() -> None:
+    """Main function to parse arguments and run training."""
+    parser = argparse.ArgumentParser(
+        description="Train an SQL agent on the Spider dataset using different model configurations"
+    )
+
+    parser.add_argument(
+        "config",
+        choices=["fast", "qwen", "llama"],
+        help="Training configuration: 'fast' (CI testing), 'qwen' (Qwen-2.5-Coder-1.5B), 'llama' (LLaMA-3.2-3B)",
+    )
+
+    parser.add_argument(
+        "--active-agent", type=str, help="Override the active agent name (default: auto-generated based on config)"
+    )
+
+    args = parser.parse_args()
+
+    # Get the appropriate configuration
+    config_functions = {"fast": config_train_fast, "qwen": config_train_qwen, "llama": config_train_llama}
+
+    config = config_functions[args.config]()
+
+    # Set active agent - use provided value or default based on config choice
+    active_agent = args.active_agent
+
+    print(f"Starting training with '{args.config}' configuration...")
+    print(f"Active agent: {active_agent}")
+
+    train(config, active_agent)
+
+
+if __name__ == "__main__":
+    main()

examples/unsloth/math_agent.py

@@ -19,7 +19,7 @@ python math_agent.py
 import json
 import os
 import re
-from typing import Any, Optional, TypedDict, cast
+from typing import Any, Optional, TypedDict
 
 import numpy as np
 from agents import Agent, ModelSettings, OpenAIChatCompletionsModel, Runner
@@ -56,7 +56,7 @@ def _download_dataset() -> None:  # pyright: ignore[reportUnusedFunction]
     Downloads the first 64 samples from the dataset and saves them to data_gsmhard.jsonl.
     This function is provided as a utility to help set up the dataset for the first time.
     """
-    ds = load_dataset("reasoning-machines/gsm-hard", split="train")
+    ds = load_dataset("reasoning-machines/gsm-hard", split="train")  # pyright: ignore[reportUnknownVariableType]
     df = ds.to_list()  # type: ignore
     with open("data_gsmhard.jsonl", "w") as f:
         for i, row in enumerate(df):  # type: ignore
@@ -79,7 +79,7 @@ def load_math_dataset(limit: Optional[int] = None) -> Dataset[GsmProblem]:
         problems = [GsmProblem(**json.loads(line)) for line in f]
     if limit is not None:
         problems = problems[:limit]
-    return cast(Dataset[GsmProblem], problems)
+    return problems
 
 
 @rollout

examples/unsloth/sft_algorithm.py

@@ -33,11 +33,10 @@ from rich.console import Console
 from unsloth_helper import unsloth_training
 
 from agentlightning import configure_logger
-from agentlightning.adapter.triplet import BaseTraceTripletAdapter, LlmProxyTripletAdapter
+from agentlightning.adapter import LlmProxyTraceToTriplet, TraceToTripletBase
 from agentlightning.llm_proxy import LLMProxy, ModelConfig
-from agentlightning.store.base import LightningStore
-from agentlightning.store.client_server import LightningStoreClient
-from agentlightning.types import Dataset, RolloutV2
+from agentlightning.store import LightningStore, LightningStoreClient
+from agentlightning.types import Dataset, Rollout
 
 console = Console()
 
@@ -140,7 +139,7 @@ async def sft_one_iter(
     model_path: str,
     train_dataset: Dataset[GsmProblem],
     llm_proxy: LLMProxy,
-    data_adapter: BaseTraceTripletAdapter,
+    data_adapter: TraceToTripletBase,
     triplet_fraction: float,
     vllm_port: int,
 ) -> str:
@@ -196,7 +195,7 @@ async def sft_one_iter(
         )
 
         # Create tasks for runners to run, associating them with the proxy address
-        rollouts: List[RolloutV2] = []
+        rollouts: List[Rollout] = []
         for data in train_dataset:
             rollouts.append(
                 await store.enqueue_rollout(
@@ -209,7 +208,7 @@ async def sft_one_iter(
         console.print(f"[bold red][Algo][/bold red] Enqueued {len(rollouts)} rollouts")
 
         # Wait for the tasks to complete
-        completed_rollouts: List[RolloutV2] = []
+        completed_rollouts: List[Rollout] = []
 
         while True:
             completed_rollouts = await store.wait_for_rollouts(
@@ -363,7 +362,7 @@ async def sft_algorithm(*, store: LightningStore) -> None:
 
     # This data adapter util is used to convert the trace data recorded by LLM proxy
     # into a format suitable for SFT
-    data_adapter = LlmProxyTripletAdapter()
+    data_adapter = LlmProxyTraceToTriplet()
 
     for iteration in range(MAX_ITERATIONS):
         model_path = await sft_one_iter(

examples/unsloth/sft_allinone.py

@@ -18,7 +18,7 @@ from rich.console import Console
 from sft_algorithm import sft_one_iter
 
 from agentlightning import Trainer, configure_logger
-from agentlightning.adapter.triplet import BaseTraceTripletAdapter
+from agentlightning.adapter import TraceToTripletBase
 from agentlightning.algorithm import BaseAlgorithm
 from agentlightning.llm_proxy import LLMProxy
 from agentlightning.types import Dataset
@@ -66,8 +66,8 @@ class UnslothSupervisedFinetuning(BaseAlgorithm):
         data_adapter = self.get_adapter()
 
         # SFT trainer relies on the adapter to convert the trace data to triplets
-        if not isinstance(data_adapter, BaseTraceTripletAdapter):
-            raise ValueError("Data adapter must be a TraceTripletAdapter.")
+        if not isinstance(data_adapter, TraceToTripletBase):
+            raise ValueError("Data adapter must be a TracerTraceToTriplet.")
         if train_dataset is None:
             raise ValueError("Train dataset must be provided.")
         if val_dataset is not None:
@@ -108,8 +108,8 @@ if __name__ == "__main__":
         llm_proxy=LLMProxy(port=12358),
         # Uncomment the following two lines if you want to rely on proxy-side trace data collection
         # Otherwise, the rollout runner will have an agentops tracer to collect the trace data,
-        # and the adapter will be a TraceTripletAdapter that parses the trace data generated by this tracer
-        # adapter=LlmProxyTripletAdapter(),
+        # and the adapter will be a TracerTraceToTriplet that parses the trace data generated by this tracer
+        # adapter=LlmProxyTraceToTriplet(),
         # tracer=OtelTracer(),
     )
-    trainer.fit_v2(math_agent, load_math_dataset())
+    trainer.fit(math_agent, load_math_dataset())

examples/unsloth/sft_rollout_runners.py

@@ -18,9 +18,8 @@ from math_agent import GsmProblem, math_agent
 from rich.console import Console
 
 from agentlightning import configure_logger
-from agentlightning.runner import AgentRunnerV2
-from agentlightning.store.base import LightningStore
-from agentlightning.store.client_server import LightningStoreClient
+from agentlightning.runner import LitAgentRunner
+from agentlightning.store import LightningStore, LightningStoreClient
 from agentlightning.tracer import OtelTracer
 
 console = Console()
@@ -37,7 +36,7 @@ def run_rollout(*, store: LightningStore, worker_id: int) -> None:
     # a simple OtelTracer to collect the rewards is enough.
     tracer = OtelTracer()
 
-    runner = AgentRunnerV2[GsmProblem](tracer=tracer)
+    runner = LitAgentRunner[GsmProblem](tracer=tracer)
 
     console.print(f"[bold green]Runners: [/bold green] Rollout runner {worker_id} started.")
 

mkdocs.yml

@@ -61,6 +61,7 @@ plugins:
             show_symbol_type_heading: true
             show_symbol_type_toc: true
             docstring_style: google
+  - autorefs
   - mike:
       version_selector: true
       css_dir: css
@@ -83,12 +84,22 @@ nav:
     - Getting Started: quickstart/getting-started.md
   - How-To Guides:
     - Train SQL Agent: how-to/train-sql-agent.md
+  - Algorithm Zoo:
+    - Overview: algorithm-zoo/index.md
+    - APO: algorithm-zoo/apo.md
+    - VERL: algorithm-zoo/verl.md
   - Deep Dive:
     - Bird's Eye View: deep-dive/birds-eye-view.md
     - Server-Client Architecture (Legacy): deep-dive/server-client-architecture.md
-  - API Reference:
-    - Core: reference/core.md
-    - RL: reference/rl.md
+  - API References:
+    - Agent: reference/agent.md
+    - Algorithm: reference/algorithm.md
+    - Command Line: reference/cli.md
+    - Instrumentation: reference/instrumentation.md
+    - Runner: reference/runner.md
+    - Store: reference/store.md
+    - Trainer: reference/trainer.md
+    - Types: reference/types.md
 
 extra_css:
   - stylesheets/extra.css

pyproject.toml

@@ -44,9 +44,11 @@ dev = [
   "mkdocs-git-revision-date-localized-plugin",
   "mkdocs-git-authors-plugin",
   "mkdocs-macros-plugin",
+  "mkdocs-autorefs",
 ]
 experiment = [
   "random-word",
+  "gdown",
 ]
 agent = [
   "autogen-agentchat",
@@ -62,9 +64,16 @@ agent = [
   "uv",
   "anthropic",
 ]
+apo = [
+  "poml",
+]
+
+# This is used for type checking.
+# Generally, we do not recommend using install .[verl] to install those dependencies.
 trl = [
-  "unsloth",
-  "unsloth_zoo",
+  # https://github.com/unslothai/unsloth/issues/3451
+  "unsloth<=2025.10.1",
+  "unsloth_zoo<=2025.10.1",
   "bitsandbytes",
   "peft",
   "datasets",
@@ -72,6 +81,10 @@ trl = [
   "trl",
   "kernels",
 ]
+verl = [
+  "verl==0.5.0",
+  "vllm>=0.8.4,<0.11.0",
+]
 
 [build-system]
 requires = ["hatchling"]
@@ -79,7 +92,7 @@ build-backend = "hatchling.build"
 
 [tool.hatch.build.targets.wheel]
 packages = ["agentlightning"]
-include = ["**/*.yaml", "**/*.yml"]
+include = ["**/*.yaml", "**/*.yml", "**/*.poml"]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]