Merge branch 'main' of github.com:microsoft/agent-lightning into apo-algorithm

agentlightning/algorithm/apo/__init__.py

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

agentlightning/algorithm/apo/apo.py

@@ -0,0 +1,603 @@
+# 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 logging
+import random
+import time
+from pathlib import Path
+from typing import Any, Dict, Generic, Iterator, List, Optional, Sequence, Tuple, TypedDict, TypeVar, cast
+
+import poml
+from openai import AsyncOpenAI
+
+from agentlightning.adapter.messages import TraceMessagesAdapter
+from agentlightning.algorithm.base import BaseAlgorithm
+from agentlightning.reward import find_final_reward
+from agentlightning.types import Dataset, NamedResources, PromptTemplate, RolloutMode, RolloutStatus, RolloutV2
+
+logger = logging.getLogger(__name__)
+
+T_task = TypeVar("T_task", bound=Dict[str, Any])
+
+
+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[Dict[str, Any]]
+
+
+GRADIENT_PROMPT_FILES = [
+    Path(__file__).parent / "prompts" / "text_gradient_variant01.poml",
+    Path(__file__).parent / "prompts" / "text_gradient_variant02.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 = 600.0,
+    ):
+        """
+        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.
+        """
+        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._history_best_prompt: Optional[PromptTemplate] = None
+        self._history_best_score: float = float("-inf")
+
+    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) -> TraceMessagesAdapter:
+        """
+        Get the adapter for converting spans to messages.
+
+        Returns:
+            The TraceMessagesAdapter instance for this algorithm.
+
+        Raises:
+            ValueError: If the adapter is not a TraceMessagesAdapter.
+        """
+        adapter = super().get_adapter()
+        if not isinstance(adapter, TraceMessagesAdapter):
+            raise ValueError("Adapter must be a TraceMessagesAdapter 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: str,
+        rollout_results: List[RolloutResultForAPO],
+    ) -> 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:
+            logger.warning(
+                f"Only {len(rollout_results)} rollouts available, but {self.gradient_batch_size} are needed. Using all rollouts."
+            )
+            sampled_rollout_results = rollout_results
+        else:
+            sampled_rollout_results = random.sample(rollout_results, self.gradient_batch_size)
+
+        logger.info(
+            f"Gradient will be computed with {self.gradient_model} for {len(sampled_rollout_results)} rollouts with template: {tg_template}"
+        )
+
+        tg_msg = poml.poml(  # type: ignore
+            tg_template,
+            context={
+                "experiments": sampled_rollout_results,
+                "prompt_template": current_prompt,
+            },
+            format="openai_chat",
+        )
+        logger.debug(f"Gradient computed with {self.gradient_model} prompt: {tg_msg}")
+        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
+        logger.debug(f"Gradient computed with {self.gradient_model} has result: {critique_text}")
+
+        return critique_text
+
+    async def textual_gradient_and_apply_edit(
+        self,
+        current_prompt: str,
+        rollout: List[RolloutResultForAPO],
+    ) -> 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)
+        if not critique_text:
+            logger.error(f"Failed to compute critique for prompt.")
+            return current_prompt
+
+        # 2) Apply edit
+        ae_template = random.choice(APPLY_EDIT_PROMPT_FILES)
+        logger.info(f"Edit will be generated by {self.apply_edit_model} with template: {ae_template}")
+        ae_msg = poml.poml(  # type: ignore
+            ae_template,
+            context={
+                "prompt_template": current_prompt,
+                "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:
+            logger.info(f"Edit generated by {self.apply_edit_model}: {new_prompt}")
+        return new_prompt
+
+    async def get_rollout_results(self, rollout: List[RolloutV2]) -> 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=[m.model_dump() for m in messages],
+            )
+            logger.info(
+                f"Rollout result for {r.rollout_id}: status {rollout_result['status']} "
+                f"with final reward {rollout_result['final_reward']}. "
+                f"{len(rollout_result['spans'])} spans and {len(rollout_result['messages'])} messages."
+            )
+            rollout_results.append(rollout_result)
+        return rollout_results
+
+    async def evaluate_prompt_on_batch(
+        self,
+        prompt: str,
+        resource_name: str,
+        dataset: Sequence[T_task],
+        mode: RolloutMode,
+    ) -> 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()
+        logger.info(f'Evaluating prompt "{prompt[:50]}..." on {len(dataset)} tasks in {mode} mode')
+
+        # Install prompt as named resource
+        resources: NamedResources = {resource_name: PromptTemplate(template=prompt, engine="f-string")}
+        await store.add_resources(resources)
+
+        rollout_ids: List[str] = []
+        for t in dataset:
+            r = await store.enqueue_rollout(input=t, mode=mode)  # task can be any dict processable by the client
+            rollout_ids.append(r.rollout_id)
+
+        deadline = time.time() + self.rollout_batch_timeout
+        finished: List[RolloutV2] = []
+        while time.time() < deadline:
+            finished = await store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=0.0)
+            if len(finished) >= len(rollout_ids):
+                logger.info(f"All {len(rollout_ids)} rollouts finished within timeout.")
+                break
+
+        rollout_results = await self.get_rollout_results(finished)
+        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)))
+
+        logger.info(f"Evaluated {len(rollout_results)} rollouts. Rewards: {final_rewards}. Average reward: {avg}")
+        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[PromptTemplate], round_num: int) -> List[PromptTemplate]:
+        """
+        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.
+        """
+        if len(beam) < self.beam_width:
+            logger.warning(
+                f"[Round {round_num + 1}] Beam width is currently {self.beam_width}, but only {len(beam)} prompts in beam. "
+                "Replicating all prompts."
+            )
+            return [beam[i % len(beam)] for i in range(self.beam_width)]
+        else:
+            return random.sample(beam, self.beam_width)
+
+    async def _generate_candidate_prompts(
+        self,
+        parent_prompts: List[PromptTemplate],
+        resource_name: str,
+        grad_dataset_iterator: Iterator[Sequence[T_task]],
+        round_num: int,
+    ) -> List[PromptTemplate]:
+        """
+        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.
+        """
+        logger.info(
+            f"[Round {round_num + 1}] Applying {self.branch_factor} edits to each of "
+            f"the {len(parent_prompts)} parents on training dataset"
+        )
+
+        candidates: List[PromptTemplate] = []
+        for prompt in parent_prompts:
+            for _ in range(self.branch_factor):
+                grad_samples = next(grad_dataset_iterator)
+                rollout_results, _ = await self.evaluate_prompt_on_batch(
+                    prompt.template, resource_name, grad_samples, mode="train"
+                )
+                new_prompt = await self.textual_gradient_and_apply_edit(prompt.template, rollout_results)
+                if not new_prompt:
+                    logger.error(f"[Round {round_num + 1}] Failed to compute edit for prompt: {prompt.template}")
+                    continue
+                new_prompt_template = PromptTemplate(template=new_prompt, engine="f-string")
+                logger.info(f"[Round {round_num + 1}] New prompt template: {new_prompt_template}")
+                candidates.append(new_prompt_template)
+
+        return candidates
+
+    async def _evaluate_and_select_beam(
+        self,
+        candidates: List[PromptTemplate],
+        resource_name: str,
+        val_dataset_iterator: Iterator[Sequence[T_task]],
+        round_num: int,
+    ) -> List[PromptTemplate]:
+        """
+        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.
+        """
+        logger.info(f"[Round {round_num + 1}] Evaluating {len(candidates)} candidates on validation dataset")
+
+        val_batch = next(val_dataset_iterator)
+        scores: List[Tuple[PromptTemplate, float]] = []
+
+        for idx, prompt in enumerate(candidates):
+            _, score = await self.evaluate_prompt_on_batch(prompt.template, resource_name, val_batch, mode="val")
+            scores.append((prompt, score))
+            logger.info(f"[Round {round_num + 1}] Candidate {idx} score: {score:.3f}")
+
+        # Sort by score (descending) and select top beam_width
+        sorted_prompts = [p for p, _ in sorted(scores, key=lambda x: x[1], reverse=True)][: self.beam_width]
+        logger.info(
+            f"[Round {round_num + 1}] Top {len(sorted_prompts)} candidates on validation dataset: {sorted_prompts}"
+        )
+
+        if len(sorted_prompts) == 0:
+            raise ValueError("No beam candidates any more")
+
+        return sorted_prompts
+
+    async def _update_best_prompt(
+        self,
+        beam: List[PromptTemplate],
+        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).
+        """
+        best_prompt = beam[0]
+        _, best_score = await self.evaluate_prompt_on_batch(
+            best_prompt.template, resource_name, cast(Sequence[T_task], val_dataset), mode="val"
+        )
+        logger.info(f"[Round {round_num + 1}] Best prompt {best_prompt} has score: {best_score:.3f}")
+
+        if best_score > self._history_best_score:
+            logger.info(
+                f"[Round {round_num + 1}] Best prompt updated. New best score: {best_score:.3f} (prev: {self._history_best_score:.3f})"
+            )
+            self._history_best_prompt = best_prompt
+            self._history_best_score = best_score
+
+    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)
+
+        # Validation datasets are guaranteed to be non-None after initialization
+        assert val_dataset is not None
+
+        # Start with seed prompt in the beam
+        beam: List[PromptTemplate] = [seed_prompt]
+
+        # Run beam search for specified number of rounds
+        for rnd in range(self.beam_rounds):
+            logger.info(f"[Round {rnd + 1}] Round {rnd + 1}/{self.beam_rounds}...")
+
+            # 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>You will revise the current prompt using the critique.</p>
+  <cp caption="Rules">
+    <list>
+      <item>Keep style concise.</item>
+      <item>Add explicit output format if helpful.</item>
+      <item>Prefer mechanism-first.</item>
+      <item>Include word limits if missing.</item>
+    </list>
+  </cp>
+  <output-format>
+    Return only the new prompt template. Keep the placeholders (in curly-brackets) as still placeholders. No fenced code blocks. No other text.
+  </output-format>
+  <human-msg>
+    <cp caption="Current Prompt Template">
+      <text>{{ prompt_template }}</text>
+    </cp>
+    <cp caption="Critique">
+      <text>{{ critique }}</text>
+    </cp>
+  </human-msg>
+</poml>

agentlightning/algorithm/apo/prompts/apply_edit_variant02.poml

@@ -0,0 +1,15 @@
+<poml>
+  <p>Revise the prompt to address all critique points. Preserve variable names in curly-brackets.</p>
+  <hint>Add a rubric: mechanism first, 1 support fact, less than 100 words.</hint>
+  <output-format>
+    Return only the new full prompt. No extra explanations, punctuations or diff.
+  </output-format>
+  <human-msg>
+    <cp caption="PROMPT" level="3">
+      <text>{{ prompt_template }}</text>
+    </cp>
+    <cp caption="CRITIQUE" level="3">
+      <text>{{ 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>{{ 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>{{ 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>{{ 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>

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
 

examples/apo/apo_existing_algo.py

@@ -0,0 +1,68 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""This sample code demonstrates how to use an existing APO algorithm to tune the prompts."""
+
+from typing import List, Optional
+
+from openai import AsyncOpenAI
+from rich.console import Console
+
+from agentlightning import Trainer, configure_logger
+from agentlightning.algorithm.base import algo
+from agentlightning.litagent.decorator import rollout
+from agentlightning.reward import find_final_reward
+from agentlightning.store.base import LightningStore
+from agentlightning.types import NamedResources, PromptTemplate, Span
+
+console = Console()
+
+
+@rollout
+async def apo_rollout(task: str, prompt_template: PromptTemplate) -> float:
+    # This relies on a public OpenAI service
+    client = AsyncOpenAI()
+
+    result = await client.chat.completions.create(
+        model="gpt-4.1-nano",
+        messages=[
+            {"role": "user", "content": prompt_template.format(any_question=task)},
+        ],
+    )
+
+    text = result.choices[0].message.content
+    console.print(f"[bold yellow][Rollout][/bold yellow] LLM returned: {text}")
+
+    return await llm_judge(task, text)
+
+
+async def llm_judge(task: str, output: Optional[str]) -> float:
+    client = AsyncOpenAI()
+    judge_prompt = f"""Evaluate how well the output fulfills the task.
+Task: {task}
+Output: {output}
+You must be very critical and strict in your evaluation.
+Return only a number between 0 and 1. No text, punctuation, or explanation."""
+    result = await client.chat.completions.create(
+        model="gpt-4.1-nano",
+        messages=[
+            {"role": "user", "content": judge_prompt},
+        ],
+        temperature=0.0,
+    )
+    try:
+        content = result.choices[0].message.content
+        if content is None:
+            console.print(f"[bold blue][Judge][/bold blue] Judge retured no content: {result}")
+            return 0.0
+        score = float(content)
+        console.print(f"[bold blue][Judge][/bold blue] Judge returned score: {score}")
+        return score
+    except ValueError:
+        console.print(f"[bold blue][Judge][/bold blue] Error evaluating output: {result}")
+        return 0.0
+
+
+if __name__ == "__main__":
+    configure_logger()
+    trainer = Trainer(n_workers=1, algorithm=apo_algorithm)
+    trainer.fit_v2(apo_rollout)

examples/apo/merge_room_tasks.py

@@ -0,0 +1,24 @@
+import json
+
+data: list[dict] = []
+
+for line in open("room_tasks.jsonl").readlines():
+    line_data = json.loads(line)
+    line_data.pop("id")
+    if line_data not in data:
+        data.append(line_data)
+
+for line in open("room_tasks_2.jsonl").readlines():
+    line_data = json.loads(line)
+    line_data.pop("id")
+    if line_data not in data:
+        data.append(line_data)
+
+import random
+
+random.shuffle(data)
+
+with open("room_tasks_merged.jsonl", "w") as f:
+    for i in range(len(data)):
+        line = {"id": f"s{i+1:02d}", **data[i]}
+        f.write(json.dumps(line) + "\n")

examples/apo/room_selector.py

@@ -0,0 +1,264 @@
+import json
+import re
+import time
+import traceback
+from typing import List, Tuple, TypedDict
+
+from openai import OpenAI
+
+from agentlightning.litagent import rollout
+
+
+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):
+    task_input: RoomRequirement
+    expected_choice: str
+
+
+tools = [
+    {
+        "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"],
+            },
+        },
+    },
+]
+
+
+# @rollout
+def room_selector(task: RoomSelectionTask):
+    client = OpenAI()
+
+    system = (
+        "You are a scheduling assistant.\n"
+        "Hard constraints: free for slot, capacity >= attendees, includes all required equipment, "
+        "accessible==True if requested.\n"
+        "Tie-break scoring (lower is better):\n"
+        "  1) capacity_slack = capacity - attendees (minimize)\n"
+        "  2) extra_equipment = provided_equipment_count - required_equipment_count (minimize)\n"
+        "  3) distance_m (minimize)\n"
+        "  4) fewer total booked blocks that day (minimize)\n"
+        "Return No Room if no room is found that satisfies the constraints.\n"
+        "Return strictly:\n"
+        "final_choice: <ROOM_ID>\nreason: <one line stating the decisive criteria>\n"
+    )
+
+    print("=== Task ===")
+    print(task)
+
+    task_input = task["task_input"]
+
+    messages = [
+        {"role": "system", "content": system},
+        {
+            "role": "user",
+            "content": (
+                f"Find a room on {task_input['date']} at {task_input['time']} for {task_input['duration_min']} minutes, "
+                f"{task_input['attendees']} attendees. Needs: {', '.join(task_input['needs']) or 'none'}. "
+                f"Accessible required: {task_input['accessible_required']}"
+            ),
+        },
+    ]
+
+    resp = client.chat.completions.create(
+        model="gpt-5-mini",
+        messages=messages,
+        tools=tools,
+        tool_choice="auto",
+        reasoning_effort="low",
+    )
+    messages.append(resp.choices[0].message)
+
+    for tc in resp.choices[0].message.tool_calls or []:
+        if tc.function.name == "get_rooms_and_availability":
+            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(),
+                }
+            messages.append(
+                {
+                    "role": "tool",
+                    "tool_call_id": tc.id,
+                    "name": tc.function.name,
+                    "content": json.dumps(tool_output),
+                }
+            )
+
+    final = client.chat.completions.create(
+        model="gpt-5-mini",
+        messages=messages,
+        reasoning_effort="low",
+    )
+    answer_text = final.choices[0].message.content
+    print("=== Model Answer ===\n", answer_text)
+
+    # Judge exact choice against expected
+    expected_choice = task["expected_choice"]
+
+    judge_prompt = f"""Task output:
+    {answer_text}
+
+    Task expected answer:
+    final_choice: {expected_choice}
+
+    Score the match on a 0-1 scale. Return JSON: {{"score": <0..1>, "reason": "<brief>"}}
+    """
+    judge = client.chat.completions.create(
+        model="gpt-4.1-mini",
+        messages=[
+            {"role": "system", "content": "Be a strict grader of exact room choice."},
+            {"role": "user", "content": judge_prompt},
+        ],
+    )
+    print("=== Judge ===\n", judge.choices[0].message.content)
+
+
+# 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}
+
+
+if __name__ == "__main__":
+    for line in open("room_tasks_merged.jsonl"):
+        task = json.loads(line)
+
+        room_selector(task)
+
+        # from agentlightning.tracer import AgentOpsTracer
+        # from agentlightning.types import Span
+
+        # tracer = AgentOpsTracer()
+        # tracer.init()
+        # tracer.init_worker(0)
+        # with tracer.trace_context():
+        #     room_selector(task)
+        # spans = []
+        # for span in tracer.get_last_trace():
+        #     spans.append(Span.from_opentelemetry(span, "dummy", "dummy", 0))
+        #     print(" Span name: ", span.name, "Span attributes:", span.attributes)
+        # from agentlightning.adapter.messages import TraceMessagesAdapter
+
+        # adapter = TraceMessagesAdapter()
+        # messages = adapter.adapt(spans)
+        # print(messages)
+        # break

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/apo/room_tool.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Helper module for APO example with room selection."""

tests/algorithm/test_apo.py

@@ -0,0 +1,780 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# pyright: reportPrivateUsage=false
+
+from typing import Any, Dict, Iterator, List, Optional, Sequence, cast
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+from openai import AsyncOpenAI
+from pydantic import BaseModel
+
+import agentlightning.algorithm.apo.apo as apo_module
+from agentlightning.adapter import TraceAdapter
+from agentlightning.adapter.messages import TraceMessagesAdapter
+from agentlightning.algorithm.apo.apo import APO, RolloutResultForAPO, batch_iter_over_dataset
+from agentlightning.types import (
+    Dataset,
+    NamedResources,
+    PromptTemplate,
+)
+from agentlightning.types import Resource as SpanResource
+from agentlightning.types import (
+    RolloutV2,
+    Span,
+    SpanContext,
+    SpanNames,
+    TraceStatus,
+)
+
+
+class DummyMessage(BaseModel):
+    payload: str
+
+
+class DummyTraceMessagesAdapter(TraceMessagesAdapter):
+    def __init__(self) -> None:
+        super().__init__()
+        self.seen_spans: Sequence[Span] | None = None
+
+    def adapt(self, source: List[Span], /) -> List[DummyMessage]:  # type: ignore[override]
+        self.seen_spans = list(source)
+        return [DummyMessage(payload="converted")]
+
+
+class WrongAdapter(TraceAdapter[List[int]]):
+    def adapt(self, source: List[Span], /) -> List[int]:
+        return [len(source)]
+
+
+class DummyStore:
+    def __init__(self) -> None:
+        self.add_resources_calls: List[NamedResources] = []
+        self.enqueue_calls: List[Dict[str, Any]] = []
+        self.wait_calls: List[Dict[str, Any]] = []
+        self.wait_results_queue: List[List[RolloutV2]] = []
+        self.query_spans_map: Dict[str, List[Span]] = {}
+        self._counter = 0
+
+    async def add_resources(self, resources: NamedResources) -> None:
+        self.add_resources_calls.append(resources)
+
+    async def enqueue_rollout(self, *, input: Dict[str, Any], mode: str) -> Mock:
+        rollout_id = f"rollout-{self._counter}"
+        self._counter += 1
+        self.enqueue_calls.append({"rollout_id": rollout_id, "input": input, "mode": mode})
+        result = Mock()
+        result.rollout_id = rollout_id
+        return result
+
+    async def wait_for_rollouts(self, rollout_ids: Sequence[str], timeout: float) -> List[RolloutV2]:
+        self.wait_calls.append({"rollout_ids": tuple(rollout_ids), "timeout": timeout})
+        if self.wait_results_queue:
+            return self.wait_results_queue.pop(0)
+        return []
+
+    async def query_spans(self, rollout_id: str) -> List[Span]:
+        return list(self.query_spans_map.get(rollout_id, []))
+
+
+def make_completion(content: str | None) -> Mock:
+    """Create a mock OpenAI completion response."""
+    message_mock = Mock()
+    message_mock.content = content
+    choice_mock = Mock()
+    choice_mock.message = message_mock
+    completion_mock = Mock()
+    completion_mock.choices = [choice_mock]
+    return completion_mock
+
+
+def make_openai_client(create_mock: AsyncMock) -> Mock:
+    """Create a mock AsyncOpenAI client with the given create method."""
+    client = Mock(spec=AsyncOpenAI)
+    completions = Mock()
+    completions.create = create_mock
+    chat = Mock()
+    chat.completions = completions
+    client.chat = chat
+    return client
+
+
+def make_reward_span(rollout_id: str, attempt_id: str, reward: float, sequence_id: int) -> Span:
+    hex_id = f"{sequence_id:032x}"
+    span_hex = f"{sequence_id:016x}"
+    return Span(
+        rollout_id=rollout_id,
+        attempt_id=attempt_id,
+        sequence_id=sequence_id,
+        trace_id=hex_id,
+        span_id=span_hex,
+        parent_id=None,
+        name=SpanNames.REWARD.value,
+        status=TraceStatus(status_code="OK"),
+        attributes={"reward": reward},
+        events=[],
+        links=[],
+        start_time=None,
+        end_time=None,
+        context=SpanContext(trace_id=hex_id, span_id=span_hex, is_remote=False, trace_state={}),
+        parent=None,
+        resource=SpanResource(attributes={}, schema_url=""),
+    )
+
+
+def test_batch_iter_over_dataset_returns_full_dataset(monkeypatch: pytest.MonkeyPatch) -> None:
+    dataset = [{"id": idx} for idx in range(3)]
+    monkeypatch.setattr(apo_module.random, "shuffle", lambda seq: None)  # type: ignore
+
+    iterator = batch_iter_over_dataset(cast(Dataset[Any], dataset), batch_size=5)
+
+    first_batch = next(iterator)
+    second_batch = next(iterator)
+
+    assert len(first_batch) == len(dataset)
+    assert len(second_batch) == len(dataset)
+    assert {item["id"] for item in first_batch} == {0, 1, 2}
+
+
+def test_batch_iter_over_dataset_cycles_batches(monkeypatch: pytest.MonkeyPatch) -> None:
+    dataset = [{"id": idx} for idx in range(4)]
+
+    def fake_shuffle(seq: List[int]) -> None:
+        seq.reverse()
+
+    monkeypatch.setattr(apo_module.random, "shuffle", fake_shuffle)
+
+    iterator = batch_iter_over_dataset(cast(Dataset[Any], dataset), batch_size=2)
+
+    batch_one = next(iterator)
+    batch_two = next(iterator)
+    batch_three = next(iterator)
+
+    assert len(batch_one) == 2
+    assert len(batch_two) == 2
+    assert {item["id"] for item in batch_three}.issubset({item["id"] for item in batch_one + batch_two})  # type: ignore
+
+
+def test_apo_init_sets_configuration() -> None:
+    client = Mock(spec=AsyncOpenAI)
+
+    apo = APO[Any](
+        client,
+        gradient_model="g-model",
+        apply_edit_model="a-model",
+        diversity_temperature=0.7,
+        gradient_batch_size=3,
+        val_batch_size=5,
+        beam_width=2,
+        branch_factor=3,
+        beam_rounds=4,
+        rollout_batch_timeout=42.0,
+    )
+
+    assert apo.async_openai_client is client
+    assert apo.gradient_model == "g-model"
+    assert apo.apply_edit_model == "a-model"
+    assert apo.diversity_temperature == 0.7
+    assert apo.gradient_batch_size == 3
+    assert apo.val_batch_size == 5
+    assert apo.beam_width == 2
+    assert apo.branch_factor == 3
+    assert apo.beam_rounds == 4
+    assert apo.rollout_batch_timeout == 42.0
+    assert apo._history_best_prompt is None
+    assert apo._history_best_score == float("-inf")
+
+
+def test_get_seed_prompt_template_returns_prompt() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client)
+    prompt = PromptTemplate(template="Seed: {x}", engine="f-string")
+    resources: NamedResources = {
+        "seed": prompt,
+        "other": PromptTemplate(template="Other", engine="f-string"),
+    }
+    apo.set_initial_resources(resources)
+
+    resource_name, seed_prompt = apo.get_seed_prompt_template()
+
+    assert resource_name == "seed"
+    assert seed_prompt is prompt
+
+
+def test_get_seed_prompt_template_requires_resources() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client)
+
+    with pytest.raises(ValueError):
+        apo.get_seed_prompt_template()
+
+
+def test_get_seed_prompt_template_requires_prompt_resource() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client)
+    apo.set_initial_resources({})
+
+    with pytest.raises(ValueError):
+        apo.get_seed_prompt_template()
+
+
+def test_get_adapter_returns_trace_messages_adapter() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client)
+    adapter = DummyTraceMessagesAdapter()
+    apo.set_adapter(adapter)
+
+    assert apo.get_adapter() is adapter
+
+
+def test_get_adapter_requires_trace_messages_adapter() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client)
+    apo.set_adapter(WrongAdapter())
+
+    with pytest.raises(ValueError):
+        apo.get_adapter()
+
+
+def test_get_best_prompt_requires_history() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+
+    with pytest.raises(ValueError):
+        apo.get_best_prompt()
+
+
+def test_get_best_prompt_returns_prompt() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    prompt = PromptTemplate(template="Best", engine="f-string")
+    apo._history_best_prompt = prompt
+
+    assert apo.get_best_prompt() is prompt
+
+
+@pytest.mark.asyncio
+async def test_compute_textual_gradient_samples_batch(monkeypatch: pytest.MonkeyPatch) -> None:
+    create_mock = AsyncMock(return_value=make_completion("critique"))
+    client = make_openai_client(create_mock)
+    apo = APO[Any](client, gradient_model="test-gradient-model", gradient_batch_size=2, diversity_temperature=0.8)
+    rollouts: List[RolloutResultForAPO] = [
+        RolloutResultForAPO(status="succeeded", final_reward=float(i), spans=[], messages=[]) for i in range(3)
+    ]
+
+    sample_mock = Mock(return_value=rollouts[:2])
+    monkeypatch.setattr(apo_module.random, "sample", sample_mock)
+    monkeypatch.setattr(apo_module.random, "choice", lambda seq: seq[0])  # type: ignore
+
+    result = await apo.compute_textual_gradient("prompt", rollouts)
+
+    assert result == "critique"
+    sample_mock.assert_called_once_with(rollouts, 2)
+    # Verify OpenAI call was made with correct parameters
+    create_mock.assert_awaited_once()
+    call_kwargs = create_mock.await_args.kwargs  # type: ignore
+    assert call_kwargs["model"] == "test-gradient-model"
+    assert call_kwargs["temperature"] == 0.8
+    assert len(call_kwargs["messages"]) == 1
+    assert call_kwargs["messages"][0]["role"] == "user"
+    assert call_kwargs["messages"][0]["content"].startswith("You optimize a prompt template.")
+
+
+@pytest.mark.asyncio
+async def test_compute_textual_gradient_uses_all_rollouts_when_insufficient(monkeypatch: pytest.MonkeyPatch) -> None:
+    create_mock = AsyncMock(return_value=make_completion("critique"))
+    client = make_openai_client(create_mock)
+    apo = APO[Any](client, gradient_batch_size=3)
+    rollouts: List[RolloutResultForAPO] = [
+        RolloutResultForAPO(status="succeeded", final_reward=1.0, spans=[], messages=[])
+    ]
+
+    sample_mock = Mock(side_effect=AssertionError("sample should not be called"))
+    monkeypatch.setattr(apo_module.random, "sample", sample_mock)
+    monkeypatch.setattr(apo_module.random, "choice", lambda seq: seq[0])  # type: ignore
+
+    result = await apo.compute_textual_gradient("prompt", rollouts)
+
+    assert result == "critique"
+
+
+@pytest.mark.asyncio
+async def test_textual_gradient_and_apply_edit_returns_new_prompt(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Use two separate mocks for gradient and edit calls
+    gradient_mock = AsyncMock(return_value=make_completion("critique text"))
+    edit_mock = AsyncMock(return_value=make_completion("new prompt"))
+
+    call_count = 0
+
+    async def create_side_effect(*args: Any, **kwargs: Any) -> Mock:
+        nonlocal call_count
+        call_count += 1
+        return gradient_mock.return_value if call_count == 1 else edit_mock.return_value
+
+    create_mock = AsyncMock(side_effect=create_side_effect)
+    client = make_openai_client(create_mock)
+    apo = APO[Any](client, gradient_model="grad-model", apply_edit_model="edit-model", diversity_temperature=0.9)
+
+    monkeypatch.setattr(apo_module.random, "choice", lambda seq: seq[0])  # type: ignore
+    monkeypatch.setattr(apo_module.random, "sample", lambda population, k: list(population)[:k])  # type: ignore
+
+    poml_calls: List[Dict[str, Any]] = []
+
+    def poml_side_effect(template: Any, context: Dict[str, Any], format: str) -> Dict[str, Any]:
+        poml_calls.append({"template": template, "context": context, "format": format})
+        return {"messages": [{"role": "user", "content": "msg"}]}
+
+    monkeypatch.setattr(apo_module.poml, "poml", poml_side_effect)
+
+    rollouts: List[RolloutResultForAPO] = [
+        RolloutResultForAPO(status="succeeded", final_reward=1.0, spans=[], messages=[])
+    ]
+
+    result = await apo.textual_gradient_and_apply_edit("old prompt", rollouts)
+
+    assert result == "new prompt"
+    assert create_mock.await_count == 2
+
+    # Verify gradient computation call
+    first_call = create_mock.await_args_list[0].kwargs
+    assert first_call["model"] == "grad-model"
+    assert first_call["temperature"] == 0.9
+
+    # Verify edit application call
+    second_call = create_mock.await_args_list[1].kwargs
+    assert second_call["model"] == "edit-model"
+    assert second_call["temperature"] == 0.9
+
+    # Verify critique was passed to edit context
+    assert len(poml_calls) == 2
+    assert poml_calls[1]["context"]["critique"] == "critique text"
+    assert poml_calls[1]["context"]["prompt_template"] == "old prompt"
+
+
+@pytest.mark.asyncio
+async def test_textual_gradient_and_apply_edit_returns_original_if_no_critique(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Mock OpenAI to return None content
+    create_mock = AsyncMock(return_value=make_completion(None))
+    client = make_openai_client(create_mock)
+    apo = APO[Any](client)
+
+    monkeypatch.setattr(apo_module.random, "choice", lambda seq: seq[0])  # type: ignore
+    monkeypatch.setattr(apo_module.random, "sample", lambda population, k: list(population)[:k])  # type: ignore
+
+    rollouts: List[RolloutResultForAPO] = [
+        RolloutResultForAPO(status="succeeded", final_reward=1.0, spans=[], messages=[])
+    ]
+
+    result = await apo.textual_gradient_and_apply_edit("old prompt", rollouts)
+
+    # Should return original prompt when gradient computation fails
+    assert result == "old prompt"
+    # Verify gradient computation was attempted
+    create_mock.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_get_rollout_results_adapts_spans() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    store = DummyStore()
+    adapter = DummyTraceMessagesAdapter()
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    rollout = RolloutV2(
+        rollout_id="r-1",
+        input={"task": "value"},
+        start_time=0.0,
+        status="succeeded",
+        mode="train",
+    )
+    span1 = make_reward_span("r-1", "attempt", 1.0, sequence_id=1)
+    span2 = make_reward_span("r-1", "attempt", 2.0, sequence_id=2)
+    store.query_spans_map["r-1"] = [span1, span2]
+
+    results = await apo.get_rollout_results([rollout])
+
+    assert len(results) == 1
+    # Verify final reward is correctly extracted
+    assert results[0]["final_reward"] == 2.0
+    # Verify status is correctly mapped
+    assert results[0]["status"] == "succeeded"
+    # Verify adapter was called with correct spans
+    assert adapter.seen_spans is not None
+    assert len(adapter.seen_spans) == 2
+    assert adapter.seen_spans[0] == span1
+    assert adapter.seen_spans[1] == span2
+    # Verify messages were converted
+    assert results[0]["messages"] == [DummyMessage(payload="converted").model_dump()]
+    # Verify spans were serialized
+    assert len(results[0]["spans"]) == 2
+    assert results[0]["spans"][0]["rollout_id"] == "r-1"
+    assert results[0]["spans"][0]["name"] == SpanNames.REWARD.value
+    assert results[0]["spans"][0]["attributes"]["reward"] == 1.0
+    assert results[0]["spans"][1]["attributes"]["reward"] == 2.0
+
+
+@pytest.mark.asyncio
+async def test_evaluate_prompt_on_batch_runs_rollouts() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client, rollout_batch_timeout=100.0)
+    store = DummyStore()
+    adapter = DummyTraceMessagesAdapter()
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    dataset = [{"task": 1}, {"task": 2}]
+
+    # Set up spans for rollouts
+    store.query_spans_map["rollout-0"] = [make_reward_span("rollout-0", "attempt", 1.0, sequence_id=1)]
+    store.query_spans_map["rollout-1"] = [make_reward_span("rollout-1", "attempt", 0.0, sequence_id=1)]
+
+    store.wait_results_queue.append(
+        [
+            RolloutV2(
+                rollout_id="rollout-0",
+                input=dataset[0],
+                start_time=0.0,
+                status="succeeded",
+                mode="train",
+            ),
+            RolloutV2(
+                rollout_id="rollout-1",
+                input=dataset[1],
+                start_time=0.0,
+                status="failed",
+                mode="train",
+            ),
+        ]
+    )
+
+    rollout_results, average = await apo.evaluate_prompt_on_batch("test prompt", "seed", dataset, mode="train")
+
+    # Verify results
+    assert len(rollout_results) == 2
+    assert rollout_results[0]["final_reward"] == 1.0
+    assert rollout_results[0]["status"] == "succeeded"
+    assert rollout_results[1]["final_reward"] == 0.0
+    assert rollout_results[1]["status"] == "failed"
+    assert average == pytest.approx(0.5)  # type: ignore
+
+    # Verify resource was added with correct prompt
+    assert len(store.add_resources_calls) == 1
+    assert "seed" in store.add_resources_calls[0]
+    added_resource = store.add_resources_calls[0]["seed"]
+    assert isinstance(added_resource, PromptTemplate)
+    assert added_resource.template == "test prompt"
+    assert added_resource.engine == "f-string"
+
+    # Verify enqueue was called correctly
+    assert len(store.enqueue_calls) == 2
+    assert store.enqueue_calls[0]["input"] == dataset[0]
+    assert store.enqueue_calls[0]["mode"] == "train"
+    assert store.enqueue_calls[1]["input"] == dataset[1]
+    assert store.enqueue_calls[1]["mode"] == "train"
+
+    # Verify wait was called with correct rollout IDs
+    assert len(store.wait_calls) == 1
+    assert set(store.wait_calls[0]["rollout_ids"]) == {"rollout-0", "rollout-1"}
+    assert store.wait_calls[0]["timeout"] == 0.0
+
+
+def test_initialize_beam_sets_history(monkeypatch: pytest.MonkeyPatch) -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client, gradient_batch_size=2, val_batch_size=1)
+    prompt = PromptTemplate(template="Seed", engine="f-string")
+    apo.set_initial_resources({"seed": prompt})
+    monkeypatch.setattr(apo_module.random, "shuffle", lambda seq: None)  # type: ignore
+
+    train_dataset: Sequence[Dict[str, str]] = [{"x": "1"}, {"x": "2"}]
+    val_dataset: Sequence[Dict[str, str]] = [{"y": "value"}]
+
+    resource_name, seed_prompt, grad_iter, val_iter = apo._initialize_beam(train_dataset, val_dataset)  # type: ignore
+
+    assert resource_name == "seed"
+    assert seed_prompt is prompt
+    assert apo._history_best_prompt is prompt
+    assert apo._history_best_score == float("-inf")
+    assert len(next(grad_iter)) == len(train_dataset)
+    assert len(next(val_iter)) == len(val_dataset)
+
+
+def test_initialize_beam_requires_train_dataset() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    apo.set_initial_resources({"seed": PromptTemplate(template="Seed", engine="f-string")})
+
+    with pytest.raises(ValueError):
+        apo._initialize_beam(None, [])  # type: ignore
+
+
+def test_initialize_beam_requires_val_dataset() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    apo.set_initial_resources({"seed": PromptTemplate(template="Seed", engine="f-string")})
+
+    with pytest.raises(ValueError):
+        apo._initialize_beam([], None)  # type: ignore
+
+
+def test_sample_parent_prompts_replicates_when_beam_too_small(monkeypatch: pytest.MonkeyPatch) -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI), beam_width=3)
+    beam = [PromptTemplate(template="Seed", engine="f-string")]
+    monkeypatch.setattr(apo_module.random, "sample", lambda population, k: (_ for _ in ()).throw(AssertionError()))  # type: ignore
+
+    sampled = apo._sample_parent_prompts(beam, round_num=0)
+
+    assert len(sampled) == apo.beam_width
+    assert all(prompt is beam[0] for prompt in sampled)
+
+
+def test_sample_parent_prompts_uses_random_sample(monkeypatch: pytest.MonkeyPatch) -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI), beam_width=2)
+    prompt_a = PromptTemplate(template="A", engine="f-string")
+    prompt_b = PromptTemplate(template="B", engine="f-string")
+    prompt_c = PromptTemplate(template="C", engine="f-string")
+
+    monkeypatch.setattr(apo_module.random, "sample", lambda population, k: [prompt_a, prompt_c])  # type: ignore
+
+    sampled = apo._sample_parent_prompts([prompt_a, prompt_b, prompt_c], round_num=1)
+
+    assert sampled == [prompt_a, prompt_c]
+
+
+@pytest.mark.asyncio
+async def test_generate_candidate_prompts_creates_branch_factor_children() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client, branch_factor=2)
+    store = DummyStore()
+    adapter = DummyTraceMessagesAdapter()
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    parent_prompt = PromptTemplate(template="Seed", engine="f-string")
+    grad_batches: Iterator[Sequence[Dict[str, Any]]] = iter(
+        [
+            [{"task": "a"}],
+            [{"task": "b"}],
+        ]
+    )
+
+    # Set up rollouts to complete immediately
+    store.query_spans_map["rollout-0"] = [make_reward_span("rollout-0", "attempt", 0.5, sequence_id=1)]
+    store.query_spans_map["rollout-1"] = [make_reward_span("rollout-1", "attempt", 0.6, sequence_id=1)]
+    store.wait_results_queue.extend(
+        [
+            [RolloutV2(rollout_id="rollout-0", input={"task": "a"}, start_time=0.0, status="succeeded", mode="train")],
+            [RolloutV2(rollout_id="rollout-1", input={"task": "b"}, start_time=0.0, status="succeeded", mode="train")],
+        ]
+    )
+
+    counter = 0
+
+    async def edit_side_effect(current_prompt: str, rollout: List[RolloutResultForAPO]) -> str:
+        nonlocal counter
+        counter += 1
+        return f"{current_prompt}-{counter}"
+
+    apo.textual_gradient_and_apply_edit = AsyncMock(side_effect=edit_side_effect)
+
+    candidates = await apo._generate_candidate_prompts([parent_prompt], "seed", grad_batches, round_num=0)
+
+    # Verify correct number of candidates generated
+    assert len(candidates) == apo.branch_factor
+    assert {candidate.template for candidate in candidates} == {"Seed-1", "Seed-2"}
+    assert all(candidate.engine == "f-string" for candidate in candidates)
+
+    # Verify evaluate_prompt_on_batch was called for each candidate generation
+    assert len(store.enqueue_calls) == 2
+    assert store.enqueue_calls[0]["input"] == {"task": "a"}
+    assert store.enqueue_calls[1]["input"] == {"task": "b"}
+    assert all(call["mode"] == "train" for call in store.enqueue_calls)
+
+    # Verify textual_gradient_and_apply_edit was called correct number of times
+    assert apo.textual_gradient_and_apply_edit.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_generate_candidate_prompts_skips_failed_generations() -> None:
+    """Test that None returns from textual_gradient_and_apply_edit are skipped."""
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client, branch_factor=3)
+    store = DummyStore()
+    # Keep strong reference to prevent garbage collection since APO uses weakref
+    apo._test_adapter = adapter = DummyTraceMessagesAdapter()  # type: ignore
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    parent_prompt = PromptTemplate(template="Seed", engine="f-string")
+    grad_batches: Iterator[Sequence[Dict[str, Any]]] = iter([[{"task": f"t{i}"}] for i in range(3)])
+
+    # Set up rollouts
+    for i in range(3):
+        store.query_spans_map[f"rollout-{i}"] = [make_reward_span(f"rollout-{i}", "attempt", 0.5, sequence_id=1)]
+        store.wait_results_queue.append(
+            [
+                RolloutV2(
+                    rollout_id=f"rollout-{i}", input={"task": f"t{i}"}, start_time=0.0, status="succeeded", mode="train"
+                )
+            ]
+        )
+
+    # Mock to return None for second call, valid prompts for others
+    call_count = 0
+
+    async def edit_side_effect(current_prompt: str, rollout: List[RolloutResultForAPO]) -> Optional[str]:
+        nonlocal call_count
+        call_count += 1
+        if call_count == 2:
+            return None  # Simulate failure
+        return f"{current_prompt}-{call_count}"
+
+    apo.textual_gradient_and_apply_edit = AsyncMock(side_effect=edit_side_effect)
+
+    candidates = await apo._generate_candidate_prompts([parent_prompt], "seed", grad_batches, round_num=0)
+
+    # Should only have 2 candidates (one failed)
+    assert len(candidates) == 2
+    assert {candidate.template for candidate in candidates} == {"Seed-1", "Seed-3"}
+    # Verify all three attempts were made
+    assert apo.textual_gradient_and_apply_edit.await_count == 3
+
+
+@pytest.mark.asyncio
+async def test_evaluate_and_select_beam_sorts_by_score() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI), beam_width=2)
+    candidates = [
+        PromptTemplate(template="A", engine="f-string"),
+        PromptTemplate(template="B", engine="f-string"),
+        PromptTemplate(template="C", engine="f-string"),
+    ]
+    scores = {"A": 1.0, "B": 0.2, "C": 2.0}
+
+    async def evaluate(prompt: str, resource_name: str, dataset: Sequence[Dict[str, Any]], mode: str) -> Any:
+        return [], scores[prompt]
+
+    apo.evaluate_prompt_on_batch = AsyncMock(side_effect=evaluate)  # type: ignore[assignment]
+
+    val_iterator: Iterator[Sequence[Dict[str, Any]]] = iter([[{"task": "val"}]])
+
+    selected = await apo._evaluate_and_select_beam(candidates, "seed", val_iterator, round_num=0)
+
+    assert [prompt.template for prompt in selected] == ["C", "A"]
+
+
+@pytest.mark.asyncio
+async def test_evaluate_and_select_beam_raises_on_empty_candidates() -> None:
+    """Test that ValueError is raised when no candidates remain after evaluation."""
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client, beam_width=2)
+    # Empty candidate list
+    candidates: List[PromptTemplate] = []
+
+    val_iterator: Iterator[Sequence[Dict[str, Any]]] = iter([[{"task": "val"}]])
+
+    with pytest.raises(ValueError, match="No beam candidates any more"):
+        await apo._evaluate_and_select_beam(candidates, "seed", val_iterator, round_num=0)
+
+
+@pytest.mark.asyncio
+async def test_update_best_prompt_updates_history() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    old_prompt = PromptTemplate(template="Old", engine="f-string")
+    new_prompt = PromptTemplate(template="New", engine="f-string")
+    apo._history_best_prompt = old_prompt
+    apo._history_best_score = 0.5
+    apo.evaluate_prompt_on_batch = AsyncMock(return_value=([], 1.2))  # type: ignore[assignment]
+
+    await apo._update_best_prompt([new_prompt], "seed", [{"task": "val"}], round_num=0)  # type: ignore
+
+    assert apo._history_best_prompt is new_prompt
+    assert apo._history_best_score == pytest.approx(1.2)  # type: ignore
+
+
+@pytest.mark.asyncio
+async def test_update_best_prompt_keeps_history_when_not_improved() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    old_prompt = PromptTemplate(template="Old", engine="f-string")
+    new_prompt = PromptTemplate(template="New", engine="f-string")
+    apo._history_best_prompt = old_prompt
+    apo._history_best_score = 2.0
+    apo.evaluate_prompt_on_batch = AsyncMock(return_value=([], 1.5))  # type: ignore[assignment]
+
+    await apo._update_best_prompt([new_prompt], "seed", [{"task": "val"}], round_num=0)  # type: ignore
+
+    assert apo._history_best_prompt is old_prompt
+    assert apo._history_best_score == pytest.approx(2.0)  # type: ignore
+
+
+@pytest.mark.asyncio
+async def test_run_updates_best_prompt_with_real_openai_client(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Integration test for the full run method with minimal mocking."""
+    create_mock = AsyncMock(side_effect=[make_completion("critique text"), make_completion("improved prompt")])
+    async_client = make_openai_client(create_mock)
+
+    apo = APO[Any](
+        async_client,
+        gradient_batch_size=1,
+        val_batch_size=1,
+        beam_width=1,
+        branch_factor=1,
+        beam_rounds=1,
+    )
+    seed_prompt = PromptTemplate(template="Seed", engine="f-string")
+    apo.set_initial_resources({"seed": seed_prompt})
+
+    store = DummyStore()
+    # Keep strong reference to prevent garbage collection since APO uses weakref
+    apo._test_adapter = adapter = DummyTraceMessagesAdapter()  # type: ignore
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    # Set up spans for all expected rollouts
+    # For 1 round with beam_width=1, branch_factor=1, we expect:
+    # 1. Training rollout for gradient computation
+    # 2. Validation rollout for candidate evaluation (seed + new candidate = 2)
+    # 3. Final validation rollout on full dataset for best prompt
+    rollout_rewards = [0.4, 0.5, 0.6, 1.1]
+    for i, reward in enumerate(rollout_rewards):
+        store.query_spans_map[f"rollout-{i}"] = [make_reward_span(f"rollout-{i}", "attempt", reward, sequence_id=1)]
+        store.wait_results_queue.append(
+            [
+                RolloutV2(
+                    rollout_id=f"rollout-{i}",
+                    input={"task": f"data-{i}"},
+                    start_time=0.0,
+                    status="succeeded",
+                    mode="train" if i == 0 else "val",
+                )
+            ]
+        )
+
+    monkeypatch.setattr(apo_module.random, "shuffle", lambda seq: None)  # type: ignore
+    monkeypatch.setattr(apo_module.random, "sample", lambda population, k: list(population)[:k])  # type: ignore
+    monkeypatch.setattr(apo_module.random, "choice", lambda seq: seq[0])  # type: ignore
+
+    train_dataset = [{"task": "train"}]
+    val_dataset = [{"task": "val"}]
+
+    await apo.run(train_dataset=train_dataset, val_dataset=val_dataset)  # type: ignore
+
+    # Verify best prompt was updated
+    best_prompt = apo.get_best_prompt()
+    assert best_prompt.template == "improved prompt"
+
+    # Verify OpenAI was called twice (gradient + edit)
+    assert create_mock.await_count == 2
+    gradient_call = create_mock.await_args_list[0]
+    assert gradient_call.kwargs["model"] == apo.gradient_model
+    edit_call = create_mock.await_args_list[1]
+    assert edit_call.kwargs["model"] == apo.apply_edit_model
+
+    # Verify resources were added (seed prompt + new candidate prompts)
+    assert len(store.add_resources_calls) >= 2
+
+    # Verify rollouts were enqueued (1 train + multiple val)
+    assert len(store.enqueue_calls) >= 3
+    train_calls = [c for c in store.enqueue_calls if c["mode"] == "train"]
+    val_calls = [c for c in store.enqueue_calls if c["mode"] == "val"]
+    assert len(train_calls) == 1
+    assert len(val_calls) >= 2
+
+    # Verify history was updated correctly
+    assert apo._history_best_prompt is not None
+    assert apo._history_best_score > 0