vendor/rllm/examples/sdk/tutorial_quickstart.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "65baac58",
   "metadata": {},
   "source": [
    "# RLLM SDK: Make Any Agent Trainable With Almost No Adaptation\n",
    "\n",
    "This tutorial will walk you through how to make **any existing agent code** trainable with RL with minimal changes. Specifically, we'll work through how to turn an existing countdown agent, and make it trainable using RLLM SDK.\n",
    "\n",
    "**Key Steps:** \n",
    "- Replace OpenAI client with our SDK provided OpenAI client\n",
    "- Write a rollout function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a72b3eb8",
   "metadata": {},
   "source": [
    "## Testing The Agent\n",
    "\n",
    "Start the proxy (litellm) for testing. During training, the Trainer manages this automatically. The proxy logs all LLM calls to a SQLite database which allow retrieval exact token ids during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4143bd1",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from rllm.sdk.proxy.proxy_manager import ProxyManager\n",
    "import os\n",
    "\n",
    "# Setup\n",
    "DB_PATH = \"/tmp/rllm_demo.db\"\n",
    "MODEL = \"gpt-4o-mini\"\n",
    "\n",
    "openai_api_key = os.environ[\"OPENAI_API_KEY\"]\n",
    "# openai_api_key=\"abc\"\n",
    "\n",
    "# Clean up\n",
    "Path(DB_PATH).unlink(missing_ok=True)\n",
    "\n",
    "# Start proxy\n",
    "proxy_manager = ProxyManager(proxy_port=4000, admin_token=\"my-shared-secret\")\n",
    "config = {\n",
    "    \"model_list\": [\n",
    "        {\n",
    "            \"model_name\": MODEL,\n",
    "            \"litellm_params\": {\n",
    "                \"model\": MODEL,\n",
    "                \"api_key\": openai_api_key,\n",
    "            },\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "proxy_manager.start_proxy_subprocess(config=config, db_path=DB_PATH, project=\"demo\")\n",
    "proxy_url = proxy_manager.get_proxy_url(include_v1=True)\n",
    "\n",
    "print(f\"✓ Proxy started at {proxy_url}\")\n",
    "print(f\"✓ Database: {DB_PATH}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67f46e74",
   "metadata": {},
   "source": [
    "### Prepare the Countdown Dataset\n",
    "**The Countdown Task:**  \n",
    "Given a set of numbers and a target, find an arithmetic expression using those numbers to reach the target. Each number can be used at most once. For example: numbers `[30, 32, 76]` and target `78` → solution could be `76 + 32 - 30 = 78`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f972d18e",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!python rllm/examples/solver_judge/prepare_countdown_data.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abbcb7aa",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from rllm.data.dataset import DatasetRegistry\n",
    "\n",
    "train_dataset = DatasetRegistry.load_dataset(\"countdown\", \"train\")\n",
    "test_dataset = DatasetRegistry.load_dataset(\"countdown\", \"test\")\n",
    "\n",
    "train_dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27791bda",
   "metadata": {},
   "source": [
    "## Your Original Agent Code\n",
    "\n",
    "A typical agent using the standard OpenAI client. This agent follows a Solver Verifier workflow: the solver generate multiple solution attempts, then the verifier select the best one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4a1d567",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from openai import AsyncOpenAI\n",
    "import asyncio\n",
    "import re\n",
    "\n",
    "judge_prompt = \"\"\"You are an expert verifier. Given a countdown problem and multiple solution attempts, select a correct solution.\n",
    "Problem:\n",
    "{problem}\n",
    "Solutions to evaluate:\n",
    "{solutions}\n",
    "A correct solution must satisfy the following criteria:\n",
    "1. The solution uses only the given numbers.\n",
    "2. Each number is used exactly once.\n",
    "3. Only basic arithmetic operations (+, -, *, /) are used.\n",
    "4. The calculation results in the target number.\n",
    "5. The final answer is clearly marked within <answer>...</answer> tags.\n",
    "Output the index of your selected solution within <answer>...</answer> tags, e.g., <answer>1</answer> for the first solution, <answer>2</answer> for the second solution, etc. If multiple solutions are correct, output the index of the first correct solution.\"\"\"\n",
    "\n",
    "\n",
    "def parse_answer_in_xml(solution_str):\n",
    "    answer_matches = re.findall(r\"<answer>(.*?)</answer>\", solution_str, re.IGNORECASE | re.DOTALL)\n",
    "    if answer_matches:\n",
    "        return answer_matches[-1].strip()\n",
    "    else:\n",
    "        return \"No solution found.\"\n",
    "\n",
    "\n",
    "class CountdownAgent:\n",
    "    \"\"\"A simple math solving agent - ORIGINAL VERSION\"\"\"\n",
    "\n",
    "    def __init__(self, api_key: str, model: str = \"gpt-4o-mini\"):\n",
    "        # Standard OpenAI client\n",
    "        self.client = AsyncOpenAI(api_key=api_key)\n",
    "        self.model = model\n",
    "\n",
    "    async def solve(self, problem: str) -> str:\n",
    "        \"\"\"Solve a math problem.\"\"\"\n",
    "        response = await self.client.chat.completions.create(\n",
    "            model=self.model,\n",
    "            messages=[\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"{problem}. Output the final answer within <answer>...</answer>\",\n",
    "                }\n",
    "            ],\n",
    "            max_tokens=1000,\n",
    "        )\n",
    "        return parse_answer_in_xml(response.choices[0].message.content)\n",
    "\n",
    "    async def judge(self, problem, solutions: list[str]) -> str:\n",
    "        \"\"\"Judge multiple solutions to a problem.\"\"\"\n",
    "        formatted_solutions = \"\\n\".join([f\"Solution {i + 1}:\\n<answer>{sol}</answer>\\n\" for i, sol in enumerate(solutions)])\n",
    "        prompt = judge_prompt.format(problem=problem, solutions=formatted_solutions)\n",
    "\n",
    "        response = await self.client.chat.completions.create(\n",
    "            model=self.model,\n",
    "            messages=[{\"role\": \"user\", \"content\": prompt}],\n",
    "            max_tokens=1000,\n",
    "        )\n",
    "        print(\"Judge:\", response.choices[0].message.content)\n",
    "        return parse_answer_in_xml(response.choices[0].message.content)\n",
    "\n",
    "    async def run(self, problem: str, n_solutions: int = 2) -> str:\n",
    "        \"\"\"Generate multiple solutions and judge them.\"\"\"\n",
    "        solutions = await asyncio.gather(*[self.solve(problem) for _ in range(n_solutions)])\n",
    "        print(\"First solution attempt:\")\n",
    "        print(solutions[0])\n",
    "\n",
    "        selected_index = await self.judge(problem, solutions)\n",
    "        if selected_index:\n",
    "            try:\n",
    "                selected_index = int(selected_index)\n",
    "                selected_solution = solutions[selected_index - 1]\n",
    "            except:\n",
    "                selected_solution = \"\"\n",
    "        else:\n",
    "            selected_solution = \"\"\n",
    "\n",
    "        return selected_solution\n",
    "\n",
    "\n",
    "# Use it\n",
    "agent = CountdownAgent(api_key=openai_api_key, model=MODEL)\n",
    "await agent.run(train_dataset[0][\"question\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa87538c",
   "metadata": {},
   "source": [
    "## Make It Trainable\n",
    "\n",
    "**Change:** Import the SDK client instead of OpenAI client  \n",
    "\n",
    "**What's `session()`?** A lightweight primitive that tracks all LLM calls within its scope and injects metadata into each call. Every LLM calls inside `with session()` can be retrieved via `sess._uid`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0f8a166",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import re\n",
    "from rllm.sdk import get_chat_client_async, session\n",
    "\n",
    "\n",
    "class TrainableAgent:\n",
    "    \"\"\"A simple math solving agent - TRAINABLE VERSION\"\"\"\n",
    "\n",
    "    def __init__(self, api_key: str, model: str = \"gpt-4o-mini\"):\n",
    "        # Replace standard OpenAI client with SDK client\n",
    "        # self.client = AsyncOpenAI(api_key=api_key)\n",
    "        self.client = get_chat_client_async(api_key=api_key, base_url=proxy_url)\n",
    "        self.model = model\n",
    "\n",
    "    async def solve(self, problem: str) -> str:\n",
    "        \"\"\"Solve a math problem.\"\"\"\n",
    "        response = await self.client.chat.completions.create(\n",
    "            model=self.model,\n",
    "            messages=[\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"{problem}. Output the final answer within <answer>...</answer>\",\n",
    "                }\n",
    "            ],\n",
    "            max_tokens=1000,\n",
    "        )\n",
    "        return parse_answer_in_xml(response.choices[0].message.content)\n",
    "\n",
    "    async def judge(self, problem, solutions: list[str]) -> str:\n",
    "        \"\"\"Judge multiple solutions to a problem.\"\"\"\n",
    "        formatted_solutions = \"\\n\".join([f\"Solution {i + 1}:\\n<answer>{sol}</answer>\\n\" for i, sol in enumerate(solutions)])\n",
    "        prompt = judge_prompt.format(problem=problem, solutions=formatted_solutions)\n",
    "\n",
    "        response = await self.client.chat.completions.create(\n",
    "            model=self.model,\n",
    "            messages=[{\"role\": \"user\", \"content\": prompt}],\n",
    "            max_tokens=1000,\n",
    "        )\n",
    "        return parse_answer_in_xml(response.choices[0].message.content)\n",
    "\n",
    "    async def run(self, problem: str, n_solutions: int = 2) -> str:\n",
    "        \"\"\"Generate multiple solutions and judge them.\"\"\"\n",
    "        solutions = await asyncio.gather(*[self.solve(problem) for _ in range(n_solutions)])\n",
    "\n",
    "        selected_index = await self.judge(problem, solutions)\n",
    "        if selected_index:\n",
    "            try:\n",
    "                selected_index = int(selected_index)\n",
    "                selected_solution = solutions[selected_index - 1]\n",
    "            except:\n",
    "                selected_solution = \"\"\n",
    "        else:\n",
    "            selected_solution = \"\"\n",
    "\n",
    "        return selected_solution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9efb951b",
   "metadata": {},
   "source": [
    "### Automatic LLM Call Tracking\n",
    "\n",
    "The `session()` primitive enables automatically capturing every LLM interaction.\n",
    "\n",
    "Access auto tracked LLM calls via `sess.steps`, each step is a concise view of a single LLM call (trace) with reward:\n",
    "\n",
    "Step obj have fields:\n",
    "- id: Trace ID, unique per trace, can be used to retrieve the full trace from the store\n",
    "- input: LLM input (from trace)\n",
    "- output: LLM response (from trace)\n",
    "- action: Parsed action (set manually by user)\n",
    "- reward: Step reward\n",
    "- metadata: Additional tracking data (can include model, tokens, latency, etc.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b730cdca",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Use it\n",
    "agent = TrainableAgent(api_key=openai_api_key, model=MODEL)\n",
    "with session() as sess:\n",
    "    solution = await agent.run(train_dataset[0][\"question\"])\n",
    "    print(\"selected solution: \", solution)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb518b07",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Access traces directly from the session\n",
    "steps = sess.steps\n",
    "print(f\"✅ Tracked {len(steps)} steps\\n\")\n",
    "\n",
    "# Inspect the first trace\n",
    "step = steps[2]\n",
    "\n",
    "print(\"=\" * 70 + \"\\nTRACE DETAILS\\n\" + \"=\" * 70)\n",
    "print(f\"\\nInput Messages:\")\n",
    "for msg in step.input[\"messages\"]:\n",
    "    print(f\"  [{msg['role']}]: {msg['content']}\")\n",
    "print(f\"\\nOutput:  {step.output['choices'][0]['message']['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d76020b3",
   "metadata": {},
   "source": [
    "## Evaluate the Agent\n",
    "\n",
    "We need to define a reward function that scores agent outputs for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efa5ae19",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def validate_equation(equation_str, available_numbers):\n",
    "    \"\"\"Validate that equation only uses available numbers and each number once.\"\"\"\n",
    "    try:\n",
    "        # Extract all numbers from the equation\n",
    "        numbers_in_eq = [int(n) for n in re.findall(r\"\\d+\", equation_str)]\n",
    "\n",
    "        # Check if all numbers in equation are available\n",
    "        available_numbers = sorted(available_numbers)\n",
    "        numbers_in_eq = sorted(numbers_in_eq)\n",
    "\n",
    "        # Each number should be used exactly once\n",
    "        return numbers_in_eq == available_numbers\n",
    "    except Exception:\n",
    "        return False\n",
    "\n",
    "\n",
    "def evaluate_equation(equation_str):\n",
    "    \"\"\"Safely evaluate the arithmetic equation using eval() with precautions.\"\"\"\n",
    "    try:\n",
    "        # Define a regex pattern that only allows numbers, operators, parentheses, and whitespace\n",
    "        allowed_pattern = r\"^[\\d+\\-*/().\\s]+$\"\n",
    "        if not re.match(allowed_pattern, equation_str):\n",
    "            raise ValueError(\"Invalid characters in equation.\")\n",
    "\n",
    "        # Evaluate the equation with restricted globals and locals\n",
    "        result = eval(equation_str, {\"__builtins__\": None}, {})\n",
    "        return result\n",
    "    except Exception:\n",
    "        return None\n",
    "\n",
    "\n",
    "def reward_fn(equation, numbers, target):\n",
    "    \"\"\"The scoring function for countdown task.\n",
    "\n",
    "    Args:\n",
    "        solution_str: the solution text\n",
    "        numbers: list of numbers\n",
    "        target: target number\n",
    "\n",
    "    Returns:\n",
    "        float: 1.0 if correct, 0.0 if incorrectet\n",
    "    \"\"\"\n",
    "\n",
    "    if equation is None or not validate_equation(equation, numbers):\n",
    "        return 0.0\n",
    "\n",
    "    # Evaluate equation\n",
    "    try:\n",
    "        result = evaluate_equation(equation)\n",
    "\n",
    "        if result is None:\n",
    "            return 0.0\n",
    "\n",
    "        if abs(result - target) < 1e-5:  # Account for floating point precision\n",
    "            return 1.0\n",
    "        else:\n",
    "            return 0.0\n",
    "    except Exception:\n",
    "        return 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "051d6dab",
   "metadata": {},
   "source": [
    "## Defining the Rollout Function\n",
    "\n",
    "Rollout function take initial input (prompts, task description etc.) , run the agent, and return a final reward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ff9fa17",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "async def rollout_v1(\n",
    "    question: str,\n",
    "    ground_truth: str,\n",
    "    nums: list,\n",
    "    target: float,\n",
    "    model=\"Qwen/Qwen3-4B-Instruct-2507\",\n",
    "    **kwargs,\n",
    ") -> float:\n",
    "    # we need to provide an rollout function that return a reward\n",
    "    agent = TrainableAgent(api_key=openai_api_key, model=model)\n",
    "    equation = await agent.run(question)\n",
    "    print(\"Target:\", target, \"\\nEquation:\", equation)\n",
    "\n",
    "    reward = reward_fn(equation, nums, target)\n",
    "    return reward\n",
    "\n",
    "\n",
    "reward = await rollout_v1(**train_dataset[0], model=\"gpt-4o-mini\")\n",
    "print(\"reward =\", reward)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c03ff3c6",
   "metadata": {},
   "source": [
    "## Train!\n",
    "\n",
    "Directly plug in the rollout function you defined to the `AgentTrainer`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2342eb0",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Training\n",
    "from rllm.trainer import AgentTrainer\n",
    "from hydra import initialize_config_dir, compose\n",
    "import os\n",
    "\n",
    "with initialize_config_dir(config_dir=os.path.abspath(\".\"), version_base=None):\n",
    "    config = compose(config_name=\"tutorial_config\")\n",
    "\n",
    "trainer = AgentTrainer(\n",
    "    agent_run_func=rollout_v1,\n",
    "    config=config,\n",
    "    train_dataset=train_dataset,\n",
    "    val_dataset=test_dataset,\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7aecfcd",
   "metadata": {},
   "source": [
    "## Bonus: Using @trajectory Decorator for Step-Level Control\n",
    "\n",
    "The `@trajectory` decorator is almost **equivalent to `with session()`**. The following are equivalent:\n",
    "\n",
    "A TrajectoryView obj is a group of steps with reward (optional) and metadata:\n",
    "- name: Trajectory name\n",
    "- steps: List of StepViews (auto-generated from traces)\n",
    "- reward: Trajectory reward (set manually)\n",
    "- input: Function arguments (dict)\n",
    "- output: Function return value (Any)\n",
    "- metadata: Additional tracking data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d17d5ae",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# These two are equivalent\n",
    "# v1\n",
    "with session(agent_name=\"countdown\") as sess:\n",
    "    output = await rollout_v1(**train_dataset[0])\n",
    "\n",
    "steps = sess.steps\n",
    "\n",
    "\n",
    "# v2 - using trajectory decorator\n",
    "@trajectory(name=\"countdown\")\n",
    "def rollout_v1(): ...\n",
    "\n",
    "\n",
    "traj = await rollout_v1(**train_dataset[0])\n",
    "steps = traj.steps\n",
    "output = traj.result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fd16dab",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from rllm.sdk import trajectory, get_chat_client_async\n",
    "from rllm.sdk.protocol import TrajectoryView\n",
    "\n",
    "\n",
    "class TrainableAgentV2:\n",
    "    \"\"\"A simple math solving agent - TRAINABLE VERSION V2\"\"\"\n",
    "\n",
    "    def __init__(self, api_key: str, model: str = \"gpt-4o-mini\"):\n",
    "        # Replace standard OpenAI client with SDK client\n",
    "        # self.client = AsyncOpenAI(api_key=api_key)\n",
    "        self.client = get_chat_client_async(api_key=api_key, base_url=proxy_url)\n",
    "        self.model = model\n",
    "\n",
    "    @trajectory(name=\"solver\")\n",
    "    async def solve(self, problem: str) -> str:\n",
    "        \"\"\"Solve a math problem.\"\"\"\n",
    "        response = await self.client.chat.completions.create(\n",
    "            model=self.model,\n",
    "            messages=[{\"role\": \"user\", \"content\": f\"{problem}. Output the final answer within <answer>...</answer>\"}],\n",
    "            max_tokens=1000,\n",
    "        )\n",
    "        return parse_answer_in_xml(response.choices[0].message.content)\n",
    "\n",
    "    @trajectory(name=\"judge\")\n",
    "    async def judge(self, problem, solutions: list[str]) -> str:\n",
    "        \"\"\"Judge multiple solutions to a problem.\"\"\"\n",
    "        formatted_solutions = \"\\n\".join([f\"Solution {i + 1}:\\n<answer>{sol}</answer>\\n\" for i, sol in enumerate(solutions)])\n",
    "        prompt = judge_prompt.format(problem=problem, solutions=formatted_solutions)\n",
    "\n",
    "        response = await self.client.chat.completions.create(\n",
    "            model=self.model,\n",
    "            messages=[{\"role\": \"user\", \"content\": prompt}],\n",
    "            max_tokens=1000,\n",
    "        )\n",
    "        return parse_answer_in_xml(response.choices[0].message.content)\n",
    "\n",
    "    async def run(self, problem: str, nums, target, n_solutions: int = 2) -> str:\n",
    "        \"\"\"Generate multiple solutions and judge them.\"\"\"\n",
    "        solver_trajs = await asyncio.gather(*[self.solve(problem) for _ in range(n_solutions)])\n",
    "        solutions = [traj.result for traj in solver_trajs]\n",
    "\n",
    "        judge_traj = await self.judge(problem, solutions)\n",
    "        selected_index = judge_traj.result\n",
    "\n",
    "        if selected_index:\n",
    "            try:\n",
    "                selected_index = int(selected_index)\n",
    "                selected_solution = solutions[selected_index - 1]\n",
    "            except:\n",
    "                selected_solution = \"\"\n",
    "        else:\n",
    "            selected_solution = \"\"\n",
    "\n",
    "        # assign rewards\n",
    "        for traj in solver_trajs:\n",
    "            traj.steps[-1].reward = reward_fn(equation=traj.result, numbers=nums, target=target)\n",
    "\n",
    "        judge_traj.steps[-1].reward = reward_fn(equation=selected_solution, numbers=nums, target=target)\n",
    "        return solver_trajs + [judge_traj]\n",
    "\n",
    "\n",
    "async def rollout_v2(question: str, nums, target, model, **kwargs) -> list[TrajectoryView]:\n",
    "    agent = TrainableAgentV2(None, model=model)\n",
    "    trajs = await agent.run(question, nums, target)\n",
    "    return trajs\n",
    "\n",
    "\n",
    "await rollout_v2(**train_dataset[0], model=\"gpt-4o-mini\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2d1cf0b",
   "metadata": {},
   "source": [
    "## How Does It Work?\n",
    "\n",
    "Here's what happens under the hood:\n",
    "\n",
    "1. **Trace Collection:** The litellm proxy captures all LLM calls, and directly get the token ids from the vllm server. This avoid retokenization, which cause off-policyness and training instability.\n",
    "2. **Trace Storage:** The SQlite based database will store all the LLM calls with appropriate metadata for fast retrieval.\n",
    "\n",
    "The following code show conceptually how the training loop works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80544ebd",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# The AgentTrainer will wrap your rollout function with a unique `rollout_id`:\n",
    "\n",
    "with session(rollout_id=rollout_id):\n",
    "    rollout(**input)\n",
    "\n",
    "steps = store.retrieve(session._uid)\n",
    "trajectories = group(steps)  # This part based on metadata\n",
    "\n",
    "trainer.step(trajectories)  # fetch all the trajectories and update model weight using GRPO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "392d8cbb",
   "metadata": {},
   "source": [
    "## Design Details (For The Curious)\n",
    "\n",
    "**Why a proxy?**  \n",
    "Transparent LLM call interception without modifying agent code. Works with any inference provider besides `OpenAI`, e.g., `Anthropic`. So you don't need to change your agent implementation, directly replace the chat client suffice.\n",
    "\n",
    "**How does session tracking work?**  \n",
    "Uses Python's **contextvar** for automatic context propagation. `with session()` or `@trajectory` creates a context that automatically groups all LLM calls inside it. Thread-safe, zero manual tracking.\n",
    "\n",
    "**Why SQLite storage?**  \n",
    "Zero installation overhead with no live service dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fb27b941602401d91542211134fc71a",
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Cleanup\n",
    "proxy_manager.shutdown_proxy()\n",
    "print(\"✓ Proxy shutdown complete\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}