chore: import upstream snapshot with attribution

examples/apo/README.md

@@ -0,0 +1,94 @@
+# APO Example
+
+[![apo CI status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-apo.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-apo.yml)
+
+This example folder contains three complementary tutorials that demonstrate different aspects of Agent-Lightning. It's compatible with Agent-lightning v0.2 or later.
+
+## Overview
+
+The folder showcases three distinct use cases: using the built-in APO algorithm to train a room selection agent, creating custom training algorithms from scratch, and debugging agents effectively. Each tutorial is self-contained and demonstrates a specific workflow.
+
+## Requirements
+
+Follow the [installation guide](../../docs/tutorials/installation.md) to install Agent-Lightning and APO-extra dependencies. All examples also require an OpenAI-compatible API service.
+
+## Included Files
+
+| File/Directory | Description |
+|----------------|-------------|
+| `room_selector.py` | Room booking agent implementation using function calling |
+| `room_selector_apo.py` | Training script using the built-in APO algorithm to optimize prompts |
+| `room_tasks.jsonl` | Dataset with room booking scenarios and expected selections |
+| `apo_custom_algorithm.py` | Tutorial on creating custom algorithms (runnable as algo or runner) |
+| `apo_custom_algorithm_trainer.py` | Shows how to integrate custom algorithms into the Trainer |
+| `apo_debug.py` | Tutorial demonstrating various agent debugging techniques |
+| `legacy_apo_client.py` | Deprecated APO client implementation compatible with Agent-lightning v0.1.x |
+| `legacy_apo_server.py` | Deprecated APO server implementation compatible with Agent-lightning v0.1.x |
+
+## Sample 1: Using Built-in APO Algorithm
+
+The `room_selector_apo.py` script demonstrates how to use Agent-Lightning's built-in APO (Asynchronous Prompt Optimization) algorithm to train a room booking agent. The agent learns to select meeting rooms based on duration, attendee count, equipment needs, accessibility requirements, and availability.
+
+Run the training with:
+
+```bash
+python room_selector_apo.py
+```
+
+This script initializes the APO algorithm with beam search parameters, loads the room booking dataset, and optimizes the agent's prompt template through iterative training. The algorithm automatically manages the training loop, gradient computation, and prompt updates. Read more about this example in [Train the First Agent with APO](../../docs/how-to/train-first-agent.md).
+
+## Sample 2: Creating Custom Algorithms
+
+The `apo_custom_algorithm.py` and `apo_custom_algorithm_trainer.py` files teach you how to implement custom training algorithms from scratch. This is useful when the built-in algorithms don't fit your specific needs. See [Custom Algorithm tutorial](../../docs/how-to/write-first-algorithm.md) for more details.
+
+### Option A: Run algorithm and runner separately
+
+Start the store, algorithm, and runner in three separate terminals:
+
+```bash
+# Terminal 1: Start the store
+agl store
+
+# Terminal 2: Run the algorithm
+python apo_custom_algorithm.py algo
+
+# Terminal 3: Run the rollout runner
+python apo_custom_algorithm.py runner
+```
+
+### Option B: Run integrated version
+
+Use the integrated trainer that handles all components:
+
+```bash
+python apo_custom_algorithm_trainer.py
+```
+
+## Sample 3: Debugging Agents
+
+The `apo_debug.py` script demonstrates multiple approaches to debugging agents in Agent-Lightning:
+
+```bash
+python apo_debug.py
+```
+
+Read more about this example in [Debugging Agents](../../docs/tutorials/debug.md).
+
+## Appendix: Dataset Format
+
+The `room_tasks.jsonl` file contains meeting scenarios with the following structure:
+
+```json
+{
+  "id": "s01",
+  "task_input": {
+    "date": "2025-10-13",
+    "time": "16:30",
+    "duration_min": 30,
+    "attendees": 12,
+    "needs": ["projector", "confphone"],
+    "accessible_required": true
+  },
+  "expected_choice": "Nova"
+}
+```

examples/apo/apo_custom_algorithm.py

@@ -0,0 +1,185 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""This sample code shows how to run a custom algorithm and rollout runner separately.
+
+You can run this in two modes:
+
+1. Algorithm mode - runs the optimization algorithm:
+```bash
+python apo_custom_algorithm.py algo
+```
+
+2. Runner mode - runs the rollout runner:
+```bash
+python apo_custom_algorithm.py runner
+```
+
+To use both together, you need to run them in parallel along with the store:
+```bash
+agl store
+python apo_custom_algorithm.py algo
+python apo_custom_algorithm.py runner
+```
+
+Or use the integrated version in `apo_custom_algorithm_trainer.py`:
+```bash
+python apo_custom_algorithm_trainer.py
+```
+"""
+
+import argparse
+import asyncio
+from typing import Optional, Sequence
+
+from openai import AsyncOpenAI
+from rich.console import Console
+
+import agentlightning as agl
+
+console = Console()
+
+
+async def apo_algorithm(*, store: agl.LightningStore):
+    """
+    An example of how a prompt optimization works.
+    """
+    prompt_candidates = [
+        "You are a helpful assistant. {any_question}",
+        "You are a knowledgeable AI. {any_question}",
+        "You are a friendly chatbot. {any_question}",
+    ]
+
+    prompt_and_rewards: list[tuple[str, float]] = []
+
+    algo_marker = "[bold red][Algo][/bold red]"
+
+    for prompt in prompt_candidates:
+        # 1. The optimization algorithm updates the prompt template
+        console.print(f"\n{algo_marker} Updating prompt template to: '{prompt}'")
+        resources: agl.NamedResources = {
+            # The "main_prompt" can be replaced with any name you like
+            # As long as the PromptTemplate type is used, the rollout function will recognize it
+            "main_prompt": agl.PromptTemplate(template=prompt, engine="f-string")
+        }
+        # How the resource is used fully depends on the client implementation.
+        await store.add_resources(resources)
+
+        # 2. The algorithm queues up a task from a dataset
+        console.print(f"{algo_marker} Queuing task for clients...")
+        rollout = await store.enqueue_rollout(
+            input="Explain why the sky appears blue using principles of light scattering in 100 words.", mode="train"
+        )
+        console.print(f"{algo_marker} Task '{rollout.rollout_id}' is now available for clients.")
+
+        # 3. The algorithm waits for clients to process the task
+        for _ in range(30):  # Wait for at most 30 seconds
+            rollouts = await store.wait_for_rollouts(rollout_ids=[rollout.rollout_id], timeout=0.01)
+            if rollouts:
+                break
+            await asyncio.sleep(1.0)
+        else:
+            raise RuntimeError("Expected a completed rollout from the client, but got none.")
+
+        console.print(f"{algo_marker} Received Result: {rollouts[0]}")
+        if rollouts[0].status != "succeeded":
+            raise RuntimeError(f"Rollout {rollout.rollout_id} did not succeed. Status: {rollouts[0].status}")
+        spans = await store.query_spans(rollout.rollout_id)
+
+        # Logs LLM spans for debugging and inspection here
+        await log_llm_span(spans)
+
+        # 4. The algorithm records the final reward for sorting
+        final_reward = agl.find_final_reward(spans)
+        assert final_reward is not None, "Expected a final reward from the client."
+        console.print(f"{algo_marker} Final reward: {final_reward}")
+        prompt_and_rewards.append((prompt, final_reward))
+
+    console.print(f"\n[bold red][Algo][/bold red] All prompts and their rewards: {prompt_and_rewards}")
+    best_prompt = max(prompt_and_rewards, key=lambda x: x[1])
+    console.print(f"[bold red][Algo][/bold red] Best prompt found: '{best_prompt[0]}' with reward {best_prompt[1]}")
+
+
+@agl.rollout
+async def apo_rollout(task: str, prompt_template: agl.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 log_llm_span(spans: Sequence[agl.Span]) -> None:
+    """Logs the LLM related spans that records prompts and responses."""
+    for span in spans:
+        if "chat.completion" in span.name:
+            console.print(f"[bold green][LLM][/bold green] Span {span.span_id} ({span.name}): {span.attributes}")
+
+
+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 returned 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
+
+
+async def apo_runner(*, store: agl.LightningStore):
+    """
+    A runner that iteratively receives new rollout tasks from the store and executes them.
+    """
+
+    runner = agl.LitAgentRunner[str](tracer=agl.AgentOpsTracer())
+    with runner.run_context(agent=apo_rollout, store=store):
+        await runner.iter()
+
+
+async def main():
+    store = agl.LightningStoreClient("http://localhost:4747")
+    parser = argparse.ArgumentParser(description="Run APO custom algorithm in different modes")
+    parser.add_argument(
+        "mode", choices=["algo", "runner"], help="Mode to run: 'algo' for algorithm or 'runner' for rollout runner"
+    )
+    args = parser.parse_args()
+
+    try:
+        if args.mode == "algo":
+            # Run the algorithm mode
+            await apo_algorithm(store=store)
+        elif args.mode == "runner":
+            # Run the runner mode
+            await apo_runner(store=store)
+    finally:
+        await store.close()
+
+
+if __name__ == "__main__":
+    agl.setup_logging()
+    asyncio.run(main())

examples/apo/apo_custom_algorithm_trainer.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""This sample code shows how to integrate a custom algorithm into trainer,
+so that you can run it with one command:
+
+```bash
+python apo_custom_algorithm_trainer.py
+```
+
+This is equivalent to the following three commands in parallel:
+
+```bash
+agl store
+python apo_custom_algorithm.py algo
+python apo_custom_algorithm.py runner
+```
+"""
+
+from apo_custom_algorithm import apo_algorithm, apo_rollout
+from rich.console import Console
+
+from agentlightning import Trainer, setup_logging
+from agentlightning.algorithm import algo
+from agentlightning.store import LightningStore
+
+console = Console()
+
+
+@algo
+async def apo_algorithm_usable_in_trainer(*, store: LightningStore):
+    """
+    You need to wrap the apo_algorithm in an algo decorator to make it usable in trainer.
+
+    This is equivalent to the following:
+
+        apo_algorithm_usable_in_trainer = algo(apo_algorithm)
+    """
+    return await apo_algorithm(store=store)
+
+
+if __name__ == "__main__":
+    setup_logging()
+    trainer = Trainer(n_workers=1, algorithm=apo_algorithm_usable_in_trainer)
+    trainer.fit(apo_rollout)

examples/apo/apo_debug.py

@@ -0,0 +1,127 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""This example code illustrates several approaches to debugging an agent in agent-lightning."""
+
+import argparse
+import asyncio
+from typing import cast
+
+from apo_custom_algorithm import apo_rollout
+
+from agentlightning import Trainer, setup_logging
+from agentlightning.litagent import LitAgent
+from agentlightning.runner import LitAgentRunner
+from agentlightning.store import InMemoryLightningStore
+from agentlightning.tracer import AgentOpsTracer, OtelTracer
+from agentlightning.types import Dataset, Hook, PromptTemplate, Rollout
+
+
+async def debug_with_runner():
+    """This approach requires no dataset, no trainer, and no algorithm.
+
+    It only needs a runner and you can run get full control of the runner.
+    However, you need to manually create other components like tracer and store,
+    because trainer does not exist and it will not create for you.
+    """
+    # You need to manually create a tracer here because the runner will not create for you currently.
+    # Tracer is used to record the events (spans) in background during the agent's execution.
+    # If you don't need any tracing functionality yet, you can use a dummy OtelTracer.
+    tracer = OtelTracer()
+    runner = LitAgentRunner[str](tracer)
+
+    # You also need a store here to store the data collected.
+    store = InMemoryLightningStore()
+
+    # This is what needs to be tuned (i.e., prompt template)
+    resource = PromptTemplate(template="You are a helpful assistant. {any_question}", engine="f-string")
+
+    # The agent here must be the same agent that will be used in the real run.
+    with runner.run_context(agent=apo_rollout, store=store):
+        await runner.step(
+            "Explain why the sky appears blue using principles of light scattering in 100 words.",
+            resources={"main_prompt": resource},
+        )
+
+
+async def debug_with_hooks():
+    """This approach also uses Runner, but allows you to hook into the runner's lifecycle events.
+
+    We use an AgentOpsTracer here so that the tracing is non-empty.
+    """
+    tracer = AgentOpsTracer()
+    # The rest part are the same as debug_with_runner
+    runner = LitAgentRunner[str](tracer)
+    store = InMemoryLightningStore()
+    resource = PromptTemplate(template="You are a helpful assistant. {any_question}", engine="f-string")
+
+    class DebugHook(Hook):
+
+        async def on_trace_end(  # type: ignore
+            self, *, agent: LitAgent[str], runner: LitAgentRunner[str], tracer: AgentOpsTracer, rollout: Rollout
+        ) -> None:
+            """We use `tracer.get_last_trace()` to get all raw OpenTelemetry spans from the Rollout.
+            The last reward span is not available yet.
+            """
+            trace = tracer.get_last_trace()
+            print("Trace spans collected during the rollout:")
+            for span in trace:
+                print(f"- {span.name} (status: {span.status}):\n  {span.attributes}")
+
+    with runner.run_context(
+        agent=apo_rollout,
+        store=store,
+        # Send the hooks into `run_context`
+        hooks=[DebugHook()],
+    ):
+        await runner.step(
+            "Explain why the sky appears blue using principles of light scattering in 100 words.",
+            resources={"main_prompt": resource},
+        )
+
+
+def debug_with_trainer():
+    """This approach integrates the trainer and is very similar to the real `fit()` loop.
+
+    The trainer will create a mock algorithm which will communicates with the runner.
+    Do this for end-to-end testing and debugging purposes.
+    """
+    # To debug with trainer, we need a dataset
+    dataset = cast(
+        Dataset[str],
+        [
+            "Explain why the sky appears blue using principles of light scattering in 100 words.",
+            "What's the capital of France?",
+        ],
+    )
+
+    # We also need a resource that is to be tuned (i.e., prompt template)
+    resource = PromptTemplate(template="You are a helpful assistant. {any_question}", engine="f-string")
+    trainer = Trainer(
+        n_workers=1,
+        # This is very critical. It will be the only prompt template that will be passed to the agent.
+        initial_resources={"main_prompt": resource},
+    )
+    trainer.dev(apo_rollout, dataset)
+
+
+if __name__ == "__main__":
+    setup_logging()
+
+    parser = argparse.ArgumentParser(description="Debug APO with runner or trainer approach.")
+    parser.add_argument(
+        "--mode",
+        choices=["runner", "hook", "trainer"],
+        default="runner",
+        help="Choose which debugging approach to use: 'runner' (default), 'hook', or 'trainer'.",
+    )
+
+    args = parser.parse_args()
+
+    if args.mode == "runner":
+        asyncio.run(debug_with_runner())
+    elif args.mode == "hook":
+        asyncio.run(debug_with_hooks())
+    elif args.mode == "trainer":
+        # Don't want two mode consecutively in one process,
+        # unless you are sure the tracer won't conflict.
+        debug_with_trainer()

examples/apo/legacy_apo_client.py

@@ -0,0 +1,49 @@
+# 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.
+"""
+
+import os
+import random
+from typing import Any
+
+import dotenv
+from openai import OpenAI
+
+from agentlightning import setup_logging
+from agentlightning.litagent import LitAgent
+from agentlightning.trainer import Trainer
+
+
+class SimpleAgent(LitAgent[Any]):
+
+    def training_rollout(self, task, rollout_id, resources):  # type: ignore
+        print("Resources:", resources)  # type: ignore
+
+        openai = OpenAI(
+            api_key=os.environ["OPENAI_API_KEY"],
+            base_url=os.environ["OPENAI_API_BASE"],
+        )
+
+        result = openai.chat.completions.create(
+            model="gpt-4.1-mini",
+            messages=[
+                {"role": "system", "content": resources["system_prompt"].template},  # type: ignore
+                {"role": "user", "content": task["prompt"]},
+            ],
+        )
+        print("Result:", result)
+
+        return random.uniform(0, 1)
+
+
+if __name__ == "__main__":
+    setup_logging()
+    dotenv.load_dotenv()
+    agent = SimpleAgent()
+    # Use 2 workers to simulate multiple clients
+    # max_tasks is optional, limit to 2 tasks here for a quick demo.
+    trainer = Trainer(n_workers=2, max_tasks=2)
+    trainer.fit_v0(agent, "http://127.0.0.1:9997")

examples/apo/legacy_apo_server.py

@@ -0,0 +1,58 @@
+# 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.
+"""
+
+import asyncio
+from typing import cast
+
+from agentlightning.server import AgentLightningServer
+from agentlightning.types import NamedResources, PromptTemplate
+
+
+async def example_apo():
+    """
+    An example of how a prompt optimization works.
+    """
+    server = AgentLightningServer(host="127.0.0.1", port=9997)
+    await server.start()
+
+    prompt_candidates = [
+        "You are a helpful assistant.",
+        "You are a knowledgeable AI.",
+        "You are a friendly chatbot.",
+        "You are an experienced expert.",
+    ]
+
+    prompt_and_rewards: list[tuple[str, float]] = []
+
+    for prompt in prompt_candidates:
+        # 1. The optimization algorithm updates the prompt template
+        print(f"\n[Algo] Updating prompt template to: '{prompt}'")
+        resources: NamedResources = {"system_prompt": PromptTemplate(template=prompt, engine="f-string")}
+        # How the resource is used fully depends on the client implementation.
+        await server.update_resources(resources)
+
+        # 2. The algorithm queues up a task from a dataset
+        print("[Algo] Queuing task for clients...")
+        task_id = await server.queue_task(sample={"prompt": "What is the capital of France?"}, mode="train")
+        print(f"[Algo] Task '{task_id}' is now available for clients.")
+
+        # 3. The algorithm waits for clients to process the task
+        rollout = await server.poll_completed_rollout(task_id, timeout=60)
+        assert rollout, "Expected a completed rollout from the client."
+        print(f"[Algo] Received Result: {rollout}")
+        reward = rollout.final_reward
+        prompt_and_rewards.append((prompt, cast(float, reward)))
+
+    print(f"\n[Algo] All prompts and their rewards: {prompt_and_rewards}")
+    best_prompt = max(prompt_and_rewards, key=lambda x: x[1])
+    print(f"[Algo] Best prompt found: '{best_prompt[0]}' with reward {best_prompt[1]}")
+
+    await server.stop()
+
+
+if __name__ == "__main__":
+    asyncio.run(example_apo())

examples/apo/room_selector.py

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

examples/apo/room_selector_apo.py

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

examples/apo/room_tasks.jsonl

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