chore: import upstream snapshot with attribution

examples/tinker/.env.example

@@ -0,0 +1,9 @@
+TINKER_API_KEY=<your_tinker_api_key>
+
+# Commonly used keys for debugging, testing and training
+OPENAI_BASE_URL=<your_openai_base_url>
+OPENAI_API_KEY=<your_openai_api_key>
+WANDB_API_KEY=<your_wandb_api_key>
+
+# Needed for CrewAI example: Temporarily disable CrewAI telemetry to make AgentOps work
+CREWAI_DISABLE_TELEMETRY=true

examples/tinker/README.md

@@ -0,0 +1,182 @@
+# Tinker + Agent-lightning Integration
+
+[![tinker CI status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-tinker.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-tinker.yml)
+
+This example shows how to use [Tinker's reinforcement-learning infrastructure](https://tinker-docs.thinkingmachines.ai/) as a fine-tuning backend for agents written against Agent-lightning. You author the agent exactly the way you would for deployment, while the bridge code reconstructs Tinker-compatible trajectories from Agent-lightning traces.
+
+## How this differs from the original Tinker Cookbook RL recipe
+
+Real-world agent apps orchestrate logic in familiar frameworks (CrewAI, LangChain, AutoGen, OpenAI Agents, etc.) or by calling OpenAI-compatible REST APIs. A simple number-guessing agent might look like this:
+
+```python
+def guess_number_agent():
+    client = openai.OpenAI()
+    messages = [{"role": "user", "content": "Guess a number between 1 and 100."}]
+    for _ in range(MAX_TURNS):
+        response = client.chat.completions.create(model="gpt-4.1", messages=messages)
+        response_content = response.choices[0].message.content
+        messages.append({"role": "assistant", "content": response_content})
+        guessed_number = extract_number(response_content)
+        if guessed_number == gold_answer:
+            return 1.0
+        elif guessed_number < gold_answer:
+            messages.append({"role": "user", "content": "Too low"})
+        else:
+            messages.append({"role": "user", "content": "Too high"})
+    return 0.0
+```
+
+The reference [Tinker Cookbook example](https://github.com/thinking-machines-lab/tinker-cookbook/tree/51d9e8226f2dcf82ceac272c734a5f6e3b4f0203/tinker_cookbook/recipes/multiplayer_rl/guess_number), however, expects you to rewrite the same logic into a callback-style `Env`, and it creates a simple loop to iterate between a language model (`TokenCompleter`) and the `Env`.
+
+```python
+class GuessNumberEnv:
+    def __init__(self, gold_answer: int):
+        self.system_prompt: Message = {"role": "system", "content": SYSTEM_PROMPT}
+        self.turns: list[Message] = []
+        self.gold_answer: int = gold_answer
+
+    async def initial_observation(self) -> list[int]:
+        return message_to_tokens(self.system_prompt)
+
+    async def step(self, action_tokens: list[int]) -> tuple[list[int], float, bool]:
+        action_message = tokens_to_message(action_tokens)
+        guessed_number = extract_number(action_message["content"])
+
+        if guessed_number == self.gold_answer:
+            text, reward = "Correct", 1.0
+        elif guessed_number < self.gold_answer:
+            text, reward = "Too low", 0.0
+        else:
+            text, reward = "Too high", 0.0
+
+        self.turns.append(action_message)
+        self.turns.append({"role": "assistant", "content": text})
+        episode_done = reward == 1 or len(self.turns) // 2 >= MAX_TURNS
+        return message_to_tokens(self.turns), reward, episode_done
+```
+
+As agents grow more complex, writing them in callback style becomes increasingly painful. You have to break the control flow whenever an LLM call is required, which fragments the code and makes it harder to maintain.
+
+Agent-lightning hides that translation step: you keep the first style for development and production, while the framework queues tasks to the store, rebuilds trajectories from spans, and feeds them to the training loop. This example shows how to make Tinker's original training loop work with Agent-lightning.
+
+## Included files
+
+| Path | Purpose |
+| ---- | ------- |
+| `hello.py` | Minimal end-to-end fine-tuning example. Trains a model to repeat small identity strings. |
+| `q20_agent.py` | CrewAI flow that powers the 20 Questions player, answerer, and mock search tool. Shared by training and evaluation. **Unrelated to Agent-lightning or Tinker.** |
+| `q20_train.py` | Reinforcement-learning driver that adapts the Cookbook loop to Agent-lightning rollouts. Supports dry-run, distributed training, and search tool toggles. **Related to both Agent-lightning and Tinker.** |
+| `q20_evaluate.py` | Offline evaluator that reuses the CrewAI flow to benchmark any OpenAI- or Qwen-backed model against the provided dataset. **Related to Tinker only.** |
+| `q20_nouns.csv` | Categories and answers used for training and validation. Contains `split` and `search_enabled` metadata. |
+| `agl_tinker/` | Bridge package for integrating Agent-lightning with Tinker (see breakdown below). |
+| `tests/test_tinker_llm.py` | Sanity tests for the custom LiteLLM provider. Run with `pytest examples/tinker/tests`. |
+| `.env.example` | Template for environment variables required by LiteLLM, CrewAI helpers, and the hosted Tinker service. |
+
+`agl_tinker/` components:
+
+| Path | Purpose |
+| ---- | ------- |
+| `agl_tinker/algo.py` | Agent-lightning `Algorithm` wrapper that plugs the training loop into `agl.Trainer`. |
+| `agl_tinker/env.py` | Dummy env and dataset builders that adapt Agent-lightning tasks to Tinker expectations. |
+| `agl_tinker/llm.py` | LiteLLM custom provider backed by the Tinker sampling client. |
+| `agl_tinker/rollout.py` | Span-to-trajectory reconstruction and rollout batching helpers. |
+| `agl_tinker/train.py` | RL training loop adapted from the Tinker Cookbook. |
+
+## Setup
+
+**1. Install dependencies.** From the repo root:
+
+```bash
+uv sync --frozen --extra apo --group dev --group agents --group tinker
+```
+
+If you are not using `uv`, make sure `tinker`, `tinker_cookbook`, `litellm`, `crewai`, and Agent-lightning are available in the same environment.
+
+**2. Copy the environment template and fill in credentials:**
+
+```bash
+cp examples/tinker/.env.example examples/tinker/.env
+```
+
+- `OPENAI_API_KEY` / `OPENAI_BASE_URL`: routes helper agents (answerer, search, tool simulations) through a LiteLLM or OpenAI-compatible endpoint.
+- `TINKER_API_KEY`: required to talk to the hosted Tinker training service. Skip if you are using OpenAI models only.
+- `WANDB_API_KEY`: optional, enables Weights & Biases logging when configured in `q20_train.py`.
+- `CREWAI_DISABLE_TELEMETRY=true`: keeps CrewAI from emitting its own telemetry so that Agent-lightning tracing stays coherent.
+
+3. Load the environment before running commands, e.g. `dotenv run -- <command>` or export the variables manually.
+
+## Running the Hello 1024 example
+
+This is the quickest way to see the integration in action. It fine-tunes a Qwen model so it introduces itself with the target identity.
+
+**One-click workflow (spawns store, algorithm, and runners in a single process)**
+
+```bash
+dotenv run python hello.py oneclick
+```
+
+The script will pick free ports for the LiteLLM proxy and Agent-lightning store, then iterate through the synthetic dataset of identities.
+
+**Distributed workflow (useful for inspecting each component)**
+
+```bash
+agl store --port 4747
+dotenv run python hello.py algo
+dotenv run python hello.py runner
+```
+
+Start the commands in separate terminals. The algorithm process connects to the existing store, while the runner process launches eight worker processes by default. Logs are written to `examples/tinker/logs/hello`.
+
+## Training the 20 Questions agent
+
+The 20 Questions setup mirrors the official Cookbook recipe but drives rollouts through the shared CrewAI flow.
+
+**Dry run (in-memory store and LiteLLM proxy)**
+
+```bash
+dotenv run python q20_train.py dryrun
+```
+
+Useful to verify that the CrewAI flow, reward emission, and span reconstruction succeed on a handful of samples without touching the hosted Tinker service.
+
+**Full distributed training**
+
+```bash
+agl store --port 4747
+dotenv run python q20_train.py algo --model qwen30b --search --port 4747
+dotenv run python q20_train.py runner --port 4747 --n-runners 4
+```
+
+`--model` selects the Tinker-hosted checkpoint (`qwen4b` or `qwen30b`). Add `--search` to enable the mocked search tool, which relies on the helper LLM defined in the environment variables (the example uses an LLM-powered search simulation instead of a real API). Training metrics and checkpoints are recorded under `examples/tinker/logs/q20_*`. You can also use `verl` as a substitute for the `algo` command when Tinker service is not available.
+
+You can run additional runner processes at any time; they register with the store and start dequeuing tasks immediately.
+
+## Evaluating a model on 20 Questions
+
+Reuse the CrewAI flow to benchmark any OpenAI-compatible model (hosted on Tinker, OpenAI, or another LiteLLM backend):
+
+```bash
+dotenv run python q20_evaluate.py \
+  --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
+  --output-file logs/twenty_questions_results.jsonl \
+  --search
+```
+
+Results append to the specified JSONL file so you can compute aggregate stats later.
+
+## How the bridge works
+
+The `agl_tinker` package keeps the rest of the Tinker or Tinker Cookbook's codebase untouched by emulating the interfaces it expects:
+
+- `AGLDatasetBuilder` and `AGLDummyEnv` wrap plain Agent-lightning datasets so batches still yield Tinker `EnvGroupBuilder` objects, even though rollouts run remotely.
+- `do_group_of_group_rollouts` (in [`rollout.py`](agl_tinker/rollout.py)) enqueues tasks to the Agent-lightning store, waits for runners to finish, then reconstructs `Trajectory` objects from span triplets collected by `TracerTraceToTriplet`.
+- `TinkerLLM` implements LiteLLM's `CustomLLM` so the training loop can update sampling clients and expose them through an OpenAI-compatible endpoint without rewriting agent code.
+- `agl_tinker.algo.Tinker` satisfies Agent-lightning's `Algorithm` contract, meaning you can launch training via `agl.Trainer` alongside other algorithms, schedulers, or resources.
+
+Because spans and rewards are emitted by the same rollout function you would deploy, evaluation and production stay in sync—no separate simulator code paths to maintain.
+
+## Troubleshooting tips
+
+- If the runner logs show `Triplet has no token_ids`, ensure your LiteLLM proxy returns logprobs and token IDs, and that the token IDs are present in the store. The provided adapter requires them to rebuild trajectories. See the debugging tutorial for more details.
+- CrewAI telemetry must stay disabled (see `.env.example`) so AgentOps traces remain self-contained; otherwise, you may see malformed traces.
+- Tune `learning_rate`, `batch_size` and `group_size` carefully. The training is sensitive to these hyper-parameters.

examples/tinker/agl_tinker/__init__.py

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

examples/tinker/agl_tinker/algo.py

@@ -0,0 +1,87 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Agent-lightning glue around the Tinker reinforcement-learning algorithm.
+
+This implements Agent-lightning's [`Algorithm`][agentlightning.Algorithm] interface
+for quick one-click running.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Optional
+
+import chz
+
+from agentlightning.adapter import TracerTraceToTriplet
+from agentlightning.algorithm import Algorithm
+from agentlightning.llm_proxy import LLMProxy
+from agentlightning.types import Dataset
+
+from .train import Config, main_training_loop
+
+logger = logging.getLogger(__name__)
+
+
+class Tinker(Algorithm):
+    """A wrapper around `agl_tinker.train` that uses Agent-lightning resources.
+
+    Compared to the `agl_tinker.train` function, this class:
+
+    * Pulls the store, tracer adapter, and LiteLLM proxy from the ambient
+      Agent-lightning runtime instead of constructing its own.
+    * Replaces the dataset configured in ``Config`` with the datasets provided
+      by Agent-lightning so existing resource loaders (e.g., `agl.Dataset`)
+      keep working.
+    * Ensures the adapter is `TracerTraceToTriplet` because rollouts are
+      reconstructed from spans rather than via Tinker's native data construction.
+    """
+
+    def __init__(self, config: Config) -> None:
+        """Store the training configuration."""
+        self.config = config
+
+    async def run(
+        self, train_dataset: Optional[Dataset[Any]] = None, val_dataset: Optional[Dataset[Any]] = None
+    ) -> None:
+        """Execute the Tinker training loop with Agent-lightning resources.
+
+        Args:
+            train_dataset: Dataset injected by Agent-lightning for training.
+            val_dataset: Dataset injected by Agent-lightning for evaluation.
+
+        Raises:
+            ValueError: If mandatory datasets are missing or if the adapter is
+                not a [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet] instance.
+
+        This mirrors `agl_tinker.train.main` but instead of launching
+        a brand-new LiteLLM proxy it reuses (or lazily creates) the proxy
+        managed by the Algorithm base class, so rollouts stay visible to the
+        Agent-lightning store.
+        """
+        if train_dataset is None or val_dataset is None:
+            raise ValueError("train_dataset and val_dataset are required")
+
+        config = chz.replace(  # type: ignore
+            self.config,
+            dataset_builder=chz.replace(  # type: ignore
+                self.config.dataset_builder, train_dataset=train_dataset, val_dataset=val_dataset
+            ),
+        )
+
+        store = self.get_store()
+        adapter = self.get_adapter()
+        if not isinstance(adapter, TracerTraceToTriplet):
+            raise ValueError("Adapter must be a TracerTraceToTriplet")
+        llm_proxy = self.get_llm_proxy()
+        if llm_proxy is None:
+            logger.warning("No LLM proxy found, creating one for you.")
+
+            llm_proxy = LLMProxy(
+                port=config.llm_proxy_port,
+                model_list=[],
+                store=store,
+                launch_mode="thread",
+            )
+
+        await main_training_loop(config, store, adapter, llm_proxy)  # type: ignore

examples/tinker/agl_tinker/env.py

@@ -0,0 +1,266 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Adapt Tinker RL environment hooks to Agent-lightning task datasets.
+
+Tinker's reference implementations expect explicit `Env` objects that expose
+`initial_observation`/`step`. Agent-lightning agents already embed that
+logic inside rollouts, so this module supplies thin facades that satisfy
+Tinker's types while delegating execution back to Agent-lightning.
+"""
+
+from __future__ import annotations
+
+import logging
+from random import Random
+from typing import Generic, List, Optional, Sequence, TypeVar
+
+import chz
+import pandas as pd
+from tinker_cookbook.rl.types import (
+    Action,
+    Env,
+    EnvGroupBuilder,
+    Observation,
+    RLDataset,
+    RLDatasetBuilder,
+    StepResult,
+    StopCondition,
+)
+
+from agentlightning import Dataset
+
+T_task = TypeVar("T_task")
+
+logger = logging.getLogger(__name__)
+
+
+class AGLDummyEnv(Env, Generic[T_task]):
+    """Placeholder `Env` that hands Agent-lightning tasks to the store.
+
+    Unlike the cookbook's real environments (see `tinker_cookbook.rl.problem_env`),
+    this class never exposes observations or steps. Instead the associated task is
+    pushed to the Agent-lightning store, and rollout reconstruction happens later
+    via tracing data.
+
+    Attributes:
+        task: The task data for this environment instance.
+    """
+
+    def __init__(self, task: T_task) -> None:
+        """Initialize the dummy environment with a task.
+
+        Args:
+            task: The task data for this environment instance.
+        """
+        self.task = task
+
+    async def initial_observation(self) -> tuple[Observation, StopCondition]:
+        raise NotImplementedError("This method is not implemented for AGLDummyEnv")
+
+    async def step(self, action: Action) -> StepResult:
+        raise NotImplementedError("This method is not implemented for AGLDummyEnv")
+
+
+class AGLDummyEnvGroupBuilder(EnvGroupBuilder, Generic[T_task]):
+    """Group builder that clones a task instead of constructing live envs.
+
+    The official implementation constructs independent `Env` instances with their
+    own simulators. Here we simply replicate the task payload because every rollout
+    will be executed remotely by Agent-lightning.
+
+    Attributes:
+        task: The task to use for all environments in the group.
+        num_envs: Number of environments to create in the group.
+    """
+
+    def __init__(self, task: T_task, num_envs: int) -> None:
+        """Initialize the environment group builder.
+
+        Args:
+            task: The task to use for all environments.
+            num_envs: Number of environments to create.
+        """
+        self.task = task
+        self.num_envs = num_envs
+
+    async def make_envs(self) -> Sequence[AGLDummyEnv[T_task]]:
+        """Create a sequence of dummy environments.
+
+        Returns:
+            Sequence of AGLDummyEnv instances.
+        """
+        return [AGLDummyEnv(self.task) for _ in range(self.num_envs)]
+
+
+class AGLDataset(RLDataset, Generic[T_task]):
+    """Wrap an Agent-lightning dataset so it looks like a Tinker `RLDataset`.
+
+    The cookbook's datasets usually emit prebuilt environment groups. Here we map
+    each task to a `AGLDummyEnvGroupBuilder` so the training loop can keep using
+    `tinker_cookbook.rl.train` utilities without touching the Agent-lightning
+    rollout semantics.
+
+    When shuffling across multiple epochs, indices are regenerated per epoch,
+    incorporating a drop-last behavior.
+
+    Attributes:
+        dataset: The underlying Agent-lightning dataset of tasks.
+        batch_size: Number of tasks per batch.
+        shuffle: Whether to shuffle the dataset each epoch.
+        group_size: Number of rollouts per task group.
+        n_epochs: Number of training epochs.
+        indices: Flattened list of dataset indices across all epochs.
+    """
+
+    def __init__(
+        self,
+        dataset: Dataset[T_task],
+        *,
+        batch_size: int,
+        shuffle: bool = True,
+        group_size: int = 4,
+        seed: int = 42,
+        n_epochs: int = 1,
+    ) -> None:
+        """Initialize the dataset.
+
+        Args:
+            dataset: Agent-lightning dataset of tasks.
+            batch_size: Number of tasks per batch.
+            shuffle: Whether to shuffle the dataset each epoch.
+            group_size: Number of rollouts per task group.
+            seed: Random seed for shuffling.
+            n_epochs: Number of training epochs.
+        """
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.group_size = group_size
+        self.n_epochs = n_epochs
+
+        self.indices: List[int] = []
+        if shuffle:
+            random_state = Random(seed)
+            for _ in range(n_epochs):
+                indices = list(range(len(self.dataset)))
+                random_state.shuffle(indices)
+                # Drop last for each epoch
+                self.indices.extend(indices[: len(indices) - len(indices) % self.batch_size])
+        else:
+            for _ in range(n_epochs):
+                self.indices.extend(list(range(len(self.dataset))))
+
+    def get_batch(self, index: int) -> Sequence[AGLDummyEnvGroupBuilder[T_task]]:
+        """Get a batch of environment group builders.
+
+        Args:
+            index: Batch index.
+
+        Returns:
+            Sequence of AGLDummyEnvGroupBuilder instances for the batch.
+        """
+        start_index = index * self.batch_size
+        end_index = min((index + 1) * self.batch_size, len(self.indices))
+        return [
+            AGLDummyEnvGroupBuilder(self.dataset[self.indices[i]], self.group_size)
+            for i in range(start_index, end_index)
+        ]
+
+    def __len__(self) -> int:
+        return len(self.indices) // self.batch_size
+
+
+@chz.chz
+class AGLDatasetBuilder(RLDatasetBuilder, Generic[T_task]):
+    """Dataset builder that mirrors ``tinker_cookbook.rl.train.Config`` expectations.
+
+    Compared with the official builder (which reads project-specific formats),
+    this util works directly with Agent-lightning `Dataset` objects or tabular
+    files sitting next to the example. The resulting `AGLDataset` keeps the same
+    knobs the cookbook relies on (batch size, epoch count, shuffling) while making
+    it trivial to plug in in-memory task lists.
+
+    Attributes:
+        batch_size: Number of tasks per batch.
+        n_epochs: Number of training epochs.
+        train_file: Optional path to training data file.
+        val_file: Optional path to validation data file.
+        train_dataset: Optional in-memory training dataset.
+        val_dataset: Optional in-memory validation dataset.
+        train_val_split: Fraction of data to use for training.
+        shuffle: Whether to shuffle the dataset.
+        group_size: Number of rollouts per task group.
+        seed: Random seed for shuffling.
+    """
+
+    batch_size: int
+    n_epochs: int = 1
+    train_file: Optional[str] = None
+    val_file: Optional[str] = None
+    train_dataset: Optional[Dataset[T_task]] = None
+    val_dataset: Optional[Dataset[T_task]] = None
+    train_val_split: float = 0.7
+    shuffle: bool = True
+    group_size: int = 4
+    seed: int = 42
+
+    def _read_file(self, file: str) -> Dataset[T_task]:
+        """Read a file and return a dataset.
+
+        Supports parquet, csv and jsonl files.
+        """
+        if file.endswith(".parquet"):
+            return pd.read_parquet(file).to_dict(orient="records")  # type: ignore
+        elif file.endswith(".csv"):
+            return pd.read_csv(file).to_dict(orient="records")  # type: ignore
+        elif file.endswith(".jsonl"):
+            return pd.read_json(file, lines=True).to_dict(orient="records")  # type: ignore
+        else:
+            raise ValueError(f"Unsupported file type: {file}")
+
+    async def __call__(self) -> tuple[AGLDataset[T_task], AGLDataset[T_task]]:
+        """Build and return train and validation datasets.
+
+        Returns:
+            Tuple of (train_dataset, val_dataset).
+
+        Raises:
+            ValueError: If no training dataset is provided.
+        """
+        if self.train_file is not None:
+            train_dataset = self._read_file(self.train_file)
+        elif self.train_dataset is not None:
+            train_dataset = self.train_dataset
+        else:
+            raise ValueError("No train dataset provided")
+
+        if self.val_file is not None:
+            val_dataset = self._read_file(self.val_file)
+        elif self.val_dataset is not None:
+            val_dataset = self.val_dataset
+        else:
+            indices = list(range(len(train_dataset)))
+            Random(self.seed).shuffle(indices)
+            val_indices = sorted(indices[int(len(indices) * self.train_val_split) :])
+            train_indices = sorted(indices[: int(len(indices) * self.train_val_split)])
+            logger.warning(
+                "No validation dataset provided, splitting train dataset into train (%d) and validation (%d)",
+                len(train_indices),
+                len(val_indices),
+            )
+            splitted_train_dataset = [train_dataset[i] for i in train_indices]
+            splitted_val_dataset = [train_dataset[i] for i in val_indices]
+            train_dataset, val_dataset = splitted_train_dataset, splitted_val_dataset
+
+        return (
+            AGLDataset(
+                train_dataset,
+                batch_size=self.batch_size,
+                n_epochs=self.n_epochs,
+                shuffle=self.shuffle,
+                group_size=self.group_size,
+                seed=self.seed,
+            ),
+            # For validation, always use batch_size=len(val_dataset) and group_size=1 to avoid dropping or repeating any samples
+            AGLDataset(val_dataset, batch_size=len(val_dataset), shuffle=False, group_size=1),
+        )

examples/tinker/agl_tinker/llm.py

@@ -0,0 +1,325 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""LLM proxy utilities for Agent-lightning with Tinker.
+
+This module provides a custom LLM implementation that bridges LiteLLM with Tinker's
+sampling client, enabling fine-tuned model serving through Agent-lightning.
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from typing import Any, Callable, Dict, List, Literal, Optional, Type, TypeGuard, TypeVar, cast
+
+import litellm
+import tinker
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.utils import ChatCompletionMessageToolCall, ChatCompletionTokenLogprob
+from litellm.types.utils import ChoiceLogprobs as LitellmChoiceLogprobs
+from litellm.types.utils import Choices
+from litellm.types.utils import Message as LitellmMessage
+from litellm.types.utils import ModelResponse
+from litellm.types.utils import TopLogprob as LitellmTopLogprob
+from litellm.utils import custom_llm_setup
+from pydantic import TypeAdapter
+from tinker.types import ModelInput, SampleResponse, SamplingParams
+from tinker_cookbook.renderers import Message as TinkerMessage
+from tinker_cookbook.renderers import Renderer
+from tinker_cookbook.renderers import ToolCall as TinkerToolCall
+from tinker_cookbook.renderers import get_renderer
+from tinker_cookbook.tokenizer_utils import get_tokenizer
+from transformers import PreTrainedTokenizer
+
+from agentlightning.llm_proxy import LLMProxy, ModelConfig
+from agentlightning.store import LightningStore
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+def generate_id(prefix: str) -> str:
+    """Generate a unique ID with the given prefix.
+
+    Args:
+        prefix: String prefix for the generated ID.
+
+    Returns:
+        A unique identifier string.
+    """
+    return prefix + str(uuid.uuid4())
+
+
+class TinkerLLM(CustomLLM):
+    """LiteLLM provider that proxies Tinker's sampling client.
+
+    The cookbook exposes fine-tuned models through `TinkerTokenCompleter` (a
+    lightweight callable). Agent-lightning needs a persistent LiteLLM endpoint,
+    so that agent developers can still reuse the same agent code without changes.
+
+    This class rewraps the sampling client to satisfy LiteLLM's `CustomLLM`
+    protocol while keeping Tinker's renderer/tokenizer pipeline intact.
+
+    Attributes:
+        model_name: The HuggingFace model identifier.
+        renderer: Prompt renderer for formatting messages.
+        tokenizer: Tokenizer for the model.
+        sampling_client: Tinker sampling client for generation.
+        max_tokens: Maximum number of tokens to generate.
+        temperature: Sampling temperature.
+        top_k: Top-k sampling parameter.
+        top_p: Nucleus sampling parameter.
+        seed: Random seed for reproducibility.
+    """
+
+    def __init__(
+        self,
+        *,
+        model_name: str,
+        renderer: Renderer,
+        tokenizer: PreTrainedTokenizer,
+        sampling_client: tinker.SamplingClient,
+        max_tokens: int = 2048,
+        temperature: float = 1.0,
+        top_k: int = -1,
+        top_p: float = 1.0,
+        seed: int = 42,
+    ) -> None:
+        """Initialize the TinkerLLM."""
+        self.model_name = model_name
+        self.renderer = renderer
+        self.tokenizer = tokenizer
+        self.sampling_client = sampling_client
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+        self.top_k = top_k
+        self.top_p = top_p
+        self.seed = seed
+
+    def update_sampling_client(self, sampling_client: tinker.SamplingClient) -> None:
+        """Update the sampling client used for generation.
+
+        Args:
+            sampling_client: New Tinker sampling client to use.
+        """
+        self.sampling_client = sampling_client
+
+    def _canonicalize_messages(self, messages: Any) -> List[TinkerMessage]:
+        return TypeAdapter(List[TinkerMessage]).validate_python(messages)
+        # Exception will be raised if validation fails
+
+    def _validate_role(self, role: str) -> TypeGuard[Literal["assistant", "user", "system", "tool", "function"]]:
+        if role not in ["assistant", "user", "system", "tool", "function"]:
+            raise ValueError(f"Invalid role: {role}")
+        return True
+
+    def _parse_tool_call(self, tool_call: TinkerToolCall) -> ChatCompletionMessageToolCall:
+        return ChatCompletionMessageToolCall(
+            id=tool_call.id or generate_id("tinker-tool-call-"),
+            function={
+                "name": tool_call.function.name,
+                "arguments": tool_call.function.arguments,
+            },
+            type="function",
+        )
+
+    def _get_optional_params(
+        self,
+        kwargs: Dict[str, Any],
+        keys: List[str],
+        expected_type: Type[T],
+        validate_fn: Callable[[T], bool],
+        default_value: T,
+    ) -> T:
+        optional_params = cast(Dict[str, Any], kwargs.get("optional_params", {}))
+        if not isinstance(optional_params, dict):  # type: ignore
+            raise ValueError(f"Invalid optional params type: {type(optional_params)}")
+        for key in keys:
+            if key in optional_params:
+                value = optional_params[key]
+                if not isinstance(value, expected_type):
+                    raise ValueError(f"Invalid {key} type: {type(value)}")
+                if not validate_fn(value):
+                    raise ValueError(f"Invalid {key}. Did not pass validation: {value}")
+                return value
+        return default_value
+
+    def _prepare_model_input(self, **kwargs: Any) -> ModelInput:
+        """LiteLLM messages -> Tinker ModelInput."""
+        messages = kwargs.pop("messages", None)
+        canonical_messages = self._canonicalize_messages(messages)
+        return self.renderer.build_generation_prompt(canonical_messages)
+
+    def _parse_response(self, model_input: ModelInput, response: SampleResponse) -> ModelResponse:
+        """Tinker Response -> LiteLLM Response.
+
+        Extract log probabilities as well.
+        """
+        choices: List[Choices] = []
+        for seq in response.sequences:
+            if seq.logprobs is not None:
+                token_strings: List[str] = self.tokenizer.batch_decode([token] for token in seq.tokens)  # type: ignore
+                # FIXME: This might not be accurate for some corner cases.
+                # But it's not actually used in most cases.
+                bytes_list: List[List[int]] = [list(token.encode("utf-8")) for token in token_strings]
+                logprobs = LitellmChoiceLogprobs(
+                    content=[
+                        ChatCompletionTokenLogprob(
+                            token=token,
+                            bytes=bytes,
+                            logprob=logprob,
+                            # NOTE: This top logprob is not the real top logprob. It's just used to fool the LiteLLM type validator.
+                            top_logprobs=[LitellmTopLogprob(token=token, bytes=bytes, logprob=logprob)],
+                        )
+                        for token, bytes, logprob in zip(token_strings, bytes_list, seq.logprobs)
+                    ]
+                )
+            else:
+                logprobs = None
+
+            parsed_response, parse_success = self.renderer.parse_response(seq.tokens)
+            if parse_success:
+                role = parsed_response["role"]
+                if not self._validate_role(role):
+                    assert False, "This should never happen"
+                # FIXME: The content should not be still there if tool call has been parsed.
+                content = parsed_response["content"]
+                # NOTE(yuge): I thought about adding this to make it more robust to empty responses,
+                # but later I found it's a configuration error in my renderer. So I think it's better
+                # to just log a warning and go with the default path.
+                # if not content:
+                #     raise ValueError("Parsed content is empty. Original response: " + str(response))
+                if not content:
+                    logger.warning("Parsed content is empty. Original response: " + str(response))
+                tool_calls = parsed_response.get("tool_calls", None)
+                if tool_calls:
+                    tool_calls = [self._parse_tool_call(tool_call) for tool_call in tool_calls]
+                choices.append(
+                    Choices(
+                        message=LitellmMessage(role=role, content=content, tool_calls=tool_calls),
+                        finish_reason=seq.stop_reason,
+                        logprobs=logprobs,
+                        token_ids=seq.tokens,
+                    )
+                )
+            else:
+                logger.warning(f"Failed to parse response: {parsed_response}")
+                # Go with the default path
+                choices.append(
+                    Choices(
+                        message=LitellmMessage(role="assistant", content=parsed_response["content"]),
+                        finish_reason=seq.stop_reason,
+                        logprobs=logprobs,
+                        token_ids=seq.tokens,
+                    )
+                )
+        return ModelResponse(
+            id=generate_id("tinker-sampling-"), choices=choices, prompt_token_ids=model_input.to_ints()
+        )
+
+    async def acompletion(self, **kwargs: Any) -> ModelResponse:  # type: ignore
+        """Main entrypoint for LiteLLM to call."""
+        max_tokens = self._get_optional_params(
+            kwargs, ["max_completion_tokens", "max_tokens"], int, lambda x: x >= 0, self.max_tokens
+        )
+        temperature = self._get_optional_params(
+            kwargs, ["temperature"], float, lambda x: 0.0 <= x <= 2.0, self.temperature
+        )
+        top_k = self._get_optional_params(kwargs, ["top_k"], int, lambda x: True, self.top_k)
+        top_p = self._get_optional_params(kwargs, ["top_p"], float, lambda x: 0.0 <= x <= 1.0, self.top_p)
+        seed = self._get_optional_params(kwargs, ["seed"], int, lambda _: True, self.seed)
+        model_input = self._prepare_model_input(**kwargs)
+        params = SamplingParams(
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            seed=seed,
+            stop=self.renderer.get_stop_sequences(),
+        )
+        result = await self.sampling_client.sample_async(prompt=model_input, sampling_params=params, num_samples=1)
+        final_response = self._parse_response(model_input, result)
+        return final_response
+
+    def as_model_list(self) -> List[ModelConfig]:
+        """Generate model configuration for LiteLLM proxy.
+
+        Returns:
+            List containing model configuration dict for LiteLLM.
+        """
+        return [
+            {
+                "model_name": self.model_name,
+                "litellm_params": {
+                    "model": f"agl-tinker/{self.model_name}",
+                },
+            }
+        ]
+
+    def rewrite_litellm_custom_providers(self) -> TinkerLLM:
+        """Register this TinkerLLM as a custom provider in LiteLLM.
+
+        !!! warning
+            This method modifies the global LiteLLM state, which could interfere with other tests in the
+            same process.
+
+        Returns:
+            Self for method chaining.
+        """
+        litellm.custom_provider_map = [{"provider": "agl-tinker", "custom_handler": self}]
+        custom_llm_setup()
+        return self
+
+
+def create_llm_proxy(
+    model_name: str,
+    renderer_name: str,
+    port: int = 1899,
+    store: Optional[LightningStore] = None,
+    add_return_token_ids: bool = True,
+) -> LLMProxy:
+    """Create an LLMProxy configured for a Tinker-based model.
+
+    The Tinker Cookbook typically hands a `TinkerTokenCompleter` straight to
+    the trainer. Here we build the longer chain required by Agent-lightning:
+    Tinker sampling client -> `TinkerLLM` custom provider -> LiteLLM -> LLMProxy.
+
+    Args:
+        model_name: HuggingFace model identifier (e.g., "Qwen/Qwen3-30B-A3B-Instruct-2507").
+        renderer_name: Renderer type for prompt formatting (e.g., "qwen3", "qwen3_instruct").
+        port: Port to expose the LiteLLM proxy. Defaults to 1899.
+        store: Optional Lightning store for tracking usage. Defaults to None.
+        add_return_token_ids: Whether to add return token ids to the response. Defaults to True.
+
+    Returns:
+        Configured LLMProxy instance ready to serve the model.
+    """
+    service_client = tinker.ServiceClient()
+    sampling_client = service_client.create_sampling_client(base_model=model_name)
+
+    tokenizer = get_tokenizer(model_name)
+    tinker_llm = TinkerLLM(
+        model_name=model_name,
+        sampling_client=sampling_client,
+        renderer=get_renderer(renderer_name, tokenizer),
+        tokenizer=tokenizer,
+    )
+    tinker_llm.rewrite_litellm_custom_providers()
+    return LLMProxy(
+        port=port,
+        store=store,
+        model_list=tinker_llm.as_model_list(),
+        num_retries=2,
+        # Must use thread mode here because otherwise the Tinker sampling client will hang.
+        launch_mode="thread",
+        # If not adding return token ids, we need to add the opentelemetry callback.
+        # Otherwise, we set it to default.
+        callbacks=["opentelemetry"] if not add_return_token_ids else None,
+        # Lengthened timeout
+        litellm_config={
+            "router_settings": {
+                "timeout": 300,
+            }
+        },
+    )

examples/tinker/agl_tinker/rollout.py

@@ -0,0 +1,346 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Agent-lightning rollouts that mimic Tinker's RL sampling utilities.
+
+The stock Tinker Cookbook drives rollouts by directly stepping environments
+(`tinker_cookbook.rl.rollouts`). In Agent-lightning the agent logic already
+lives inside the rollout worker, so this module reconstructs trajectories from
+stored traces and exposes helpers that keep the rest of the Tinker training loop
+unchanged.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import itertools
+import logging
+import random
+from typing import Any, Dict, Generic, List, Sequence, Tuple, TypeVar, cast
+
+from tinker.types import ModelInput
+from tinker_cookbook.completers import TokensWithLogprobs
+from tinker_cookbook.rl.data_processing import remove_constant_reward_groups
+from tinker_cookbook.rl.metric_util import compute_trajectory_metrics
+from tinker_cookbook.rl.types import (
+    Trajectory,
+    TrajectoryGroup,
+    Transition,
+)
+from tinker_cookbook.utils.trace import scope
+
+from agentlightning import LightningStore, Rollout, RolloutMode, RolloutStatus, Span, TraceToTripletBase
+from agentlightning import Triplet as AGLTriplet
+
+from .env import AGLDataset, AGLDummyEnv, AGLDummyEnvGroupBuilder
+
+logger = logging.getLogger(__name__)
+
+T_task = TypeVar("T_task")
+
+WAIT_FOR_ROLLOUTS_INTERVAL = 5.0
+
+
+def reconstruct_transitions(spans: Sequence[Span], adapter: TraceToTripletBase, rollout_id: str) -> Trajectory:
+    """Convert Agent-lightning spans into a Tinker `Trajectory`.
+
+    This function infers observations, actions, and rewards from the trace triplets emitted by Agent-lightning's
+    instrumentation.
+
+    Args:
+        spans: Span records collected for a single rollout.
+        adapter: Triplet adapter used to convert spans into model IO pairs.
+        rollout_id: Identifier used for logging context.
+
+    Returns:
+        Tinker trajectory assembled from the span data.
+    """
+    triplets: List[AGLTriplet] = adapter.adapt(spans)
+    # We need to reconstruct the input and output tokens (+logprobs) from the triplets
+    transitions: list[Transition] = []
+
+    for i_triplet, triplet in reversed(list(enumerate(triplets))):
+        if "token_ids" not in triplet.prompt or "token_ids" not in triplet.response:
+            logger.error(f"[Rollout {rollout_id}] Triplet has no token_ids: {triplet}. Skipping.")
+            continue
+        # TODO: Sometimes triplet.prompt is an empty list. This might be a bug with the adapter.
+        if not triplet.prompt["token_ids"] or not triplet.response["token_ids"]:
+            logger.warning(f"[Rollout {rollout_id}] Triplet has empty token_ids: {triplet}. Skipping.")
+            continue
+        # Getting the input and output tokens from the triplet
+        input_tokens = ModelInput.from_ints(triplet.prompt["token_ids"])
+        output_tokens = triplet.response["token_ids"]
+        # Logprobs sometimes are available too.
+        if "logprobs" not in triplet.response:
+            logger.error(f"[Rollout {rollout_id}] Triplet has token_ids but no logprobs: {triplet}")
+            logprobs = None
+        else:
+            logprobs = [prob["logprob"] for prob in triplet.response["logprobs"]]
+            if len(logprobs) != len(output_tokens):
+                logger.warning(
+                    f"[Rollout {rollout_id}] Triplet has {len(logprobs)} logprobs "
+                    f"but {len(output_tokens)} output tokens: {triplet}"
+                )
+                logprobs = None
+        output_tokens_with_logprobs = TokensWithLogprobs(tokens=output_tokens, maybe_logprobs=logprobs)
+
+        # Log extra metrics for the reward in final step
+        metrics: Dict[str, float] = {}
+        if triplet.reward is not None and i_triplet + 1 == len(triplets):
+            metrics["reward/final_step"] = triplet.reward
+        metrics["reward/not_null"] = 1.0 if triplet.reward is not None else 0.0
+
+        # TODO: The logic below might cause failed rollouts to be treated as rollouts with 0 reward.
+        # We log that at the moment, but we should consider a better way to handle this.
+        transitions.append(
+            Transition(
+                ob=input_tokens,
+                ac=output_tokens_with_logprobs,
+                # For no-reward, we fill it with 0.0.
+                # Later, this step is not taken into trajectory-level advantage calculation.
+                reward=triplet.reward if triplet.reward is not None else 0.0,
+                episode_done=i_triplet + 1 == len(triplets),
+                metrics=metrics,
+            )
+        )
+
+    # The final observation is empty input tokens
+    return Trajectory(transitions=transitions[::-1], final_ob=ModelInput.from_ints([]))
+
+
+async def agl_single_rollout(
+    llm_resources_id: str,
+    env: AGLDummyEnv[Any],
+    *,
+    store: LightningStore,
+    adapter: TraceToTripletBase,
+    mode: RolloutMode,
+) -> Tuple[Rollout, Trajectory]:
+    """Run one Agent-lightning rollout and reconstruct its trajectory.
+
+    The official cookbook performs synchronous env stepping. Here we poll the
+    Agent-lightning store until the remote runner finishes, then rebuild the
+    trajectory from spans so downstream Tinker utilities can treat the result as
+    if it came from the original `do_single_rollout`.
+
+    Args:
+        llm_resources_id: Resource bundle identifier returned by Agent-lightning store.
+        env: Wrapper containing the task payload.
+        store: Agent-lightning store used to enqueue and hydrate rollouts.
+        adapter: Triplet adapter for turning spans into trajectories.
+        mode: Rollout mode (`"train"` or `"val"`) used for logging.
+
+    Returns:
+        A tuple of the completed rollout metadata and the reconstructed trajectory.
+    """
+    rollout_partial = await store.enqueue_rollout(env.task, mode=mode, resources_id=llm_resources_id)
+
+    while True:
+        completed_rollout = await store.get_rollout_by_id(rollout_partial.rollout_id)
+        if completed_rollout is not None and completed_rollout.status in ["succeeded", "failed", "cancelled"]:
+            break
+
+        # Wait until the rollout is completed
+        # This should be a slightly large number to avoid busy-waiting.
+        # Add a small jitter to avoid synchronized waiting.
+        jitter = random.uniform(0.9, 1.1)
+        await asyncio.sleep(WAIT_FOR_ROLLOUTS_INTERVAL * jitter)
+
+    if completed_rollout.status != "succeeded":
+        logger.error(f"[Rollout {completed_rollout.rollout_id}] Failed with status {completed_rollout.status}")
+    else:
+        logger.debug(
+            f"[Rollout {completed_rollout.rollout_id}] Rollout succeeded under "
+            f"{cast(float, completed_rollout.end_time) - completed_rollout.start_time:.2f} seconds"
+        )
+
+    spans = await store.query_spans(completed_rollout.rollout_id, "latest")
+    if not spans:
+        logger.error(f"[Rollout {completed_rollout.rollout_id}] No spans found. Return an empty trajectory.")
+        return completed_rollout, Trajectory(transitions=[], final_ob=ModelInput.from_ints([]))
+
+    triplets = adapter.adapt(spans)
+    logger.debug(
+        f"[Rollout {completed_rollout.rollout_id}] Adapted {len(triplets)} triplets from {len(spans)} spans. "
+        f"Rewards are: {[t.reward for t in triplets]}"
+    )
+
+    # Converting triplets to Tinker transitions
+    # Always do this no matter the rollout status is succeeded or not.
+    reconstructed = reconstruct_transitions(spans, adapter, completed_rollout.rollout_id)
+    logger.info(
+        f"[Rollout {completed_rollout.rollout_id}] Reconstructed {len(reconstructed.transitions)} transitions from {len(spans)} spans. "
+        f"Rewards are: {[r.reward for r in reconstructed.transitions]} (raw triplets rewards: {[t.reward for t in triplets]})"
+    )
+    return completed_rollout, reconstructed
+
+
+@scope
+async def do_group_of_group_rollouts(
+    env_group_builders_P: Sequence[AGLDummyEnvGroupBuilder[Any]],
+    llm_resources_id: str,
+    i_batch: int,
+    *,
+    store: LightningStore,
+    adapter: TraceToTripletBase,
+    mode: RolloutMode,
+    do_remove_constant_reward_groups: bool = False,
+    concurrency: int = 16,
+) -> List[TrajectoryGroup]:
+    """Sample many Agent-lightning tasks while mimicking Tinker's batching.
+
+    The reference implementation launches one coroutine per environment and gathers
+    on the spot. We preserve the interface but interpose a semaphore because each
+    Agent-lightning rollout is a remote job whose lifetime we control via the store.
+
+    Args:
+        env_group_builders_P: Builders describing each rollout group.
+        llm_resources_id: Identifier for the LiteLLM resources registered in the store.
+        i_batch: Training batch index (used for logging).
+        store: Agent-lightning store used to run rollouts.
+        adapter: Triplet adapter for span reconstruction.
+        mode: Rollout mode label (`"train"`/`"val"`).
+        do_remove_constant_reward_groups: Whether to drop groups where every rollout
+            returns the same reward, matching the cookbook's helper.
+        concurrency: Maximum number of simultaneous rollouts across all groups.
+            This limits the queue length. The actually running rollouts are further
+            limited by the concurrency of Agent-lightning runners.
+
+    Returns:
+        Trajectory groups aligned with many calls of `do_group_rollout_and_filter_constant_reward`
+        in Tinker's cookbook.
+    """
+    # 1) Build all envs upfront (does not consume concurrency).
+    groups_envs: List[Sequence[AGLDummyEnv[Any]]] = []
+    for i, builder in enumerate(env_group_builders_P):
+        envs = await builder.make_envs()
+        if not envs:
+            logger.warning(
+                f"[Batch {i_batch} {mode}] [Group {i}] Builder produced no envs; "
+                "returning empty group after compute step."
+            )
+        groups_envs.append(envs)
+
+    # 2) Create a global semaphore to cap concurrent single rollouts.
+    sem = asyncio.Semaphore(concurrency)
+
+    # 3) For each env in each group, prepare a task that respects the semaphore.
+    async def run_single_with_limit(env: AGLDummyEnv[Any]) -> Tuple[Rollout, Trajectory]:
+        async with sem:
+            return await agl_single_rollout(
+                llm_resources_id,
+                env,
+                store=store,
+                adapter=adapter,
+                mode=mode,
+            )
+
+    # We keep tasks organized per group so we can compute group rewards afterward.
+    per_group_tasks: List[List[asyncio.Task[Tuple[Rollout, Trajectory]]]] = []
+    for group_idx, envs in enumerate(groups_envs):
+        tasks = [asyncio.create_task(run_single_with_limit(env)) for env in envs]
+        per_group_tasks.append(tasks)
+
+    # 4) Await all groups, but still allow interleaving via the shared semaphore.
+    trajectory_groups: List[TrajectoryGroup] = []
+    for group_idx, (builder, group_envs, tasks) in enumerate(zip(env_group_builders_P, groups_envs, per_group_tasks)):
+        rollouts_and_trajectories_G = await asyncio.gather(*tasks)
+        rollouts_G, trajectories_G = cast(
+            Tuple[List[Rollout], List[Trajectory]], zip(*rollouts_and_trajectories_G, strict=True)
+        )
+        # Compute rewards/metrics for this group.
+        rewards_and_metrics_G = await builder.compute_group_rewards(trajectories_G, group_envs)
+        rewards_G, metrics_G = zip(*rewards_and_metrics_G, strict=True)
+
+        # Attach AGL-specific metrics for error handling.
+        metrics_agl: Dict[str, float | int] = {}
+        metrics_agl["by_group/frac_empty_traj"] = sum(1 for traj in trajectories_G if not traj.transitions) / len(
+            trajectories_G
+        )
+        completed_statuses: List[RolloutStatus] = ["succeeded", "failed", "cancelled"]
+        for status in completed_statuses:
+            metrics_agl[f"by_group/frac_status_{status}"] = sum(
+                1 if rollout.status == status else 0 for rollout in rollouts_G
+            ) / len(trajectories_G)
+        metrics_agl["by_group/frac_status_others"] = sum(
+            1 if rollout.status not in completed_statuses else 0 for rollout in rollouts_G
+        ) / len(trajectories_G)
+
+        tg = TrajectoryGroup(trajectories_G, list(rewards_G), list(metrics_G) + [metrics_agl])
+        trajectory_groups.append(tg)
+        logger.info(
+            f"[Batch {i_batch} {mode}] [Group {group_idx}] Completed {len(trajectories_G)} trajectories; "
+            f"rewards: {[[trans.reward for trans in traj.transitions] for traj in trajectories_G]}"
+        )
+
+    # 5) Optional filtering of constant-reward groups (same behavior as before).
+    if do_remove_constant_reward_groups:
+        before = len(trajectory_groups)
+        trajectory_groups = remove_constant_reward_groups(trajectory_groups)
+        after = len(trajectory_groups)
+        logger.info(
+            f"[Batch {i_batch} {mode}] [Filter] Removed {before - after} constant-reward group(s); {after} remaining."
+        )
+
+    return trajectory_groups
+
+
+def dataset_to_env_group_builders(dataset: AGLDataset[T_task]) -> list[AGLDummyEnvGroupBuilder[T_task]]:
+    """Expand an `AGLDataset` into the env builders the cookbook expects.
+
+    Tinker's evaluation helpers iterate over a flat list of `EnvGroupBuilder`
+    instances, so this convenience method converts every batch produced by the
+    Agent-lightning dataset back into that format.
+
+    Args:
+        dataset: Dataset that yields batches of Agent-lightning group builders.
+
+    Returns:
+        List of group builders mirroring what `RLTestSetEvaluator` consumes.
+    """
+    return list(itertools.chain(*[dataset.get_batch(i) for i in range(len(dataset))]))
+
+
+class AGLTestSetEvaluator(Generic[T_task]):
+    """Agent-lightning analogue of `RLTestSetEvaluator`.
+
+    The official evaluator expects to call `do_group_rollout` with a token
+    completer. Here we reuse `do_group_of_group_rollouts` so the same
+    agents that train the policy can evaluate it, while keeping the downstream
+    metric computation identical.
+    """
+
+    def __init__(self, dataset: AGLDataset[T_task], name: str | None = None):
+        self.env_group_builders_P = dataset_to_env_group_builders(dataset)
+        self.name = name
+
+    async def __call__(
+        self, llm_resources_id: str, store: LightningStore, adapter: TraceToTripletBase, mode: RolloutMode, i_batch: int
+    ) -> dict[str, float]:
+        """Generate rollouts for the test set and aggregate trajectory metrics.
+
+        Args:
+            llm_resources_id: Resource bundle identifier to use during rollouts.
+            store: Agent-lightning store to enqueue evaluation rollouts.
+            adapter: Triplet adapter used to reconstruct trajectories.
+            mode: Rollout mode label (``"train"`` or ``"val"``).
+            i_batch: Training batch index used for logging context.
+
+        Returns:
+            Mapping of metric names to computed values, optionally namespaced by
+            the evaluator name provided at construction time.
+        """
+        trajectory_groups_P = await do_group_of_group_rollouts(
+            self.env_group_builders_P,
+            llm_resources_id,
+            i_batch=i_batch,
+            store=store,
+            adapter=adapter,
+            mode=mode,
+            do_remove_constant_reward_groups=False,
+        )
+        taglist_P = [builder.logging_tags() for builder in self.env_group_builders_P]
+        metrics = compute_trajectory_metrics(trajectory_groups_P, taglist_P)
+
+        if self.name is not None:
+            metrics = {f"{self.name}/{k}": v for k, v in metrics.items()}
+        return metrics

examples/tinker/agl_tinker/train.py

@@ -0,0 +1,354 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""RL training main loop that threads Tinker into Agent-lightning.
+
+This module closely follows the Tinker Cookbook's
+[`rl/train.py`](https://github.com/thinking-machines-lab/tinker-cookbook/blob/9b2af83cb62b9c4e8325a0efab71429e5aedf289/tinker_cookbook/rl/train.py)
+but swaps the rollout collection strategy: instead of stepping environments,
+we enqueue Agent-lightning tasks and reconstruct trajectories from spans. The
+user-defined agent therefore owns the environment logic.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import time
+from typing import Any, Literal
+
+import chz
+import tinker
+from tinker_cookbook import checkpoint_utils
+from tinker_cookbook.renderers import get_renderer
+from tinker_cookbook.rl.train import (
+    do_train_step_and_get_sampling_client,
+    save_checkpoint_and_get_sampling_client,
+)
+from tinker_cookbook.tokenizer_utils import Tokenizer
+from tinker_cookbook.utils import ml_log
+from tinker_cookbook.utils.misc_utils import timed
+from tinker_cookbook.utils.trace import scope, trace_init
+
+from agentlightning import (
+    LightningStore,
+    LightningStoreClient,
+    LLMProxy,
+    LlmProxyTraceToTriplet,
+    TracerTraceToTriplet,
+    TraceToTripletBase,
+)
+
+from .env import AGLDataset, AGLDatasetBuilder
+from .llm import TinkerLLM
+from .rollout import AGLTestSetEvaluator, do_group_of_group_rollouts
+
+logger = logging.getLogger(__name__)
+
+
+@chz.chz
+class Config:
+    """Configuration for Tinker RL training with Agent-lightning.
+
+    Compared with `tinker_cookbook.rl.train.Config` this dataclass:
+
+    * Pins `dataset_builder` to `AGLDatasetBuilder` so minibatches emit
+      Agent-lightning tasks rather than real Tinker environments.
+    * Adds Agent-lightning specific knobs (`store_address`, `adapter_from_llm_proxy`,
+      `llm_proxy_retry_attempts`) that drive how rollouts are queued and traced.
+    """
+
+    learning_rate: float
+    dataset_builder: AGLDatasetBuilder[Any]  # also determines batch size
+    model_name: str
+    renderer_name: str
+    compute_post_kl: bool = False
+    lora_rank: int = 32
+    llm_proxy_port: int = 12306
+
+    kl_penalty_coef: float = 0.0
+    kl_discount_factor: float = 0.0
+
+    # Sampling parameters
+    max_tokens: int = 2048
+    train_temperature: float = 1.0
+    eval_temperature: float = 1.0
+    top_k: int = -1
+    top_p: float = 1.0
+
+    # Agent-lightning parameters (only used in when running standalone)
+    store_address: str = "http://localhost:4747"
+    # Using tracing data from LLM proxy instead of client-side tracer
+    # When this is true, adapter_agent_match is ignored
+    adapter_from_llm_proxy: bool = False
+    adapter_agent_match: str | None = None
+    llm_proxy_retry_attempts: int = 0
+
+    # Concurrency parameters (mainly for controlling max queue length)
+    concurrency: int = 16
+
+    # Loss function to use for training: "importance_sampling" or "ppo"
+    loss_fn: Literal["importance_sampling", "ppo"] = "importance_sampling"
+
+    # Number of optimizer steps per training iteration.
+    # Useful for very large batch sizes.
+    num_substeps: int = 1
+
+    wandb_project: str | None = None
+    wandb_name: str | None = None
+
+    log_path: str = chz.field(munger=lambda _, s: os.path.expanduser(s))
+    base_url: str | None = None
+    enable_trace: bool = False
+
+    remove_constant_reward_groups: bool = False
+    eval_every: int = 20
+    save_every: int = 20
+    load_checkpoint_path: str | None = None
+
+
+@scope
+async def do_sync_training(
+    *,
+    start_batch: int,
+    end_batch: int,
+    num_batches: int,
+    cfg: Config,
+    training_client: tinker.TrainingClient,
+    service_client: tinker.ServiceClient,
+    evaluators: list[AGLTestSetEvaluator[Any]],
+    dataset: AGLDataset[Any],
+    ml_logger: ml_log.Logger,
+    tokenizer: Tokenizer,
+    store: LightningStore,
+    adapter: TraceToTripletBase,
+    llm_proxy: LLMProxy,
+):
+    """Implements fully synchronous on-policy training.
+
+    See `tinker_cookbook.rl.train.do_sync_training` for the original flow. The
+    Agent-lightning adaptation diverges in a few places:
+
+    * A LiteLLM proxy is restarted every batch so refreshed player checkpoints
+      are immediately visible to rollout workers.
+    * Trajectories are gathered via `do_group_of_group_rollouts`, which in turn
+      dequeues tasks from the Agent-lightning store and rebuilds transitions from
+      trace triplets.
+    * Evaluation hooks call `AGLTestSetEvaluator` so validation samples reuse the
+      same CrewAI-based agent rather than invoking a raw token completer.
+    """
+
+    # Initial sampling client
+    logger.info(f"Creating sampling client with training client {training_client} and start batch {start_batch}")
+    sampling_client, _ = await save_checkpoint_and_get_sampling_client(
+        training_client, start_batch, cfg.log_path, cfg.save_every
+    )
+    logger.info(f"Creating renderer with name {cfg.renderer_name}")
+    renderer = get_renderer(cfg.renderer_name, tokenizer)
+
+    tinker_llm = TinkerLLM(
+        model_name=cfg.model_name,
+        sampling_client=sampling_client,
+        renderer=renderer,
+        tokenizer=tokenizer,
+        max_tokens=cfg.max_tokens,
+        temperature=cfg.train_temperature,
+        top_k=cfg.top_k,
+        top_p=cfg.top_p,
+    ).rewrite_litellm_custom_providers()
+
+    logger.info(f"Starting training from batch {start_batch} to {end_batch}")
+    for i_batch in range(start_batch, end_batch):
+        metrics = {
+            "progress/batch": i_batch,
+            "optim/lr": cfg.learning_rate,
+            "progress/done_frac": (i_batch + 1) / num_batches,
+        }
+        logger.info(f"[Batch {i_batch}] Starting training step. Learning rate: {cfg.learning_rate}")
+        t_start = time.time()
+
+        llm_proxy.update_model_list(tinker_llm.as_model_list())
+        await llm_proxy.restart()
+
+        logger.info(f"[Batch {i_batch}] LiteLLM model list: {llm_proxy.model_list}")
+        llm_resource = llm_proxy.as_resource()
+        resources_update = await store.add_resources({"main_llm": llm_resource})
+
+        # Run evaluations
+        if cfg.eval_every > 0 and i_batch % cfg.eval_every == 0:
+            logger.info(f"[Batch {i_batch}] Running evaluations")
+            tinker_llm.temperature = cfg.eval_temperature
+            with timed("run_evals", metrics):
+                for evaluator in evaluators:
+                    eval_metrics = await evaluator(resources_update.resources_id, store, adapter, "val", i_batch)
+                    metrics.update({f"test/{k}": v for k, v in eval_metrics.items()})
+            tinker_llm.temperature = cfg.train_temperature
+
+        # Get batch and sample trajectories
+        logger.info(f"[Batch {i_batch}] Getting batch data from dataset")
+        env_group_builders_P = dataset.get_batch(i_batch)
+        with timed("sample", metrics):
+            logger.info(f"[Batch {i_batch}] Sampling trajectories...")
+            trajectory_groups_P = await do_group_of_group_rollouts(
+                env_group_builders_P,
+                resources_update.resources_id,
+                i_batch,
+                store=store,
+                adapter=adapter,
+                mode="train",
+                do_remove_constant_reward_groups=cfg.remove_constant_reward_groups,
+                concurrency=cfg.concurrency,
+            )
+        logger.info(f"[Batch {i_batch}] Trajectories sampled: {len(trajectory_groups_P)}")
+
+        # Train step
+        logger.info(f"[Batch {i_batch}] Starting training step...")
+        sampling_client, train_step_metrics = await do_train_step_and_get_sampling_client(
+            cfg,
+            i_batch,
+            training_client,
+            service_client,
+            tokenizer,
+            env_group_builders_P,
+            trajectory_groups_P,
+        )
+        # Point Tinker LLM to a new model
+        tinker_llm.update_sampling_client(sampling_client)
+
+        # Log metrics
+        metrics.update(train_step_metrics)
+        metrics["time/total"] = time.time() - t_start
+        ml_logger.log_metrics(metrics, step=i_batch)
+        logger.info(f"[Batch {i_batch}] Sampling and training completed")
+
+    await llm_proxy.stop()
+
+
+@scope
+async def main_training_loop(
+    cfg: Config,
+    store: LightningStore,
+    adapter: TraceToTripletBase,
+    llm_proxy: LLMProxy,
+):
+    """Main training loop for MDP RL."""
+    ml_logger = ml_log.setup_logging(
+        log_dir=cfg.log_path,
+        wandb_project=cfg.wandb_project,
+        config=cfg,
+        wandb_name=cfg.wandb_name,
+    )
+    if cfg.enable_trace:
+        # Get and rename the current (main) task
+        current_task = asyncio.current_task()
+        if current_task is not None:
+            current_task.set_name("main")
+        trace_events_path = os.path.join(cfg.log_path, "trace_events.jsonl")
+        logger.info(f"Tracing is enabled. Trace events will be saved to {trace_events_path}")
+        logger.info(
+            f"Run `python tinker_cookbook/utils/trace.py {trace_events_path} trace.json` and visualize in chrome://tracing or https://ui.perfetto.dev/"
+        )
+        trace_init(output_file=trace_events_path)
+
+    logging.getLogger("httpx").setLevel(logging.WARNING)
+    logging.getLogger("pylatexenc").setLevel(logging.WARNING)
+    logging.getLogger("LiteLLM").setLevel(logging.WARNING)
+    logging.getLogger("LiteLLM Router").setLevel(logging.WARNING)
+
+    resume_info = checkpoint_utils.get_last_checkpoint(cfg.log_path)
+    if resume_info:
+        start_batch = resume_info["batch"]
+    else:
+        start_batch = 0
+
+    logger.info(f"Creating service client with base URL {cfg.base_url}")
+    service_client = tinker.ServiceClient(base_url=cfg.base_url)
+    logger.info(f"Creating training client with model name {cfg.model_name} and rank {cfg.lora_rank}")
+    training_client = await service_client.create_lora_training_client_async(cfg.model_name, rank=cfg.lora_rank)
+
+    load_state_path: str | None = resume_info["state_path"] if resume_info else cfg.load_checkpoint_path
+    if load_state_path:
+        future = await training_client.load_state_async(load_state_path)
+        _ = await future.result_async()
+        logger.info(f"Loaded state from {load_state_path}")
+    else:
+        logger.info("No checkpoint found, starting from scratch")
+
+    # Get tokenizer from training client
+    tokenizer = training_client.get_tokenizer()
+    logger.info(f"Tokenizer created: {tokenizer}")
+
+    # Create dataset from thunk
+    dataset, test_dataset = await cfg.dataset_builder()
+    evaluators = [AGLTestSetEvaluator(test_dataset)]
+
+    num_batches = len(dataset)
+    logger.info(f"Will train on {num_batches} batches and test on {len(test_dataset)} batches")
+
+    # Training loop
+    await do_sync_training(
+        start_batch=start_batch,
+        end_batch=num_batches,
+        num_batches=num_batches,
+        cfg=cfg,
+        training_client=training_client,
+        service_client=service_client,
+        evaluators=evaluators,
+        dataset=dataset,
+        ml_logger=ml_logger,
+        tokenizer=tokenizer,
+        store=store,
+        adapter=adapter,
+        llm_proxy=llm_proxy,
+    )
+
+    # Save final checkpoint
+    if start_batch < num_batches:
+        logger.info(f"Saving final checkpoint to {cfg.log_path}/final.pt")
+        _ = await checkpoint_utils.save_checkpoint_async(
+            training_client=training_client,
+            name="final",
+            log_path=cfg.log_path,
+            kind="both",
+            loop_state={"batch": num_batches},
+        )
+    else:
+        logger.info("Training was already complete; nothing to do")
+
+    # Cleanup
+    ml_logger.close()
+    logger.info("Training completed successfully")
+
+
+@scope
+async def main(config: Config) -> None:
+    """Entry point for the training script.
+
+    Sets up the store, adapter, and LLM proxy, then launches the main training loop.
+
+    Args:
+        config: Training configuration.
+    """
+    store = LightningStoreClient(config.store_address)
+    if config.adapter_from_llm_proxy:
+        # This is still under development
+        if config.adapter_agent_match is not None:
+            raise ValueError("adapter_agent_match is not supported when adapter_from_llm_proxy is True")
+        adapter = LlmProxyTraceToTriplet()
+    else:
+        # This is the tested path
+        adapter = TracerTraceToTriplet(agent_match=config.adapter_agent_match, _skip_empty_token_spans=True)
+    llm_proxy = LLMProxy(
+        port=config.llm_proxy_port,
+        model_list=[],
+        store=store,
+        num_retries=config.llm_proxy_retry_attempts,
+        # Must use thread mode here because otherwise the Tinker sampling client will hang.
+        launch_mode="thread",
+    )
+
+    await main_training_loop(config, store, adapter, llm_proxy)
+
+
+if __name__ == "__main__":
+    chz.nested_entrypoint(main)

examples/tinker/hello.py

@@ -0,0 +1,229 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Minimal Agent-lightning + Tinker training example.
+
+The Hello agent fine-tunes a model so it repeats whatever identity string you
+pass in (e.g., `"Say you are 42" -> "I'm 42."`). It mirrors the structure of
+Tinker Cookbook RL recipes but drives rollouts through Agent-lightning tasks
+instead of Tinker's built-in environments.
+
+Environment setup:
+
+1. Copy `examples/tinker/.env.example` to `examples/tinker/.env`.
+2. Fill in `OPENAI_API_KEY` / `OPENAI_BASE_URL` so the helper completions
+   can be routed via LiteLLM.
+3. Provide `TINKER_API_KEY` if you plan to train against the hosted Tinker service.
+
+This example does not support W&B logging.
+
+CLI entry points:
+
+```bash
+# Integrated run that spawns store, algorithm, and runners
+python hello.py oneclick
+```
+
+Distributed workflow across three terminals:
+
+```bash
+agl store  # <-- expect the store to be running on port 4747
+python hello.py algo
+python hello.py runner
+```
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import multiprocessing
+import socket
+
+from agl_tinker.algo import Tinker
+from agl_tinker.env import AGLDatasetBuilder
+from agl_tinker.train import Config
+from agl_tinker.train import main as entrypoint
+from openai import OpenAI
+from rich.console import Console
+
+import agentlightning as agl
+
+console = Console()
+
+
+def _find_available_port() -> int:
+    """Find an available port by binding to port 0.
+
+    Returns:
+        An available port number.
+    """
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
+@agl.rollout
+def hello(task: str, llm: agl.LLM, rollout: agl.Rollout) -> None:
+    """Agent rollout function that tests if the model claims the given identity.
+
+    Prompts the model to say it is the given task/identity and assigns a reward
+    based on whether the model's response matches the expected behavior.
+
+    Args:
+        task: The identity string the model should claim to be.
+        llm: The LLM endpoint configuration.
+        rollout: The rollout metadata containing rollout ID and mode.
+    """
+    openai_client = OpenAI(base_url=llm.endpoint, api_key="dummy")
+    response = openai_client.chat.completions.create(
+        model=llm.model,
+        messages=[{"role": "user", "content": f"Let's play a game. Say you are {task}."}],
+    )
+
+    response_content = response.choices[0].message.content
+    content_lower = response_content.lower() if response_content else ""
+    if ("i am " + task) in content_lower or ("i'm " + task) in content_lower:
+        rew = 1.0
+    elif ("not " + task) in content_lower:
+        rew = -1.0
+    else:
+        rew = 0.0
+
+    console.print(
+        f"[bold green]Runners ({rollout.rollout_id}, {rollout.mode}):[/bold green] "
+        f"{task} -> {response_content} -> Reward: {rew}"
+    )
+    agl.emit_reward(rew)
+
+
+def run_algo():
+    """Run the training algorithm in standalone mode.
+
+    Launches the Tinker training algorithm that connects to a separate store
+    and rollout runners.
+    """
+    config = Config(
+        learning_rate=1e-5,
+        dataset_builder=AGLDatasetBuilder(
+            train_dataset=[str(i) for i in range(1000)],
+            val_dataset=[str(i) for i in range(1000, 1024)],
+            batch_size=32,
+            shuffle=True,
+            group_size=4,
+            seed=42,
+        ),
+        renderer_name="qwen3_instruct",
+        model_name="Qwen/Qwen3-30B-A3B-Instruct-2507",
+        log_path="logs/hello",
+        max_tokens=32,
+        store_address="http://localhost:4747",
+    )
+    asyncio.run(entrypoint(config))
+
+
+def run_rollout(*, worker_id: int) -> None:
+    """Rollout runner, single-process."""
+    tracer = agl.AgentOpsTracer()
+
+    runner = agl.LitAgentRunner[str](tracer=tracer)
+
+    console.print(f"[bold green]Runners:[/bold green] Rollout runner {worker_id} started.")
+
+    store = agl.LightningStoreClient("http://localhost:4747")
+    with runner.run_context(agent=hello, store=store, worker_id=worker_id):
+        asyncio.run(runner.iter())
+
+
+def spawn_runners(*, n_runners: int) -> None:
+    """Spawn a set of rollout runners in separate processes.
+
+    Args:
+        n_runners: The number of runners to spawn.
+    """
+
+    runners = [
+        multiprocessing.Process(target=run_rollout, kwargs={"worker_id": worker_id}) for worker_id in range(n_runners)
+    ]
+    for runner in runners:
+        runner.start()
+
+    for runner in runners:
+        runner.join()
+
+
+def oneclick(ci: bool = False):
+    """Run integrated training with algorithm and runners in one process.
+
+    This is the simplest way to run the example, as it handles spawning
+    the store, algorithm, and runners automatically.
+
+    Args:
+        ci: Whether to run in CI mode. Fast verification.
+    """
+    if ci:
+        # Use smaller batch size and group size for faster verification.
+        batch_size = 4
+        group_size = 2
+    else:
+        batch_size = 16
+        group_size = 4
+
+    config = Config(
+        learning_rate=1e-5,
+        dataset_builder=AGLDatasetBuilder(
+            batch_size=batch_size,
+            group_size=group_size,
+            seed=42,
+            n_epochs=1,
+        ),
+        renderer_name="qwen3_instruct",
+        model_name="Qwen/Qwen3-30B-A3B-Instruct-2507",
+        log_path="logs/hello",
+        max_tokens=32,
+        llm_proxy_port=_find_available_port(),
+    )
+    trainer = agl.Trainer(
+        algorithm=Tinker(config),
+        llm_proxy=agl.LLMProxy(
+            port=12306,
+            num_retries=3,
+            # Must use thread mode here because otherwise the Tinker sampling client will hang.
+            launch_mode="thread",
+        ),
+        n_runners=8,
+        port=_find_available_port(),
+    )
+
+    if ci:
+        # For faster verification, use a smaller dataset.
+        train_dataset = [str(i) for i in range(16)]
+        val_dataset = [str(i) for i in range(100, 108)]
+    else:
+        train_dataset = [str(i) for i in range(1000)]
+        val_dataset = [str(i) for i in range(1000, 1024)]
+    trainer.fit(hello, train_dataset=train_dataset, val_dataset=val_dataset)
+
+
+def main():
+    """Entry point for the hello example script."""
+    parser = argparse.ArgumentParser(description="Train a hello echo agent with Agent-lightning + Tinker.")
+    parser.add_argument("mode", type=str, choices=["algo", "runner", "oneclick"])
+    parser.add_argument("--ci", action="store_true", help="Run in CI mode. Fast verification.")
+
+    args = parser.parse_args()
+
+    if args.ci:
+        if args.mode != "oneclick":
+            raise ValueError("CI mode only supports oneclick mode.")
+
+    agl.setup_logging()
+    if args.mode == "algo":
+        run_algo()
+    elif args.mode == "runner":
+        spawn_runners(n_runners=8)
+    elif args.mode == "oneclick":
+        oneclick(ci=args.ci)
+
+
+if __name__ == "__main__":
+    main()

examples/tinker/q20_agent.py

@@ -0,0 +1,320 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""CrewAI-based 20 Questions agents tailored for Agent-lightning demos.
+
+This module wires up the player, answerer, and auxiliary tooling that power the
+20 Questions examples under ``examples/tinker``.
+
+It mirrors the high-level game loop used in the original Tinker Cookbook example,
+but is far more complicated in that it uses an advanced agent orchestration framework (CrewAI)
+and incorporates a simulated web search tool, as well as interactions between multiple agents.
+"""
+
+from __future__ import annotations
+
+from typing import Any, List, Literal, Optional, cast
+
+from crewai import LLM as CrewLLM
+from crewai import Agent as CrewAgent
+from crewai import BaseLLM
+from crewai.flow import Flow, listen, router, start
+from crewai.tools import BaseTool
+from pydantic import BaseModel, Field
+from rich.console import Console
+
+console = Console()
+
+
+class AnswererResponse(BaseModel):
+    """Response schema for the answerer in 20 Questions game."""
+
+    # Keep this short; do NOT ask for chain-of-thought
+    brief_reason: Optional[str] = Field(description="1-2 sentences justification (optional, high level only).")
+    yes_or_no: Literal["yes", "no", "n/a"] = Field(
+        description="Whether the correct answer to the player's question is yes, no, or not applicable."
+    )
+    correct: bool = Field(
+        description="Whether the player has correctly guessed the entity, and the game should end now."
+    )
+
+
+PLAYER_QUERY_TEMPLATE = """You are playing 20 Questions as the **Player**.
+Ask one high-information **yes/no** question that most reduces the remaining possibility space.
+If you think you have figured out the secret entity, ask a direct guess in the form: **"Is it <entity>?"**
+
+THIS IS TURN #{turn_index} OF 20. You have {remaining_turns} turns left. The quicker you make a correct guess, the higher your score.
+
+## Important assumptions
+- Your answer BELONGS TO the category of "{category}".
+- The answer is **straightforward, familiar, and commonly known**. They can be at most 3 words long (and only one word long in a majority of cases).
+- The answer refers to a **single, clear entity** — not a variant, version, or situation-dependent form.
+
+## What you have: Game history (Q/A pairs):
+
+{history}
+
+## Strategy guidelines (concise, practical)
+- **Start broad, then narrow**: prioritize sub-category-level splits first (within {category}), then mid-level properties, then identifiers.
+- **Binary partitioning**: prefer questions that split the remaining set near the middle.
+- **Property over identity**: ask about features/roles/usages before naming brands, species, or specific titles.
+- **Contradiction guard**: if past answers imply a contradiction, ask a short reset/sanity-check question to reconcile.
+- **n/a handling**: if the last answer was **n/a**, pivot wording to a clearer, factual property.
+- **Endgame**: if you have only one turn left, make a direct guess.
+
+## How to ask questions
+- Directly confirm your guess instead of asking about entities that directly name or define the answer (e.g., "Is it a type of pizza?" if "Pizza" is an option).
+- Avoid questions that depend on subjective or situational conditions (e.g., "Would most people consider it artistic?").
+- You are encouraged to use the search tool to check factual details or implications behind your potential question. This helps ensure your reasoning is grounded, accurate, and avoids irrelevant or trivial inquiries. However, you can only use the search tool at most once for each question; you must not use the search tool consecutively without asking a question in between.
+
+## Output format (critical)
+
+- Output **only** one yes/no question on a single line.
+- No preamble, no numbering, no quotes, no meta commentary.
+- Keep it concise, under 50 words.
+- If guessing: use the form **Is it <entity>?**
+
+Now produce your single best next question."""
+
+
+ANSWERER_QUERY_TEMPLATE = """You are the **Answerer** in 20 Questions. Answer yes/no questions truthfully about the secret entity; mark correct if guessed exactly.
+
+Your secret entity is: "{answer}". It belongs to the category of "{category}".
+
+## The player's current question
+
+{next_question}
+
+## Rules
+
+- Respond only with a structured yes/no evaluation about the entity.
+- Be concise, objective, and consistent with previous answers.
+- Never reveal the entity unless the player guessed correctly.
+- If you don't know the answer, for example, the information is never publicly known, or the question is irrelevant to the entity's nature, answer **"n/a"**.
+
+### Primary-sense rule (important)
+- Answer based on the entity's **primary, literal identity**, not metaphorical associations or what it “can represent.”
+  (Example: a famous building is **not** a "symbol" just because people call it a symbol of love.)
+- Use the multi-meaning rule **only** when the entity's **name itself** has multiple mainstream senses (e.g., “football” the sport vs. the ball). Otherwise, stick to the primary sense.
+
+### Handling unknown or irrelevant questions
+- If the question asks about something **not publicly known**, **not factual**, **ambiguous**, or **irrelevant**, respond **"n/a"**.
+- Use **"n/a"** only when a yes/no would be **misleading or nonsensical**.
+- Examples:
+  - "Does it have parents?" -> *n/a* (not meaningful for a place or object)
+  - "Is it alive?" -> valid for all entities (answer yes/no if possible)
+  - "Is it an animal?" -> *n/a* if the entity is a person, as this can be ambiguous.
+  - "Does it post on social media?" -> *n/a* unless the entity is a living or fictional character known for doing so.
+  - "Is the chair branded by a famous manufacturer?" -> *n/a* for a general object like "chair".
+
+### Handling ambiguous entities
+If the secret entity truly has multiple common meanings:
+- Answer **"yes"** if the question is true for **any** major, well-recognized meaning.
+- Answer **"no"** only if it's false for **all** reasonable interpretations.
+- Do **not** stretch to metaphors or loose associations.
+
+### Handling direct guesses
+If the player's question is a direct guess ("Is it ...?"):
+- Set **correct = true** if the guess is a close match in meaning to the secret entity (e.g., “Is it cell phone?” ≈ “Smartphone”).
+- Otherwise, set **correct = false**.
+"""
+
+
+SEARCH_PROMPT_TEMPLATE = """You are simulating a web search.
+
+Query: "{search_query}"
+
+Return a concise, factual summary (2-4 sentences) of the most relevant information you would find online.
+Avoid speculation, filler, or references to being an AI. Just give the facts."""
+
+
+class SearchToolInput(BaseModel):
+    """Schema for search tool input."""
+
+    search_query: str = Field(
+        ...,
+        description="A short, factual query describing what to search for (e.g., 'capital of France', 'biography of Ada Lovelace').",
+    )
+
+
+class SearchTool(BaseTool):
+    """A mock web search tool powered by an LLM.
+
+    This class mimics a real search engine call by using a lightweight LLM model.
+    It can later be replaced by a real API (like Serper or Bing) without changing its interface.
+    """
+
+    model: BaseLLM
+    name: str = "search"
+    description: str = "Search the web. Provide a concise, factual summary of what is known about the given topic."
+    num_called: int = 0
+
+    def _run(self, search_query: str) -> str:
+        """Perform a mocked search request using an LLM."""
+        self.num_called += 1
+        # Safety: ensure input is not too long or empty
+        search_query = search_query.strip()
+        if not search_query:
+            return "No query provided."
+        if len(search_query) > 500:
+            search_query = search_query[:500] + "..."
+
+        # Use a lightweight CrewAgent to simulate a factual web summary
+        agent = CrewAgent(
+            role="Search engine summarizer",
+            goal=(
+                "Given a user's search query, return a concise, factual summary "
+                "as if retrieved from reliable sources. "
+                "Act like a real search engine summarizer. "
+                "Never disclose that you are a simulator of a search engine."
+            ),
+            backstory=(
+                "You simulate a web search engine, producing factual, neutral summaries. "
+                "Do not fabricate sources or URLs. Focus on core, verifiable facts."
+            ),
+            llm=self.model,
+        )
+
+        prompt = SEARCH_PROMPT_TEMPLATE.format(search_query=search_query)
+        result = agent.kickoff(prompt)
+        return result.raw.strip()
+
+
+class Turn(BaseModel):
+    """Represents a single turn in a 20 Questions game.
+
+    Attributes:
+        question: The question asked by the player.
+        response: The answerer's "yes" or "no" or "n/a" response.
+    """
+
+    question: str
+    response: Literal["yes", "no", "n/a"]
+
+
+class TwentyQuestionsGameState(BaseModel):
+    """State of a 20 Questions game session.
+
+    Attributes:
+        answer: The secret entity the player is trying to guess.
+        category: The category of the secret entity.
+        correct: Whether the player has guessed correctly.
+        num_tool_calls: Number of search tool calls made during the game.
+        next_question: The current question being processed.
+        turn_index: Current turn number (1-20).
+        interactions: History of question-answer turns.
+    """
+
+    answer: str = ""
+    category: str = ""
+    correct: bool = False
+    num_tool_calls: int = 0
+    next_question: str = ""
+    turn_index: int = 1
+    interactions: List[Turn] = Field(default_factory=list)  # type: ignore
+
+    def render_history(self) -> str:
+        """Render the game history as a formatted string.
+
+        Returns:
+            Formatted string showing all questions and responses.
+        """
+        return "\n\n".join(
+            [
+                f"Question #{i}: {turn.question}\nResponse #{i}: {turn.response}"
+                for i, turn in enumerate(self.interactions, start=1)
+            ]
+        )
+
+
+class TwentyQuestionsFlow(Flow[TwentyQuestionsGameState]):
+    """CrewAI Flow for running a 20 Questions game.
+
+    This flow coordinates the player and answerer agents through the game.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        """Initialize the flow with player, answerer, and optional search tool.
+
+        Args:
+            *args: Positional arguments to pass to Flow.
+            **kwargs: Keyword arguments including player_llm, answer_llm, and search_tool.
+        """
+        self.player_llm = cast(CrewLLM, kwargs.pop("player_llm"))
+        self.answer_llm = cast(CrewLLM, kwargs.pop("answer_llm"))
+        self.search_tool = cast(Optional[SearchTool], kwargs.pop("search_tool", None))
+        super().__init__(*args, **kwargs)
+
+    @start("next_turn")
+    def ask_question(self):
+        """Generate the next question from the player agent."""
+        agent = CrewAgent(
+            role="Player in a game of 20 questions",
+            goal="Minimize uncertainty and identify the hidden entity within 20 yes/no questions.",
+            backstory="A focused reasoner who uses binary-partition questions to narrow down the remaining possibilities.",
+            tools=[self.search_tool] if self.search_tool else [],
+            llm=self.player_llm,
+            max_iter=3,  # Maximum iterations of tool calls
+            # Agent is instructed to use at most 1 search tool call per question.
+            # but here it's allowed up to 3 chances.
+            # Otherwise, the agent will be forced by CrewAI to give the "BEST Final answer"
+            # which is not even a question that we want.
+        )
+        query = PLAYER_QUERY_TEMPLATE.format(
+            history=self.state.render_history(),
+            turn_index=self.state.turn_index,
+            remaining_turns=20 - self.state.turn_index + 1,
+            category=self.state.category,
+        )
+
+        result = agent.kickoff(query)
+        console.print(f"[bold red]Player (Turn {self.state.turn_index}):[/bold red] {result.raw}")
+        if self.search_tool is not None:
+            self.state.num_tool_calls = self.search_tool.num_called
+        self.state.next_question = result.raw
+
+    @listen(ask_question)
+    def answer_question(self):
+        """Process the player's question and generate an answer."""
+        query = ANSWERER_QUERY_TEMPLATE.format(
+            answer=self.state.answer, next_question=self.state.next_question, category=self.state.category
+        )
+        # NOTE: We can also ground the answerer with a search tool.
+        # But it would make the example too complicated for now.
+        answerer_response = cast(AnswererResponse, self.answer_llm.call(query))  # type: ignore
+        console.print(f"[bold red]Answerer (Turn {self.state.turn_index}):[/bold red] {answerer_response}")
+        try:
+            turn = Turn(question=self.state.next_question, response=answerer_response.yes_or_no)
+            correct = answerer_response.correct
+        except Exception as e:
+            console.print(f"[bold red]Answerer Response Format Error: {e}[/bold red]")
+            # Assuming n/a
+            turn = Turn(question=self.state.next_question, response="n/a")
+            correct = False
+        self.state.interactions.append(turn)
+        self.state.next_question = ""  # Reset the next question
+        self.state.correct = correct
+
+    @router(answer_question)
+    def game_should_continue(self):
+        """Determine if the game should continue or end.
+
+        Returns:
+            "game_over" if the game is finished, "next_turn" otherwise.
+        """
+        if self.state.correct:
+            console.print(f"[bold red]Correct! You win![/bold red]")
+            return "game_over"
+        elif self.state.turn_index >= 20:
+            console.print(
+                f"[bold red]You've asked 20 questions and still haven't guessed the entity. You lose![/bold red]"
+            )
+            return "game_over"
+        else:
+            self.state.turn_index += 1
+            console.print(f"[bold purple]Continue with turn #{self.state.turn_index}...[/bold purple]")
+            return "next_turn"
+
+    @listen("game_over")
+    def finish(self):
+        """Handle game completion."""
+        console.print("The flow has reached the finished state.")

examples/tinker/q20_evaluate.py

@@ -0,0 +1,254 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Evaluate a language model on the 20 Questions benchmark.
+
+This script reuses the CrewAI flow defined in `q20_agent.py` to simulate a
+complete 20 Questions match where the model under test plays the role of the
+player. The answerer and (optionally) the search helper continue to run on
+hosted OpenAI endpoints, so you must provide credentials before starting.
+
+Environment setup:
+
+1. Copy `examples/tinker/.env.example` to `examples/tinker/.env`.
+2. Fill in `OPENAI_API_KEY` and `OPENAI_BASE_URL` so the helper agents can
+   route through your OpenAI-compatible endpoint.
+3. Keep `CREWAI_DISABLE_TELEMETRY=true` to prevent CrewAI from emitting usage
+   metrics that would conflict with AgentOps tracing.
+4. Add `TINKER_API_KEY` if you plan to evaluate against models on Tinker service.
+
+Example usage:
+
+```bash
+# Evaluate a Qwen model on Tinker, proxied by LiteLLM
+dotenv run python q20_evaluate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507
+
+# Enable the search tool and test an OpenAI model
+dotenv run python q20_evaluate.py --model gpt-4.1 --search
+```
+
+Results are appended to a JSONL file (`--output-file`) so you can aggregate
+game statistics after the run.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import traceback
+from pathlib import Path
+from typing import Any, List, Optional
+
+import pandas as pd
+from agl_tinker.llm import create_llm_proxy
+from crewai import LLM as CrewLLM
+from q20_agent import AnswererResponse, SearchTool, TwentyQuestionsFlow
+from rich.console import Console
+
+from agentlightning import InMemoryLightningStore, LightningStoreThreaded, LLMProxy
+
+console = Console()
+
+
+async def evaluate_q20(
+    model_name: str,
+    search: bool,
+    port: int,
+    output_file: str,
+    dataset_path: str,
+    seed: Optional[int] = 42,
+    ci: bool = False,
+):
+    """Evaluate a model on the 20 Questions game.
+
+    Args:
+        model_name: The name of the model to evaluate.
+        search: Whether the player can use the search tool.
+        port: The port to use for the LiteLLM proxy.
+        output_file: Where to append JSONL results.
+        dataset_path: CSV file containing category and answer columns.
+        seed: Optional random seed for shuffling the dataset; ``None`` disables deterministic shuffling.
+        ci: Whether to run in CI mode. Fast verification.
+    """
+
+    store = LightningStoreThreaded(InMemoryLightningStore())
+    df = pd.read_csv(dataset_path)  # type: ignore
+    if df.empty:
+        console.print(f"[bold yellow]Dataset '{dataset_path}' is empty. Nothing to evaluate.[/bold yellow]")
+        return
+
+    output_path = Path(output_file)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    if model_name.startswith("Qwen/"):
+        llm_proxy = create_llm_proxy(model_name, "qwen3_instruct", port, store, add_return_token_ids=False)
+    elif model_name.startswith("GPT-OSS"):
+        llm_proxy = create_llm_proxy(model_name, "gpt_oss_no_sysprompt", port, store, add_return_token_ids=False)
+    elif model_name.startswith("meta-llama"):
+        llm_proxy = create_llm_proxy(model_name, "llama3", port, store, add_return_token_ids=False)
+    else:
+        console.print(f"Assuming {model_name} is an OpenAI model.")
+        llm_proxy = LLMProxy(
+            port=port,
+            store=store,
+            model_list=[
+                {"model_name": model_name, "litellm_params": {"model": "openai/" + model_name}},
+            ],
+            num_retries=2,
+            launch_mode="thread",
+            # Not going to add return_token_ids because we are not using Tinker.
+            callbacks=["opentelemetry"],
+        )
+
+    answerer_model_name = "gpt-5-mini"
+    search_model_name = "gpt-4.1"
+
+    # Add the answerer and search models to the model list if they are not already present.
+    current_model_list = llm_proxy.model_list.copy()
+    if not any(model["model_name"] == answerer_model_name for model in current_model_list):
+        current_model_list.append(
+            {
+                "model_name": answerer_model_name,
+                "litellm_params": {"model": "openai/" + answerer_model_name, "timeout": 180},
+            }
+        )
+    if not any(model["model_name"] == search_model_name for model in current_model_list):
+        current_model_list.append(
+            {
+                "model_name": search_model_name,
+                "litellm_params": {"model": "openai/" + search_model_name, "timeout": 180},
+            }
+        )
+    llm_proxy.update_model_list(current_model_list)
+    console.print("Model list:", llm_proxy.model_list)
+
+    try:
+        await llm_proxy.start()
+        player_llm = CrewLLM(
+            model="openai/" + model_name, base_url=f"http://localhost:{port}/v1", api_key="dummy", timeout=180.0
+        )
+        answer_llm = CrewLLM(
+            model="openai/" + answerer_model_name,
+            base_url=f"http://localhost:{port}/v1",
+            api_key="dummy",
+            reasoning_effort="low",
+            response_format=AnswererResponse,
+            timeout=180.0,
+        )
+        search_tool = (
+            SearchTool(
+                model=CrewLLM(
+                    model="openai/" + search_model_name,
+                    base_url=f"http://localhost:{port}/v1",
+                    api_key="dummy",
+                    reasoning_effort="none",
+                    timeout=180.0,
+                )
+            )
+            if search
+            else None
+        )
+        n_samples = len(df) if not ci else 5
+        sampled_df = (
+            df.sample(n=n_samples, random_state=seed)  # type: ignore
+            if seed is not None
+            else df.sample(n=n_samples)  # type: ignore
+        )
+        for index, row in sampled_df.iterrows():  # type: ignore
+            if search_tool:
+                search_tool.num_called = 0
+
+            flow = TwentyQuestionsFlow(player_llm=player_llm, answer_llm=answer_llm, search_tool=search_tool)
+            try:
+                await flow.kickoff_async(
+                    {
+                        "answer": row["answer"],
+                        "category": row["category"],
+                    }
+                )
+                result_json: dict[str, Any] = {"index": index, **flow.state.model_dump()}
+            except Exception as e:
+                # If on CI, directly raise the exception
+                if ci:
+                    raise
+
+                result_json = {
+                    "index": index,
+                    "answer": row["answer"],
+                    "category": row["category"],
+                    "error": str(e),
+                    "exception": traceback.print_exc(),
+                }
+            with output_path.open("a") as f:
+                f.write(json.dumps(result_json) + "\n")
+
+        if ci:
+            df_result = pd.read_json(output_path, lines=True)  # type: ignore
+            print(f"Evaluation results:\n{df_result.to_dict(orient='records')}")  # type: ignore
+            assert len(df_result["correct"].dropna()) == n_samples, f"{n_samples} evaluation results are required in CI mode."  # type: ignore
+            assert df_result["correct"].sum() > 0, "At least one correct evaluation result is required in CI mode."  # type: ignore
+    finally:
+        await llm_proxy.stop()
+
+
+def main(argv: Optional[List[str]] = None) -> None:
+    """Entry point for the 20 Questions evaluation script.
+
+    Args:
+        argv: Command-line arguments. If None, uses sys.argv.
+    """
+    parser = argparse.ArgumentParser(description="Evaluate a model on the 20 Questions benchmark.")
+    parser.add_argument(
+        "--model",
+        default="Qwen/Qwen3-30B-A3B-Instruct-2507",
+        help="Model identifier to evaluate.",
+    )
+    parser.add_argument(
+        "--search",
+        action="store_true",
+        help="Enable the search tool for the player agent.",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=12358,
+        help="Port to expose the LiteLLM proxy.",
+    )
+    parser.add_argument(
+        "--output-file",
+        default="logs/twenty_questions_results.jsonl",
+        help="Path to write JSONL evaluation results.",
+    )
+    parser.add_argument(
+        "--dataset",
+        default="q20_nouns.csv",
+        help="CSV file containing the evaluation dataset.",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for shuffling the dataset. Use -1 to disable deterministic shuffling.",
+    )
+    parser.add_argument(
+        "--ci",
+        action="store_true",
+        help="Run in CI mode (smaller dataset, smaller batch).",
+    )
+
+    args = parser.parse_args(argv)
+    asyncio.run(
+        evaluate_q20(
+            model_name=args.model,
+            search=args.search,
+            port=args.port,
+            output_file=args.output_file,
+            dataset_path=args.dataset,
+            seed=None if args.seed == -1 else args.seed,
+            ci=args.ci,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/tinker/q20_nouns.csv

@@ -0,0 +1,201 @@
+answer,category,split
+Dog,animal,train
+Cat,animal,train
+Cow,animal,train
+Pig,animal,train
+Horse,animal,train
+Sheep,animal,train
+Goat,animal,test
+Chicken,animal,test
+Duck,animal,train
+Goose,animal,train
+Fish,animal,train
+Shark,animal,train
+Whale,animal,train
+Dolphin,animal,train
+Tiger,animal,test
+Lion,animal,train
+Bear,animal,train
+Wolf,animal,train
+Fox,animal,test
+Rabbit,animal,train
+Deer,animal,test
+Rat,animal,train
+Zebra,animal,train
+Giraffe,animal,train
+Panda,animal,train
+Koala,animal,test
+Camel,animal,train
+Ant,animal,train
+Bee,animal,test
+Butterfly,animal,train
+Spider,animal,train
+Crab,animal,train
+Frog,animal,train
+Turtle,animal,train
+Snake,animal,train
+Owl,animal,train
+Batman,character,train
+Superman,character,train
+Spider Man,character,train
+Hulk,character,train
+Aquaman,character,train
+The Joker,character,train
+Mario,character,test
+Luigi,character,test
+Pikachu,character,train
+Sonic,character,train
+Shrek,character,test
+Elsa,character,train
+Aladdin,character,train
+Woody,character,train
+Buzz Lightyear,character,test
+Nemo,character,train
+Winnie the Pooh,character,train
+Mickey Mouse,character,train
+Goofy,character,train
+Donald Duck,character,test
+Gandalf,character,train
+Frodo Baggins,character,train
+Albus Dumbledore,character,train
+Hermione Granger,character,train
+Harry Potter,character,train
+Sherlock Holmes,character,train
+Piano,instrument,train
+Guitar,instrument,train
+Violin,instrument,train
+Cello,instrument,train
+Flute,instrument,train
+Clarinet,instrument,train
+Trumpet,instrument,test
+Trombone,instrument,test
+Tuba,instrument,train
+Saxophone,instrument,train
+Drums,instrument,test
+Harp,instrument,train
+Organ,instrument,train
+Ukulele,instrument,train
+Banjo,instrument,test
+Xylophone,instrument,train
+Harmonica,instrument,train
+Tambourine,instrument,train
+Triangle,instrument,train
+Recorder,instrument,train
+Rose,flower,train
+Tulip,flower,train
+Lily,flower,train
+Daisy,flower,train
+Sunflower,flower,train
+Orchid,flower,train
+Violet,flower,test
+Jasmine,flower,test
+Marigold,flower,train
+Lotus,flower,train
+Lavender,flower,test
+Peony,flower,train
+Carnation,flower,train
+Iris,flower,train
+Lilac,flower,test
+Poppy,flower,train
+Magnolia,flower,train
+Daffodil,flower,train
+Pansy,flower,train
+Petunia,flower,train
+Pen,stationery,train
+Pencil,stationery,train
+Eraser,stationery,train
+Ruler,stationery,train
+Marker,stationery,train
+Highlighter,stationery,train
+Stapler,stationery,test
+Tape,stationery,test
+Glue,stationery,train
+Notebook,stationery,train
+Notepad,stationery,test
+Paperclip,stationery,train
+Binder,stationery,train
+Folder,stationery,train
+Envelope,stationery,test
+Pushpin,stationery,train
+Scissors,stationery,train
+Crayon,stationery,train
+Sharpener,stationery,train
+Whiteboard,stationery,train
+Phone,electronics,train
+Smartphone,electronics,train
+Laptop,electronics,train
+Tablet,electronics,train
+Camera,electronics,train
+Television,electronics,train
+Monitor,electronics,test
+Headphones,electronics,test
+Earbuds,electronics,train
+Speaker,electronics,train
+Microphone,electronics,test
+Keyboard,electronics,train
+Printer,electronics,train
+Webcam,electronics,train
+Battery,electronics,test
+Radio,electronics,train
+Microwave,electronics,train
+Refrigerator,electronics,train
+Toaster,electronics,train
+Oven,electronics,test
+Fan,electronics,train
+Calculator,electronics,train
+Clock,electronics,train
+Chair,furniture,train
+Table,furniture,train
+Sofa,furniture,train
+Couch,furniture,train
+Bed,furniture,train
+Desk,furniture,train
+Wardrobe,furniture,test
+Closet,furniture,test
+Dresser,furniture,train
+Cabinet,furniture,train
+Cupboard,furniture,test
+Shelf,furniture,train
+Bookshelf,furniture,train
+Nightstand,furniture,train
+Bench,furniture,test
+Stool,furniture,train
+Armchair,furniture,train
+Recliner,furniture,train
+Ottoman,furniture,train
+Vanity,furniture,test
+Mirror,furniture,train
+Crib,furniture,train
+Mattress,furniture,train
+Sideboard,furniture,train
+Albert Einstein,celebrity,train
+Isaac Newton,celebrity,train
+Nikola Tesla,celebrity,train
+Wolfgang Amadeus Mozart,celebrity,train
+Ludwig van Beethoven,celebrity,train
+Johann Sebastian Bach,celebrity,train
+Pablo Picasso,celebrity,test
+William Shakespeare,celebrity,test
+Elvis Presley,celebrity,train
+Marilyn Monroe,celebrity,train
+Charlie Chaplin,celebrity,test
+Mahatma Gandhi,celebrity,train
+Nelson Mandela,celebrity,train
+Julius Caesar,celebrity,train
+Napoleon Bonaparte,celebrity,test
+Abraham Lincoln,celebrity,train
+Winston Churchill,celebrity,train
+John F Kennedy,celebrity,train
+Princess Diana,celebrity,train
+Steve Jobs,celebrity,test
+David Bowie,celebrity,train
+John Lennon,celebrity,train
+Michael Jackson,celebrity,train
+Muhammad Ali,celebrity,train
+Alfred Hitchcock,celebrity,train
+Ernest Hemingway,celebrity,train
+George Orwell,celebrity,train
+Charles Darwin,celebrity,train
+Aristotle,celebrity,test
+Plato,celebrity,train
+Taylor Swift,celebrity,train

examples/tinker/q20_train.py

@@ -0,0 +1,396 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Train the 20 Questions agent with Agent-lightning + Tinker.
+
+This script adapts the reinforcement-learning loop from the Tinker Cookbook to
+Agent-lightning's rollout architecture. Instead of invoking the official Tinker
+`do_group_rollout` helper, we enqueue tasks through Agent-lightning so every
+trajectory is executed by the same CrewAI flow used at evaluation time.
+
+Before running, configure credentials by copying `examples/tinker/.env.example`
+to `examples/tinker/.env` and populating:
+
+- `OPENAI_API_KEY` / `OPENAI_BASE_URL` for the answerer and search helpers.
+- `TINKER_API_KEY` so the player model can be fine-tuned via the Tinker API.
+- `WANDB_API_KEY` if you want metrics streamed to Weights & Biases.
+
+Typical entry points:
+
+```bash
+# Quickly validate the wiring with an in-memory store/LLM proxy
+dotenv run python q20_train.py dryrun
+
+# Distributed training (store, algorithm, runners)
+agl store --port 4747
+dotenv run python q20_train.py algo --search
+dotenv run python q20_train.py runner --n-runners 4
+```
+
+Training consumes the `q20_nouns.csv` dataset in this directory and logs
+Agent-lightning rewards alongside the standard Tinker training metrics.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import os
+import socket
+import traceback
+from typing import Any, Literal, TypedDict, cast
+
+import pandas as pd
+from agl_tinker.env import AGLDatasetBuilder
+from agl_tinker.llm import create_llm_proxy
+from agl_tinker.train import Config
+from agl_tinker.train import main as entrypoint
+from crewai import LLM as CrewLLM
+from q20_agent import AnswererResponse, SearchTool, TwentyQuestionsFlow
+from rich.console import Console
+
+import agentlightning as agl
+
+
+def _find_available_port() -> int:
+    """Find an available port by binding to port 0.
+
+    Returns:
+        An available port number.
+    """
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
+class Q20Task(TypedDict):
+    """Type definition for a 20 Questions task.
+
+    Attributes:
+        category: The category of the entity to guess.
+        answer: The secret entity.
+        search_enabled: Whether the player can use the search tool.
+    """
+
+    category: str
+    answer: str
+    search_enabled: bool
+
+
+LLM_TIMEOUT = 120.0
+
+console = Console()
+
+
+@agl.rollout
+async def q20_agent(task: Q20Task, llm: agl.LLM, rollout: agl.Rollout) -> None:
+    """Rollout function for the 20 Questions agent during training.
+
+    Args:
+        task: The 20 Questions task containing category, answer, and search settings.
+        llm: The LLM being trained (player model).
+        rollout: Rollout metadata from Agent-lightning.
+    """
+    answer_llm_setting = os.getenv("ANSWERER_LLM", "gpt-5-mini")
+    search_llm_setting = os.getenv("SEARCH_LLM", "gpt-4.1")
+    player_llm = CrewLLM(model="openai/" + llm.model, base_url=llm.endpoint, api_key="dummy", timeout=LLM_TIMEOUT)
+    answer_llm = CrewLLM(
+        model="openai/" + answer_llm_setting,
+        base_url=os.getenv("OPENAI_BASE_URL"),
+        api_key=os.getenv("OPENAI_API_KEY"),
+        reasoning_effort="low",
+        response_format=AnswererResponse,
+        timeout=LLM_TIMEOUT,
+    )
+    if task["search_enabled"]:
+        search_tool = SearchTool(
+            model=CrewLLM(
+                model="openai/" + search_llm_setting,
+                base_url=os.getenv("OPENAI_BASE_URL"),
+                api_key=os.getenv("OPENAI_API_KEY"),
+                reasoning_effort="none",
+                timeout=LLM_TIMEOUT,
+            )
+        )
+    else:
+        search_tool = None
+
+    flow = TwentyQuestionsFlow(player_llm=player_llm, answer_llm=answer_llm, search_tool=search_tool)
+    try:
+        await flow.kickoff_async(cast(Any, task))
+        agl.emit_reward(1.0 if flow.state.correct else 0.0)
+    except Exception:
+        console.print(f"Error in q20_agent: {traceback.format_exc()}")
+        raise
+        # Above, the exception is re-raised, so the rollout will appear failed, but reward will be none.
+        # The handling below is another approach that will make the rollout appear succeeded, but with 0 reward.
+        # I think algorithm should handle the case instead.
+        # agl.emit_exception(e)
+        # agl.emit_reward(0.0)
+
+
+def dry_run(model: Literal["qwen4b", "qwen30b"]):
+    """Run a quick dry-run test of the 20 Questions training setup.
+
+    Uses in-memory store and processes 4 sample tasks to verify the setup works.
+    """
+    store = agl.LightningStoreThreaded(agl.InMemoryLightningStore())
+    if model == "qwen4b":
+        model_name = "Qwen/Qwen3-4B-Instruct-2507"
+    elif model == "qwen30b":
+        model_name = "Qwen/Qwen3-30B-A3B-Instruct-2507"
+    else:
+        raise ValueError(f"Invalid model: {model}")
+    llm_proxy = create_llm_proxy(model_name, "qwen3_instruct", store=store)
+    trainer = agl.Trainer(
+        n_runners=2,
+        initial_resources={"llm": llm_proxy.as_resource()},
+        store=store,
+    )
+    try:
+        asyncio.run(llm_proxy.start())
+        sampled_csv = pd.read_csv("q20_nouns.csv").sample(n=4, random_state=42)  # type: ignore
+        sampled_csv["search_enabled"] = False
+        dataset = sampled_csv.to_dict(orient="records")  # type: ignore
+        trainer.dev(q20_agent, cast(agl.Dataset[Q20Task], dataset))
+    finally:
+        asyncio.run(llm_proxy.stop())
+
+
+async def algo(search: bool, model: Literal["qwen4b", "qwen30b"], port: int, ci: bool = False):
+    """Run the training algorithm for 20 Questions.
+
+    Args:
+        search: Whether to enable the search tool for the player.
+        model: Model variant to use ("qwen4b" or "qwen30b").
+        port: Port where the Agent-lightning store is running.
+    """
+    raw_data = pd.read_csv("q20_nouns.csv")  # type: ignore
+    raw_data["search_enabled"] = search
+    train_data, test_data = raw_data[raw_data["split"] == "train"], raw_data[raw_data["split"] == "test"]  # type: ignore
+
+    train_dataset = cast(agl.Dataset[Q20Task], train_data.to_dict(orient="records"))  # type: ignore
+    test_dataset = cast(agl.Dataset[Q20Task], test_data.to_dict(orient="records"))  # type: ignore
+
+    if model == "qwen4b":
+        model_name = "Qwen/Qwen3-4B-Instruct-2507"
+        renderer_name = "qwen3_instruct"
+    elif model == "qwen30b":
+        model_name = "Qwen/Qwen3-30B-A3B-Instruct-2507"
+        renderer_name = "qwen3_instruct"
+    else:
+        raise ValueError(f"Invalid model: {model}")
+
+    experiment_name = f"q20_{'search' if search else 'no_search'}_{model}"
+
+    llm_proxy_port = _find_available_port()
+
+    if ci:
+        train_dataset = cast(agl.Dataset[Q20Task], train_dataset[:2])  # type: ignore
+        test_dataset = cast(agl.Dataset[Q20Task], test_dataset[:2])  # type: ignore
+        group_size = 2
+        batch_size = 2
+        n_epochs = 1
+    else:
+        group_size = 16
+        batch_size = 16
+        n_epochs = 10
+
+    config = Config(
+        learning_rate=1e-4,
+        dataset_builder=AGLDatasetBuilder(
+            train_dataset=train_dataset,
+            val_dataset=test_dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            group_size=group_size,
+            seed=17,
+            n_epochs=n_epochs,
+        ),
+        lora_rank=16,
+        renderer_name=renderer_name,
+        model_name=model_name,
+        log_path=f"logs/{experiment_name}",
+        concurrency=32,
+        eval_every=4,
+        wandb_project="AgentLightningQ20",
+        wandb_name=experiment_name,
+        store_address=f"http://localhost:{port}",
+        llm_proxy_port=llm_proxy_port,
+        adapter_from_llm_proxy=False,
+        llm_proxy_retry_attempts=5,
+    )
+    await entrypoint(config)
+
+
+def algo_verl(search: bool, model: Literal["qwen25", "qwen3"], port: int):
+    """Alternatively, you can use VERL to train the 20 Questions agent locally.
+    Use this as a substitute for the `algo` function when Tinker service is not available.
+
+    Args:
+        search: Whether to enable the search tool for the player.
+        model: Specifies the model variant ('qwen25' or 'qwen3').
+        port: Port where the Agent-lightning store is running.
+    """
+    store = agl.LightningStoreClient(f"http://localhost:{port}")
+
+    if model == "qwen25":
+        model_name = "Qwen/Qwen2.5-3B-Instruct"
+    elif model == "qwen3":
+        model_name = "Qwen/Qwen3-4B-Instruct-2507"
+    else:
+        raise ValueError(f"Invalid model: {model}")
+
+    experiment_name = f"q20_{'search' if search else 'no_search'}_{model}"
+
+    verl_config = {
+        "algorithm": {
+            "adv_estimator": "grpo",
+            "use_kl_in_reward": False,
+        },
+        "data": {
+            "train_batch_size": 16,
+            "max_prompt_length": 8192,
+            "max_response_length": 1024,
+        },
+        "actor_rollout_ref": {
+            "rollout": {
+                "tensor_model_parallel_size": 1,
+                "n": 8,
+                "log_prob_micro_batch_size_per_gpu": 4,
+                "multi_turn": {"format": "hermes"},
+                "name": "vllm",
+                "gpu_memory_utilization": 0.8,
+            },
+            "actor": {
+                "ppo_mini_batch_size": 16,
+                "ppo_micro_batch_size_per_gpu": 2,
+                "optim": {"lr": 5e-7},
+                "use_kl_loss": False,
+                "kl_loss_coef": 0.0,
+                "entropy_coeff": 0,
+                "clip_ratio_low": 0.2,
+                "clip_ratio_high": 0.3,
+                "fsdp_config": {
+                    "param_offload": True,
+                    "optimizer_offload": True,
+                },
+            },
+            "ref": {
+                "log_prob_micro_batch_size_per_gpu": 4,
+                "fsdp_config": {"param_offload": True},
+            },
+            "model": {
+                "path": model_name,
+                "use_remove_padding": True,
+                "enable_gradient_checkpointing": True,
+                "enable_activation_offload": True,
+            },
+        },
+        "trainer": {
+            "n_gpus_per_node": 1,
+            "val_before_train": True,
+            "critic_warmup": 0,
+            "logger": ["console", "wandb"],
+            "project_name": "AgentLightningQ20VERL",
+            "experiment_name": experiment_name,
+            "nnodes": 1,
+            "test_freq": 4,
+            "total_epochs": 10,
+        },
+    }
+
+    verl = agl.VERL(verl_config)
+    # Use the data recorded at the proxy side
+    adapter = agl.LlmProxyTraceToTriplet()
+    verl.set_adapter(adapter)
+    verl.set_store(store)
+
+    raw_data = pd.read_csv("q20_nouns.csv")  # type: ignore
+    raw_data["search_enabled"] = search
+    train_data, test_data = raw_data[raw_data["split"] == "train"], raw_data[raw_data["split"] == "test"]  # type: ignore
+
+    train_dataset = cast(agl.Dataset[Q20Task], train_data.to_dict(orient="records"))  # type: ignore
+    test_dataset = cast(agl.Dataset[Q20Task], test_data.to_dict(orient="records"))  # type: ignore
+
+    verl.run(train_dataset=train_dataset, val_dataset=test_dataset)
+
+
+def runner(port: int = 4747, n_runners: int = 2):
+    """Run rollout runners that execute the 20 Questions game.
+
+    Args:
+        port: Port where the Agent-lightning store is running.
+        n_runners: Number of parallel runners to spawn.
+    """
+    # Run only the runners without algorithm
+    store = agl.LightningStoreClient(f"http://localhost:{port}")
+    trainer = agl.Trainer(
+        algorithm=None,
+        store=store,
+        strategy={"type": "cs", "managed_store": False, "n_runners": n_runners, "role": "runner"},
+    )
+    trainer.fit(q20_agent)
+
+
+def _run_dryrun(args: argparse.Namespace) -> None:
+    dry_run(model=args.model)
+
+
+def _run_algo(args: argparse.Namespace) -> None:
+    asyncio.run(algo(search=args.search, model=args.model, port=args.port, ci=args.ci))
+
+
+def _run_runner(args: argparse.Namespace) -> None:
+    runner(port=args.port, n_runners=args.n_runners)
+
+
+def _run_algo_verl(args: argparse.Namespace) -> None:
+    algo_verl(search=args.search, model=args.model, port=args.port)
+
+
+def main() -> None:
+    """Entry point for the 20 Questions training script."""
+    parser = argparse.ArgumentParser(description="Run the Q20 AgentLightning experiments.")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    dryrun_parser = subparsers.add_parser("dryrun", help="Run the in-memory dry run.")
+    dryrun_parser.add_argument(
+        "--model", choices=("qwen4b", "qwen30b"), default="qwen30b", help="Model variant to train."
+    )
+    dryrun_parser.set_defaults(func=_run_dryrun)
+
+    algo_parser = subparsers.add_parser("algo", help="Launch the full training algorithm.")
+    algo_parser.add_argument("--port", type=int, default=4747, help="Port for the AgentLightning store.")
+    algo_parser.add_argument("--search", action="store_true", help="Enable search tool.")
+    algo_parser.add_argument(
+        "--model",
+        choices=("qwen4b", "qwen30b"),
+        default="qwen30b",
+        help="Model variant to train.",
+    )
+    algo_parser.add_argument("--ci", action="store_true", help="Run in CI mode (smaller dataset, smaller batch).")
+    algo_parser.set_defaults(func=_run_algo)
+
+    algo_verl_parser = subparsers.add_parser("verl", help="Launch the full training algorithm with VERL.")
+    algo_verl_parser.add_argument("--port", type=int, default=4747, help="Port for the AgentLightning store.")
+    algo_verl_parser.add_argument(
+        "--model",
+        choices=("qwen25", "qwen3"),
+        default="qwen3",
+        help="Model variant to train.",
+    )
+    algo_verl_parser.add_argument("--search", action="store_true", help="Enable search tool.")
+    algo_verl_parser.set_defaults(func=_run_algo_verl)
+
+    runner_parser = subparsers.add_parser("runner", help="Run only the rollout runners.")
+    runner_parser.add_argument("--port", type=int, default=4747, help="Port for the AgentLightning store.")
+    runner_parser.add_argument("--n-runners", type=int, default=2, help="Number of runners to use.")
+    runner_parser.set_defaults(func=_run_runner)
+
+    args = parser.parse_args()
+    agl.setup_logging()
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

examples/tinker/tests/__init__.py

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

examples/tinker/tests/test_tinker_llm.py

@@ -0,0 +1,274 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""This file includes some basic tests for the integration of Tinker's sampling client
+with LiteLLM and Agent-lightning.
+
+It should be included in CI in future if we decided to maintain this example.
+"""
+
+import argparse
+import asyncio
+import json
+from typing import Any, Awaitable, Callable, Dict, cast
+
+import openai
+import tinker
+from agl_tinker.llm import TinkerLLM
+from agl_tinker.rollout import reconstruct_transitions
+from rich.console import Console
+from tinker_cookbook.renderers import Qwen3InstructRenderer
+from transformers import AutoTokenizer, PreTrainedTokenizer
+
+from agentlightning import (
+    AgentOpsTracer,
+    InMemoryLightningStore,
+    LLMProxy,
+    LlmProxyTraceToTriplet,
+    TracerTraceToTriplet,
+    emit_reward,
+    setup_logging,
+)
+from agentlightning.store import LightningStoreThreaded
+
+setup_logging(apply_to=["agl_tinker"])
+
+_tool_call_system_prompt = """
+You must call the provided tool once before responding to the user.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{"name": "echo_text", "description": "Echo back any provided text.", "parameters": {"type": "object", "properties": {"text": {"type": "string", "description": "Text to repeat back."}}, "required": ["text"]}}
+</tools>
+
+For each function call, return a json object with function name and args within <tool_call></tool_call> XML tags:
+<tool_call>
+{"name": <function-name>, "args": <args-json-object>}
+</tool_call>
+"""
+
+
+def _run_tool_call_roundtrip(client: openai.OpenAI, *, model_name: str) -> None:
+    """Force a tool call, parse the args, and feed back the tool result."""
+    prompt_messages: list[Dict[str, str]] = [
+        # FIXME: Currently the tool call definition needs to be hard-coded into the system prompt.
+        {"role": "system", "content": _tool_call_system_prompt},
+        {"role": "user", "content": "Use the tool to echo 'Agent Lightning loves tool calls'."},
+    ]
+    response = client.chat.completions.create(
+        model=model_name,
+        messages=cast(Any, prompt_messages),
+        max_tokens=256,
+        temperature=0.0,
+        # tools=cast(Any, tools),
+        # tool_choice="auto",
+    )
+    print("First response:", response)
+    tool_calls = response.choices[0].message.tool_calls or []
+    if not tool_calls:
+        raise AssertionError("Model did not emit a tool call when forced to do so.")
+    tool_call = tool_calls[0]
+    if tool_call.type != "function" or tool_call.function is None:  # pyright: ignore[reportUnnecessaryComparison]
+        raise AssertionError("Unexpected tool call payload from model.")
+    arguments = tool_call.function.arguments or "{}"
+    tool_args = json.loads(arguments)
+    tool_result = tool_args.get("text", "")
+    followup_messages: list[Dict[str, Any]] = [
+        *prompt_messages,
+        {
+            "role": "assistant",
+            "content": "",  # FIXME: Content must be here to make validation happy.
+            "tool_calls": [
+                {
+                    "id": tool_call.id,
+                    "type": "function",
+                    "function": {
+                        "name": tool_call.function.name,
+                        "arguments": tool_call.function.arguments,
+                    },
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "tool_call_id": tool_call.id,
+            "name": tool_call.function.name,
+            "content": f"Echoed text: {tool_result}",
+        },
+    ]
+    followup_response = client.chat.completions.create(
+        model=model_name,
+        messages=cast(Any, followup_messages),
+        max_tokens=64,
+        temperature=0.5,
+    )
+    print("Followup response:", followup_response)
+
+
+def _run_text_completion(client: openai.OpenAI, *, model_name: str) -> None:
+    """Simple text-only completion to contrast the tool-call scenario."""
+    response = client.chat.completions.create(
+        model=model_name,
+        messages=[{"role": "user", "content": "Hello world!"}],
+        max_tokens=20,
+        temperature=0.5,
+        top_p=0.9,
+        seed=11,
+    )
+    print(response)
+
+
+async def _run_tracer_test(*, use_tool_call: bool) -> None:
+    console = Console()
+    model_name = "Qwen/Qwen3-30B-A3B-Instruct-2507"
+
+    tokenizer = cast(PreTrainedTokenizer, AutoTokenizer.from_pretrained(model_name))  # type: ignore
+    renderer = Qwen3InstructRenderer(tokenizer)  # type: ignore
+    service_client = tinker.ServiceClient()
+    sampling_client = service_client.create_sampling_client(base_model=model_name)
+    tinker_llm = TinkerLLM(
+        model_name=model_name, renderer=renderer, tokenizer=tokenizer, sampling_client=sampling_client, max_tokens=20
+    )
+    tinker_llm.rewrite_litellm_custom_providers()
+
+    store = LightningStoreThreaded(InMemoryLightningStore())
+    rollout = await store.start_rollout("dummy", "train")
+    llm_proxy = LLMProxy(
+        port=4000,
+        store=store,
+        model_list=tinker_llm.as_model_list(),
+        num_retries=0,
+        launch_mode="thread",
+    )
+
+    scenario = "tool-call" if use_tool_call else "text-only"
+    console.print(f"Running tracer test scenario: {scenario}")
+
+    try:
+        tracer = AgentOpsTracer()
+        tracer.init()
+        tracer.init_worker(worker_id=0, store=store)
+
+        # init tracer before llm_proxy to avoid tracer provider being not active.
+        console.print("Starting LLM proxy...")
+        await llm_proxy.start()
+        console.print("LLM proxy started")
+
+        client = openai.OpenAI(base_url="http://localhost:4000/v1", api_key="dummy")
+
+        async with tracer.trace_context(
+            name=f"test_llm_{scenario}", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
+        ):
+            if use_tool_call:
+                _run_tool_call_roundtrip(client, model_name=model_name)
+            else:
+                _run_text_completion(client, model_name=model_name)
+            emit_reward(8.0)
+
+        print(f"Found {len(tracer.get_last_trace())} spans in the tracer")
+
+        tracer.teardown_worker(0)
+        tracer.teardown()
+
+        for store_span in await store.query_spans(rollout.rollout_id):
+            print(store_span)
+
+        spans = await store.query_spans(rollout.rollout_id)
+        console.print(f"Found {len(spans)} spans")
+        adapter = TracerTraceToTriplet()
+        trajectory = reconstruct_transitions(spans, adapter, rollout.rollout_id)
+        print(trajectory)
+        assert len(trajectory.transitions) > 0
+        assert len(trajectory.transitions[0].ac.tokens) > 0
+    finally:
+        console.print("Stopping LLM proxy...")
+        await llm_proxy.stop()
+        console.print("LLM proxy stopped")
+
+
+async def test_tracer_text_only():
+    await _run_tracer_test(use_tool_call=False)
+
+
+async def test_tracer_tool_call():
+    await _run_tracer_test(use_tool_call=True)
+
+
+async def test_tracer():
+    await test_tracer_tool_call()
+
+
+async def test_llm_proxy():
+    # FIXME: The llm proxy adapter needs some fixes to make this test work
+    console = Console()
+    model_name = "Qwen/Qwen3-30B-A3B-Instruct-2507"
+
+    tokenizer = cast(PreTrainedTokenizer, AutoTokenizer.from_pretrained(model_name))  # type: ignore
+    renderer = Qwen3InstructRenderer(tokenizer)  # type: ignore
+    service_client = tinker.ServiceClient()
+    sampling_client = service_client.create_sampling_client(base_model=model_name)
+    tinker_llm = TinkerLLM(
+        model_name=model_name, renderer=renderer, tokenizer=tokenizer, sampling_client=sampling_client, max_tokens=20
+    )
+    tinker_llm.rewrite_litellm_custom_providers()
+
+    store = LightningStoreThreaded(InMemoryLightningStore())
+    rollout = await store.start_rollout("dummy", "train")
+    llm_proxy = LLMProxy(
+        port=4000,
+        store=store,
+        model_list=tinker_llm.as_model_list(),
+        num_retries=0,
+        launch_mode="thread",
+    )
+
+    try:
+        # init tracer before llm_proxy to avoid tracer provider being not active.
+        console.print("Starting LLM proxy...")
+        await llm_proxy.start()
+        console.print("LLM proxy started")
+
+        client = openai.OpenAI(
+            base_url=f"http://localhost:4000/rollout/{rollout.rollout_id}/attempt/{rollout.attempt.attempt_id}",
+            api_key="dummy",
+        )
+
+        response = client.chat.completions.create(
+            model=model_name,
+            messages=[{"role": "user", "content": "Hello world!"}],
+            max_tokens=10,
+            temperature=0.5,
+            top_p=0.9,
+            seed=43,
+        )
+        print(response)
+
+        for store_span in await store.query_spans(rollout.rollout_id):
+            print(store_span)
+
+        spans = await store.query_spans(rollout.rollout_id)
+        console.print(f"Found {len(spans)} spans")
+        adapter = LlmProxyTraceToTriplet()
+        trajectory = reconstruct_transitions(spans, adapter, rollout.rollout_id)
+        print(trajectory)
+    finally:
+        console.print("Stopping LLM proxy...")
+        await llm_proxy.stop()
+        console.print("LLM proxy stopped")
+
+
+CLI_VARIANTS: Dict[str, Callable[[], Awaitable[None]]] = {
+    "tracer-tool": test_tracer_tool_call,
+    "tracer-text": test_tracer_text_only,
+    "llm-proxy": test_llm_proxy,
+}
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Manually run the async Tinker LLM integration tests.")
+    parser.add_argument(
+        "variant",
+        choices=sorted(CLI_VARIANTS.keys()),
+        help="Which async test to run.",
+    )
+    args = parser.parse_args()
+    asyncio.run(CLI_VARIANTS[args.variant]())  # type: ignore