chore: import upstream snapshot with attribution

agentlightning/algorithm/__init__.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .base import Algorithm
+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__ = ["Algorithm", "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,895 @@
+# 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)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import random
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    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 Algorithm
+from agentlightning.algorithm.utils import batch_iter_over_dataset, with_llm_proxy, with_store
+from agentlightning.reward import find_final_reward
+from agentlightning.types import Dataset, NamedResources, PromptTemplate, Rollout, RolloutMode, RolloutStatus
+
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.store.base import LightningStore
+
+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",
+]
+
+
+class APO(Algorithm, 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,
+        gradient_prompt_files: Optional[List[Path]] = None,
+        apply_edit_prompt_files: Optional[List[Path]] = None,
+        # 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.
+            gradient_prompt_files: Prompt templates used to compute textual gradients (critiques).
+            apply_edit_prompt_files: Prompt templates used to apply edits based on critiques.
+        """
+        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.gradient_prompt_files = gradient_prompt_files or GRADIENT_PROMPT_FILES
+        self.apply_edit_prompt_files = apply_edit_prompt_files or APPLY_EDIT_PROMPT_FILES
+
+        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(self.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(self.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
+
+    @with_store
+    async def get_rollout_results(
+        self,
+        store: LightningStore,
+        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] = []
+        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,
+            )
+
+    @with_llm_proxy()
+    @with_store
+    async def run(
+        self,
+        store: LightningStore,  # Injected by decorator - callers should not provide this parameter
+        llm_proxy: Optional[LLMProxy],  # Injected by decorator - callers should not provide this parameter
+        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

@@ -0,0 +1,162 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import inspect
+import weakref
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Optional,
+    Union,
+)
+
+from agentlightning.adapter import TraceAdapter
+from agentlightning.client import AgentLightningClient
+from agentlightning.store.base import LightningStore
+from agentlightning.types import Dataset, NamedResources
+
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.trainer import Trainer
+
+
+class Algorithm:
+    """Algorithm is the strategy, or tuner to train the agent."""
+
+    _trainer_ref: weakref.ReferenceType[Trainer] | None = None
+    _llm_proxy_ref: weakref.ReferenceType["LLMProxy"] | None = None
+    _store: LightningStore | None = None
+    _initial_resources: NamedResources | None = None
+    _adapter_ref: weakref.ReferenceType[TraceAdapter[Any]] | None = None
+
+    def is_async(self) -> bool:
+        """Return True if the algorithm is asynchronous."""
+        return inspect.iscoroutinefunction(self.run)
+
+    def set_trainer(self, trainer: Trainer) -> None:
+        """
+        Set the trainer for this algorithm.
+
+        Args:
+            trainer: The Trainer instance that will handle training and validation.
+        """
+        self._trainer_ref = weakref.ref(trainer)
+
+    def get_trainer(self) -> Trainer:
+        """
+        Get the trainer for this algorithm.
+
+        Returns:
+            The Trainer instance associated with this agent.
+        """
+        if self._trainer_ref is None:
+            raise ValueError("Trainer has not been set for this agent.")
+        trainer = self._trainer_ref()
+        if trainer is None:
+            raise ValueError("Trainer reference is no longer valid (object has been garbage collected).")
+        return trainer
+
+    def set_llm_proxy(self, llm_proxy: LLMProxy | None) -> None:
+        """
+        Set the LLM proxy for this algorithm to reuse when available.
+
+        Args:
+            llm_proxy: The LLMProxy instance configured by the trainer, if any.
+        """
+        self._llm_proxy_ref = weakref.ref(llm_proxy) if llm_proxy is not None else None
+
+    def get_llm_proxy(self) -> Optional[LLMProxy]:
+        """
+        Retrieve the configured LLM proxy instance, if one has been set.
+
+        Returns:
+            The active LLMProxy instance or None when not configured.
+        """
+        if self._llm_proxy_ref is None:
+            return None
+
+        llm_proxy = self._llm_proxy_ref()
+        if llm_proxy is None:
+            raise ValueError("LLM proxy reference is no longer valid (object has been garbage collected).")
+
+        return llm_proxy
+
+    def set_adapter(self, adapter: TraceAdapter[Any]) -> None:
+        """
+        Set the adapter for this algorithm to collect and convert traces.
+        """
+        self._adapter_ref = weakref.ref(adapter)
+
+    def get_adapter(self) -> TraceAdapter[Any]:
+        """
+        Retrieve the adapter for this algorithm to communicate with the runners.
+        """
+        if self._adapter_ref is None:
+            raise ValueError("Adapter has not been set for this algorithm.")
+        adapter = self._adapter_ref()
+        if adapter is None:
+            raise ValueError("Adapter reference is no longer valid (object has been garbage collected).")
+        return adapter
+
+    def set_store(self, store: LightningStore) -> None:
+        """
+        Set the store for this algorithm to communicate with the runners.
+
+        Store is set directly instead of using weakref because its copy is meant to be
+        maintained throughout the algorithm's lifecycle.
+        """
+        self._store = store
+
+    def get_store(self) -> LightningStore:
+        """
+        Retrieve the store for this algorithm to communicate with the runners.
+        """
+        if self._store is None:
+            raise ValueError("Store has not been set for this algorithm.")
+        return self._store
+
+    def get_initial_resources(self) -> Optional[NamedResources]:
+        """
+        Get the initial resources for this algorithm.
+        """
+        return self._initial_resources
+
+    def set_initial_resources(self, resources: NamedResources) -> None:
+        """
+        Set the initial resources for this algorithm.
+        """
+        self._initial_resources = resources
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self.run(*args, **kwargs)
+
+    def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> Union[None, Awaitable[None]]:
+        """Subclasses should implement this method to implement the algorithm.
+
+        Args:
+            train_dataset: The dataset to train on. Not all algorithms require a training dataset.
+            val_dataset: The dataset to validate on. Not all algorithms require a validation dataset.
+
+        Returns:
+            Algorithm should refrain from returning anything. It should just run the algorithm.
+        """
+        raise NotImplementedError("Subclasses must implement run().")
+
+    def get_client(self) -> AgentLightningClient:
+        """Get the client to communicate with the algorithm.
+
+        If the algorithm does not require a server-client communication, it can also create a mock client
+        that never communicates with itself.
+
+        Deprecated and will be removed in a future version.
+
+        Returns:
+            The AgentLightningClient instance associated with this algorithm.
+        """
+        raise NotImplementedError("Subclasses must implement get_client().")

agentlightning/algorithm/decorator.py

@@ -0,0 +1,264 @@
+# 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 Algorithm
+
+# 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(Algorithm, Generic[AF]):
+    """An algorithm wrapper built from a callable implementation.
+
+    Functional algorithms let you provide an ordinary function instead of
+    subclassing [`Algorithm`][agentlightning.Algorithm]. The wrapper inspects
+    the callable signature to supply optional dependencies
+    such as the store, adapter, and LLM proxy.
+    """
+
+    @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:
+        """Wrap a function that implements algorithm behaviour.
+
+        Args:
+            algorithm_func: Sync or async callable implementing the algorithm
+                contract. Arguments are detected automatically based on the
+                function signature.
+        """
+        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 wrapped function with injected dependencies.
+
+        Args:
+            train_dataset: Optional training dataset passed through when the
+                callable declares a `train_dataset` parameter.
+            val_dataset: Optional validation dataset passed through when the
+                callable declares a `val_dataset` parameter.
+
+        Returns:
+            None for sync callables or an awaitable when the callable is async.
+
+        Raises:
+            TypeError: If a dataset is provided but the function signature does
+                not accept the corresponding argument.
+        """
+        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]]]:
+    """Convert a callable into a [`FunctionalAlgorithm`][agentlightning.algorithm.decorator.FunctionalAlgorithm].
+
+    The decorator inspects the callable signature to decide which dependencies
+    to inject at runtime, enabling concise algorithm definitions that still
+    leverage the full training runtime.
+
+    Args:
+        func: Function implementing the algorithm logic. May be synchronous or
+            asynchronous. The function can expect all of, or a subset of the following parameters:
+
+            - `store`: [`LightningStore`][agentlightning.store.base.LightningStore],
+            - `train_dataset`: [`Dataset`][agentlightning.Dataset],
+            - `val_dataset`: [`Dataset`][agentlightning.Dataset],
+            - `llm_proxy`: [`LLMProxy`][agentlightning.LLMProxy],
+            - `adapter`: [`TraceAdapter`][agentlightning.TraceAdapter],
+            - `initial_resources`: [`NamedResources`][agentlightning.NamedResources],
+
+            If the function does not expect a parameter, the wrapper will not inject it into the call.
+            Using `*args` and `**kwargs` will not work and no parameters will be injected.
+
+    Returns:
+        FunctionalAlgorithm that proxies the callable while exposing the
+        `Algorithm` interface.
+
+    Examples:
+        ```python
+        from agentlightning.algorithm.decorator import algo
+
+        @algo
+        def batching_algorithm(*, store, train_dataset, val_dataset):
+            for sample in train_dataset:
+                store.enqueue_rollout(input=sample, mode="train")
+
+        @algo
+        async def async_algorithm(*, store, train_dataset=None, val_dataset=None):
+            await store.enqueue_rollout(input={"prompt": "hello"}, mode="train")
+        ```
+    """
+    return FunctionalAlgorithm(func)

agentlightning/algorithm/fast.py

@@ -0,0 +1,250 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, List, Literal, Optional
+
+from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span
+
+from .base import Algorithm
+from .utils import with_llm_proxy, with_store
+
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.store.base import LightningStore
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["FastAlgorithm", "Baseline"]
+
+
+class FastAlgorithm(Algorithm):
+    """Base class for lightweight algorithms optimised for developer workflows.
+
+    Fast algorithms prioritise short feedback loops so an agent developer can run
+    small-scale experiments without waiting for long-running training jobs to
+    finish.
+    """
+
+
+def _timestamp_to_iso_str(timestamp: float) -> str:
+    return datetime.fromtimestamp(timestamp).isoformat()
+
+
+class Baseline(FastAlgorithm):
+    """Reference implementation that streams the full dataset through the rollout queue.
+
+    The baseline algorithm batches task submissions, waits for each rollout to
+    finish, and logs every collected span and reward. It is primarily useful as
+    a smoke test for the platform plumbing rather than a performant trainer.
+
+    The baseline algorithm will auto-start a LLM proxy if one is provided and not yet started.
+
+    Args:
+        n_epochs: Number of dataset passes to execute for both the train and val
+            splits during developer experiments.
+        train_split: Fraction of the concatenated dataset to treat as training
+            data. Must be strictly between 0 and 1.
+        polling_interval: Interval, in seconds, to poll the store for queue
+            depth and rollout completion.
+        max_queue_length: Number of rollouts allowed to wait in the queue before
+            throttling additional submissions.
+        span_verbosity: Level of detail to include when logging span metadata.
+
+    Raises:
+        ValueError: If `train_split` falls outside the `(0, 1)` interval.
+
+    Examples:
+        ```python
+        from agentlightning.algorithm.fast import Baseline
+
+        algorithm = Baseline(n_epochs=2, train_split=0.8, span_verbosity="key_values")
+        trainer.fit(algorithm, train_dataset=my_train, val_dataset=my_val)
+        ```
+    """
+
+    def __init__(
+        self,
+        *,
+        n_epochs: int = 1,
+        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
+
+    def _span_to_string(self, rollout_id: str, attempt: Attempt, span: Span) -> str:
+        """Format a span for logging based on the configured verbosity."""
+        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:
+        """Log attempt metadata and emit adapted traces when a rollout ends."""
+        store = self.get_store()
+
+        rollout_id = rollout.rollout_id
+        rollout_end_time = rollout.end_time or asyncio.get_event_loop().time()
+        logger.info(
+            f"[Rollout {rollout_id}] Finished with status {rollout.status} in {rollout_end_time - rollout.start_time:.2f} seconds."
+        )
+
+        # Logs all the attempts and their corresponding spans
+        attempts = await store.query_attempts(rollout_id)
+        for attempt in attempts:
+            logger.info(
+                "[Rollout %s | Attempt %s] ID: %s. Status: %s. Worker: %s",
+                rollout_id,
+                attempt.sequence_id,
+                attempt.attempt_id,
+                attempt.status,
+                attempt.worker_id,
+            )
+            spans = await store.query_spans(rollout_id=rollout_id)
+            for span in spans:
+                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:
+            adapter = self.get_adapter()
+        except ValueError:
+            logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.")
+            adapter = None
+        if adapter is not None:
+            spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
+            transformed_data = adapter.adapt(spans)
+            logger.info(f"[Rollout {rollout_id}] Adapted data: {transformed_data}")
+
+    async def _enqueue_rollouts(
+        self, dataset: Dataset[Any], train_indices: List[int], val_indices: List[int], resources_id: str
+    ) -> None:
+        """Submit rollouts while respecting the maximum queue length."""
+        store = self.get_store()
+
+        for index in train_indices + val_indices:
+            queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"])
+            if len(queuing_rollouts) <= 1:
+                # Only enqueue a new rollout when there is at most 1 rollout in the queue.
+                sample = dataset[index]
+                mode = "train" if index in train_indices else "val"
+                rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
+                logger.info(f"[Rollout {rollout.rollout_id}] Enqueued in {mode} mode with sample: {sample}")
+            await asyncio.sleep(self.polling_interval)
+
+    async def _harvest_rollout_spans(self, rollout_id: str):
+        """Poll rollout status updates until completion and log transitions."""
+        store = self.get_store()
+        last_status: Optional[RolloutStatus] = None
+
+        while True:
+            rollout = await store.get_rollout_by_id(rollout_id)
+            if rollout is not None:
+                if rollout.status in ["succeeded", "failed", "cancelled"]:
+                    # Rollout is finished, log all the data.
+                    await self._handle_rollout_finish(rollout)
+                    # We are done here.
+                    self._finished_rollout_count += 1
+                    logger.info(f"Finished {self._finished_rollout_count} rollouts.")
+                    break
+
+                if last_status != rollout.status:
+                    if last_status is not None:
+                        logger.info(f"[Rollout {rollout_id}] Status changed to {rollout.status}.")
+                    else:
+                        logger.info(f"[Rollout {rollout_id}] Status is initialized to {rollout.status}.")
+                    last_status = rollout.status
+
+                else:
+                    logger.debug(f"[Rollout {rollout_id}] Status is still {rollout.status}.")
+
+            await asyncio.sleep(self.polling_interval)
+
+    @with_llm_proxy()
+    @with_store
+    async def run(
+        self,
+        store: LightningStore,  # Injected by decorator - callers should not provide this parameter
+        llm_proxy: Optional[LLMProxy],  # Injected by decorator - callers should not provide this parameter
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> None:
+        """Execute the baseline loop across the provided datasets."""
+        train_dataset_length = len(train_dataset) if train_dataset is not None else 0
+        val_dataset_length = len(val_dataset) if val_dataset is not None else 0
+        if train_dataset_length == 0 and val_dataset_length == 0:
+            logger.error(
+                "MockAlgorithm requires at least one dataset. Provide train_dataset or val_dataset before running."
+            )
+            return
+
+        concatenated_dataset = [train_dataset[i] for i in range(train_dataset_length) if train_dataset is not None] + [
+            val_dataset[i] for i in range(val_dataset_length) if val_dataset is not None
+        ]
+        train_indices = list(range(0, train_dataset_length))
+        val_indices = list(range(train_dataset_length, train_dataset_length + val_dataset_length))
+        logger.debug(f"Train indices: {train_indices}")
+        logger.debug(f"Val indices: {val_indices}")
+
+        # Currently we only supports a single resource update at the start.
+        initial_resources = self.get_initial_resources()
+        if initial_resources is not None:
+            resource_update = await store.update_resources("default", initial_resources)
+            resources_id = resource_update.resources_id
+            logger.info(f"Initial resources set: {initial_resources}")
+        else:
+            logger.warning("No initial resources provided. Skip initializing resources.")
+            resources_id = None
+
+        for epoch in range(self.n_epochs):
+            harvest_tasks: List[asyncio.Task[None]] = []
+            logger.info(f"Proceeding epoch {epoch + 1}/{self.n_epochs}.")
+            for index in train_indices + val_indices:
+                logger.info(
+                    f"Processing index {index}. {len(train_indices)} train indices and {len(val_indices)} val indices in total."
+                )
+                while True:
+                    queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"])
+                    if len(queuing_rollouts) <= self.max_queue_length:
+                        # Only enqueue a new rollout when there is at most "max_queue_length" rollout in the queue.
+                        sample = concatenated_dataset[index]
+                        mode = "train" if index in train_indices else "val"
+                        rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
+                        harvest_tasks.append(asyncio.create_task(self._harvest_rollout_spans(rollout.rollout_id)))
+                        logger.info(f"Enqueued rollout {rollout.rollout_id} in {mode} mode with sample: {sample}")
+                        break
+                    else:
+                        # Sleep a bit and try again later.
+                        await asyncio.sleep(self.polling_interval)
+
+            # Wait for all harvest tasks to complete
+            logger.info(f"Waiting for {len(harvest_tasks)} harvest tasks to complete...")
+            if len(harvest_tasks) > 0:
+                await asyncio.gather(*harvest_tasks)

agentlightning/algorithm/utils.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import functools
+import logging
+import random
+from collections.abc import Coroutine
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Concatenate,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    ParamSpec,
+    Sequence,
+    TypeVar,
+    overload,
+)
+
+from agentlightning.types import Dataset
+
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.store.base import LightningStore
+
+    from .base import Algorithm
+
+T_task = TypeVar("T_task")
+T_algo = TypeVar("T_algo", bound="Algorithm")
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+logger = logging.getLogger(__name__)
+
+
+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 = []
+
+
+def with_store(
+    func: Callable[Concatenate[T_algo, LightningStore, P], Coroutine[Any, Any, R]],
+) -> Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]]:
+    """Inject the algorithm's `LightningStore` into coroutine methods.
+
+    The decorator calls `Algorithm.get_store()` once per invocation and passes the
+    resulting store as an explicit argument to the wrapped coroutine. Decorated
+    methods therefore receive the resolved store even when invoked by helper
+    utilities rather than directly by the algorithm.
+
+    Args:
+        func: The coroutine that expects `(self, store, *args, **kwargs)`.
+
+    Returns:
+        A coroutine wrapper that automatically retrieves the store and forwards it
+        to `func`.
+    """
+
+    @functools.wraps(func)
+    async def wrapper(self: T_algo, *args: P.args, **kwargs: P.kwargs) -> R:
+        store = self.get_store()
+        return await func(self, store, *args, **kwargs)
+
+    return wrapper
+
+
+@overload
+def with_llm_proxy(
+    required: Literal[False] = False,
+    auto_start: bool = True,
+) -> Callable[
+    [Callable[Concatenate[T_algo, Optional[LLMProxy], P], Coroutine[Any, Any, R]]],
+    Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]],
+]: ...
+
+
+@overload
+def with_llm_proxy(
+    required: Literal[True],
+    auto_start: bool = True,
+) -> Callable[
+    [Callable[Concatenate[T_algo, LLMProxy, P], Coroutine[Any, Any, R]]],
+    Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]],
+]: ...
+
+
+def with_llm_proxy(
+    required: bool = False,
+    auto_start: bool = True,
+) -> Callable[
+    [Callable[..., Coroutine[Any, Any, Any]]],
+    Callable[..., Coroutine[Any, Any, Any]],
+]:
+    """Resolve and optionally lifecycle-manage the configured LLM proxy.
+
+    Args:
+        required: When True, raises `ValueError` if the algorithm does not have an
+            [`LLMProxy`][agentlightning.LLMProxy] set. When False, the wrapped coroutine receives
+            `None` if no proxy is available.
+        auto_start: When True, [`LLMProxy.start()`][agentlightning.LLMProxy.start] is invoked if the proxy is not
+            already running before calling `func` and [`LLMProxy.stop()`][agentlightning.LLMProxy.stop] is
+            called afterwards.
+
+    Returns:
+        A decorator that injects the [`LLMProxy`][agentlightning.LLMProxy] (or `None`) as the first
+        argument after `self` and manages automatic startup/shutdown when requested.
+    """
+
+    def decorator(
+        func: Callable[..., Coroutine[Any, Any, Any]],
+    ) -> Callable[..., Coroutine[Any, Any, Any]]:
+        @functools.wraps(func)
+        async def wrapper(self: Algorithm, *args: Any, **kwargs: Any) -> Any:
+            llm_proxy = self.get_llm_proxy()
+
+            if required and llm_proxy is None:
+                raise ValueError(
+                    "LLM proxy is required but not configured. Call set_llm_proxy() before using this method."
+                )
+
+            auto_started = False
+            if auto_start and llm_proxy is not None:
+                if llm_proxy.is_running():
+                    logger.info("Proxy is already running, skipping start")
+                else:
+                    logger.info("Starting proxy, managed by the algorithm")
+                    await llm_proxy.start()
+                    auto_started = True
+
+            try:
+                # At type level, overloads guarantee that if `required=True`
+                # then `func` expects a non-optional LLMProxy.
+                return await func(self, llm_proxy, *args, **kwargs)
+            finally:
+                if auto_started and llm_proxy is not None:
+                    logger.info("Stopping proxy, managed by the algorithm")
+                    await llm_proxy.stop()
+
+        return wrapper
+
+    return decorator

agentlightning/algorithm/verl/__init__.py

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

agentlightning/algorithm/verl/interface.py

@@ -0,0 +1,202 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Optional, Type
+
+from hydra import compose, initialize
+from omegaconf import OmegaConf
+
+from agentlightning.algorithm.base import Algorithm
+from agentlightning.client import AgentLightningClient
+from agentlightning.types import Dataset
+from agentlightning.verl.entrypoint import run_ppo  # type: ignore
+
+if TYPE_CHECKING:
+    from agentlightning.verl.daemon import AgentModeDaemon
+    from agentlightning.verl.trainer import AgentLightningTrainer
+
+
+class VERL(Algorithm):
+    """VERL-powered algorithm that delegates training to the VERL PPO runner.
+
+    !!! warning
+        Advanced customisation currently requires copying the VERL source and
+        modifying it directly. Native hooks for overriding training behaviour
+        will land in a future release.
+
+    Args:
+        config: Dictionary mirroring the overrides passed to the VERL CLI. The
+            overrides are merged with VERL's packaged defaults via Hydra before
+            launching training.
+        trainer_cls: Optional override for the trainer class. Experimental.
+        daemon_cls: Optional override for the daemon class. Experimental.
+
+    !!! note "Trajectory aggregation (experimental)"
+
+        Trajectory-level aggregation merges an entire multi-turn rollout into a single,
+        masked training sample so GPU time is spent once per trajectory rather than N times
+        per turn. Enable it via:
+
+        ```python
+        config["agentlightning"]["trace_aggregator"] = {
+            "level": "trajectory",
+            "trajectory_max_prompt_length": 4096,
+            "trajectory_max_response_length": 34384,
+        }
+        ```
+
+        Keep conversations structured (message lists rather than manual string
+        concatenation) so prefix matching can stitch traces. `trajectory_max_prompt_length`
+        should be set to the maximum length of the prompt for the first turn, and
+        `trajectory_max_response_length` should be set to the maximum cumulative
+        length of agent responses in the full trajectory.
+        Toggle `debug=True` plus `mismatch_log_dir` when you need to inspect
+        retokenization or chat-template mismatches. See
+        [this blog post](https://agent-lightning.github.io/posts/trajectory_level_aggregation/)
+        for more details.
+
+    Examples:
+        ```python
+        from agentlightning.algorithm.verl import VERL
+
+        algorithm = VERL(
+            config={
+                "algorithm": {
+                    "adv_estimator": "grpo",
+                    "use_kl_in_reward": False,
+                },
+                "data": {
+                    "train_batch_size": 32,
+                    "max_prompt_length": 4096,
+                    "max_response_length": 2048,
+                },
+                "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.6,
+                    },
+                    "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-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": "calc_x",
+                    "nnodes": 1,
+                    "save_freq": 64,
+                    "test_freq": 32,
+                    "total_epochs": 2,
+                },
+            }
+        )
+        trainer.fit(algorithm, train_dataset=my_train_dataset)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: dict[str, Any],
+        trainer_cls: Optional[Type[AgentLightningTrainer]] = None,
+        daemon_cls: Optional[Type[AgentModeDaemon]] = None,
+    ):
+        super().__init__()
+
+        # Compose the base config exactly like your decorator:
+        with initialize(version_base=None, config_path="pkg://agentlightning/verl"):
+            base_cfg = compose(config_name="config")
+
+        # Merge your dict overrides
+        override_conf = OmegaConf.create(config)
+        # Allow adding new fields
+        OmegaConf.set_struct(base_cfg, False)
+        self.config = OmegaConf.merge(base_cfg, override_conf)
+        self.trainer_cls = trainer_cls
+        self.daemon_cls = daemon_cls
+
+    def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> None:
+        """Launch the VERL PPO entrypoint with the configured runtime context.
+
+        Args:
+            train_dataset: Optional dataset forwarded to VERL for training.
+            val_dataset: Optional dataset forwarded to VERL for evaluation.
+
+        Raises:
+            ValueError: If required dependencies such as the store, LLM proxy, or
+                adapter have been garbage-collected when using the V1 execution
+                mode.
+        """
+        from agentlightning.verl.daemon import AgentModeDaemon
+        from agentlightning.verl.trainer import AgentLightningTrainer
+
+        trainer_cls = self.trainer_cls or AgentLightningTrainer
+        daemon_cls = self.daemon_cls or AgentModeDaemon
+        try:
+            store = self.get_store()
+        except Exception:
+            print("Store is not set. Assuming v0 execution mode.")
+            run_ppo(
+                self.config,
+                train_dataset=train_dataset,
+                val_dataset=val_dataset,
+                store=None,
+                llm_proxy=None,
+                adapter=None,
+                trainer_cls=trainer_cls,
+                daemon_cls=daemon_cls,
+            )
+        else:
+            print("Store is set. Assuming v1 execution mode.")
+            llm_proxy = self.get_llm_proxy()
+            adapter = self.get_adapter()
+            run_ppo(
+                self.config,
+                train_dataset=train_dataset,
+                val_dataset=val_dataset,
+                store=store,
+                llm_proxy=llm_proxy,
+                adapter=adapter,
+                trainer_cls=trainer_cls,
+                daemon_cls=daemon_cls,
+            )
+
+    def get_client(self) -> AgentLightningClient:
+        """Create a client bound to the VERL-managed Agent Lightning server.
+
+        Deprecated:
+            Since v0.2.
+        """
+        port = self.config.agentlightning.port
+        return AgentLightningClient(endpoint=f"http://localhost:{port}")