chore: import upstream snapshot with attribution

examples/minimal/README.md

@@ -0,0 +1,20 @@
+# Minimal Component Showcase
+
+[![minimal CI status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml)
+
+`examples/minimal` provides bite-sized programs that demonstrate how individual Agent-lightning building blocks behave in isolation.
+
+Each module have been documented with its own CLI usage in the module-level docstring. Use this directory as a reference when wiring the same pieces into a larger system.
+
+## What’s Included?
+
+| Component | Demonstrated In | Highlights |
+| --- | --- | --- |
+| LightningStore + OTLP ingestion | `write_traces.py` | Shows how `OtelTracer` and `AgentOpsTracer` open rollouts, emit spans, and optionally forward them to a remote store client. |
+| MultiMetrics backend | `write_metrics.py` | Emits counters/histograms through `ConsoleMetricsBackend` and `PrometheusMetricsBackend` simultaneously, exposing `/metrics` for scraping. |
+| LLM proxying | `llm_proxy.py` | Guards either OpenAI or a local vLLM deployment with `LLMProxy`, proving how requests are routed through `/rollout/<id>/attempt/<id>` namespaces and captured in the store. |
+| vLLM lifecycle | `vllm_server.py` | Minimal context manager that shells out to `vllm serve`, monitors readiness, and tears down the process safely. |
+
+All runtime instructions (CLI arguments, required environment variables, etc.) are embedded directly in each script’s top-level docstring so the source stays self-documenting.
+
+For full-fledged training workflows or multi-component experiments, browse the other subdirectories under `examples/`. This `minimal` folder deliberately keeps each demonstration focused on a single component so you can understand and test them independently.

examples/minimal/llm_proxy.py

@@ -0,0 +1,254 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Examples to serve an LLM proxy for a vLLM server or an OpenAI service.
+
+Usage: run one of the following commands to start a server.
+
+```bash
+python llm_proxy.py vllm Qwen/Qwen2.5-0.5B-Instruct
+```
+
+Use the following command to test the LLM proxy.
+
+```bash
+python llm_proxy.py test Qwen/Qwen2.5-0.5B-Instruct
+```
+
+You can also test the OpenAI Proxy path (`OPENAI_API_KEY` environment variable is required).
+
+```bash
+dotenv run python llm_proxy.py openai gpt-4.1-mini
+```
+"""
+
+import argparse
+import asyncio
+import os
+from typing import Sequence, no_type_check
+
+import aiohttp
+from portpicker import pick_unused_port
+from rich.console import Console
+from vllm_server import vllm_server
+
+import agentlightning as agl
+
+console = Console()
+
+
+async def serve_llm_proxy_with_vllm(model_name: str, store_port: int = 43887):
+    """Serve an LLM proxy for a vLLM server."""
+    # Create a store to store the traces
+    store = agl.InMemoryLightningStore()
+    store_server = agl.LightningStoreServer(store, "127.0.0.1", store_port)
+    await store_server.start()
+
+    # Create a vLLM server
+    vllm_port = pick_unused_port()
+    with vllm_server(model_name, vllm_port) as vllm_endpoint:
+        # Server is up.
+
+        # Create an LLM proxy to guard the vLLM server and catch the traces
+        llm_proxy = agl.LLMProxy(
+            port=43886,
+            model_list=[
+                {
+                    "model_name": model_name,
+                    "litellm_params": {
+                        "model": f"hosted_vllm/{model_name}",
+                        "api_base": vllm_endpoint,
+                    },
+                }
+            ],
+            store=store_server,
+        )
+
+        try:
+            await llm_proxy.start()
+
+            # Wait forever
+            await asyncio.sleep(float("inf"))
+
+        finally:
+            # Stop the LLM proxy and the store server
+            await llm_proxy.stop()
+            await store_server.stop()
+
+
+async def serve_llm_proxy_with_openai(model_name: str, store_port: int = 43887):
+    """Serve an LLM proxy for an OpenAI server."""
+    # Create a store to store the traces
+    store = agl.InMemoryLightningStore()
+    store_server = agl.LightningStoreServer(store, "127.0.0.1", store_port)
+    await store_server.start()
+
+    if not os.getenv("OPENAI_API_KEY"):
+        raise ValueError("OPENAI_API_KEY environment variable is not set")
+
+    # Create an LLM proxy to guard the OpenAI server and catch the traces
+    llm_proxy = agl.LLMProxy(
+        port=43886,
+        model_list=[
+            {
+                "model_name": model_name,
+                "litellm_params": {
+                    "model": "openai/" + model_name,
+                    # Must have OpenAI API key set in the environment variable
+                },
+            }
+        ],
+        store=store_server,
+        callbacks=["opentelemetry"],
+    )
+
+    try:
+        await llm_proxy.start()
+        # Wait forever
+        await asyncio.sleep(float("inf"))
+    finally:
+        # Stop the LLM proxy and the store server
+        await llm_proxy.stop()
+        await store_server.stop()
+
+
+async def test_llm_proxy(model_name: str, store_port: int = 43887):
+    """Test the LLM proxy by sending a request to the proxy and checking the response.
+
+    We do it via aiohttp here. This can also be done with OpenAI client.
+    """
+    # We first connect to the store server and start a rollout.
+    store = agl.LightningStoreClient(f"http://localhost:{store_port}")
+    rollout = await store.start_rollout(input={"origin": "test_llm_proxy"})
+
+    # The chat completion URL is simply /v1/chat/completions under the namespace of current rollout and attempt.
+    # This ensures the traces are properly put into the correct bucket.
+    chat_completion_url = (
+        f"http://localhost:43886/rollout/{rollout.rollout_id}/attempt/{rollout.attempt.attempt_id}/v1/chat/completions"
+    )
+
+    async with aiohttp.ClientSession() as session:
+        async with session.post(
+            chat_completion_url,
+            json={
+                "model": model_name,
+                "messages": [{"role": "user", "content": "Hello, what's your name?"}],
+            },
+        ) as response:
+            response_body = await response.json()
+            console.print("Response body:", response_body)
+            _verify_response_body(response_body, model_name)
+
+    spans = await store.query_spans(rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id)
+    for span in spans:
+        console.print("Span:", span)
+    _verify_span(spans)
+
+    await store.close()
+
+
+@no_type_check
+def _verify_response_body(response_body: dict, model_name: str):
+    """Expect Response body to be something like this:
+
+    ```python
+    {
+        'id': 'chatcmpl-996a90a8678e4ed0a0d2724df2c0bba5',
+        'created': 1763178218,
+        'model': 'hosted_vllm/Qwen/Qwen2.5-0.5B-Instruct',
+        'object': 'chat.completion',
+        'choices': [
+            {
+                'finish_reason': 'stop',
+                'index': 0,
+                'message': {
+                    'content': 'Hello! I am Qwen, an AI language model created by Alibaba Cloud. My name is Qwen, and I can assist you with
+    various tasks and provide information on a wide range of topics. How may I help you today?',
+                    'role': 'assistant'
+                },
+                'provider_specific_fields': {
+                    'stop_reason': None,
+                    'token_ids': [9707, 0, ...],
+                }
+            }
+        ],
+        'usage': {'completion_tokens': 48, 'prompt_tokens': 36, 'total_tokens': 84},
+        'prompt_token_ids': [151644, 8948, ...],
+    }
+    ```
+    """
+    if "qwen" in model_name.lower():
+        assert "qwen" in response_body["choices"][0]["message"]["content"].lower()
+        assert (
+            "provider_specific_fields" in response_body["choices"][0]
+        ), "provider_specific_fields not found in response body"
+        assert (
+            "token_ids" in response_body["choices"][0]["provider_specific_fields"]
+        ), "token_ids not found in response body"
+        assert "prompt_token_ids" in response_body, "prompt_token_ids not found in response body"
+    else:
+        assert "chatgpt" in response_body["choices"][0]["message"]["content"].lower()
+
+
+def _verify_span(spans: Sequence[agl.Span]):
+    """Only a few spans are checked here.
+
+    `raw_gen_ai_request` span:
+
+    ```python
+    Span(
+        rollout_id='ro-4c68a7e686a1',
+        attempt_id='at-308eb814',
+        sequence_id=1,
+        name='raw_gen_ai_request',
+        attributes={
+            'llm.hosted_vllm.messages': '[{\'role\': \'user\', \'content\': "Hello, what\'s your name?"}]',
+            'llm.hosted_vllm.extra_body': "{'return_token_ids': True}",
+            'llm.hosted_vllm.choices': '... \'token_ids\': [40, 1079, 1207, 16948, ...',
+            'llm.hosted_vllm.model': 'Qwen/Qwen2.5-0.5B-Instruct',
+            'llm.hosted_vllm.prompt_token_ids': '[151644, 8948, ...]',
+        },
+        resource=OtelResource(
+            attributes={
+                'agentlightning.rollout_id': 'ro-4c68a7e686a1',
+                'agentlightning.attempt_id': 'at-308eb814',
+                'agentlightning.span_sequence_id': 1
+            },
+        )
+    )
+    ```
+    """
+
+    assert len(spans) > 1
+    has_raw_gen_ai_request = False
+    for span in spans:
+        if span.name == "raw_gen_ai_request":
+            has_raw_gen_ai_request = True
+            if "llm.hosted_vllm.messages" in span.attributes:
+                assert "return_token_ids" in span.attributes["llm.hosted_vllm.extra_body"]  # type: ignore
+                assert "token_ids" in span.attributes["llm.hosted_vllm.choices"]  # type: ignore
+                assert span.attributes["llm.hosted_vllm.prompt_token_ids"]
+        assert "agentlightning.rollout_id" in span.resource.attributes
+        assert "agentlightning.attempt_id" in span.resource.attributes
+        assert "agentlightning.span_sequence_id" in span.resource.attributes
+
+    assert has_raw_gen_ai_request, "raw_gen_ai_request span not found"
+
+
+if __name__ == "__main__":
+    agl.setup_logging()
+    parser = argparse.ArgumentParser(description="LLM Proxy runner")
+    parser.add_argument(
+        "mode",
+        choices=["vllm", "openai", "test"],
+        help="Which function to run",
+    )
+    parser.add_argument("model", type=str, help="Model name to serve.")
+
+    args = parser.parse_args()
+
+    if args.mode == "vllm":
+        asyncio.run(serve_llm_proxy_with_vllm(args.model))
+    elif args.mode == "openai":
+        asyncio.run(serve_llm_proxy_with_openai(args.model))
+    elif args.mode == "test":
+        asyncio.run(test_llm_proxy(args.model))

examples/minimal/vllm_server.py

@@ -0,0 +1,91 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Programmatically launch and stop an vLLM server."""
+
+import subprocess
+import time
+from contextlib import contextmanager
+from typing import Optional
+
+import httpx
+from openai import OpenAI
+from rich.console import Console
+
+console = Console()
+
+
+@contextmanager
+def vllm_server(
+    model_path: str,
+    port: int,
+    startup_timeout: float = 300.0,
+    terminate_timeout: float = 10.0,
+    gpu_memory_utilization: float = 0.7,
+    auto_tool_choice: bool = True,
+    tool_call_parser: Optional[str] = "hermes",
+):
+    """Serves a vLLM model from command line.
+
+    Args:
+        model_path: The path to the vLLM model. It can be either a local path or a Hugging Face model ID.
+        port: The port to serve the model on.
+        startup_timeout: The timeout for the server to start.
+        terminate_timeout: The timeout for the server to terminate.
+        gpu_memory_utilization: The GPU memory utilization for the server. Set it lower to avoid OOM.
+        auto_tool_choice: Whether to enable auto tool choice.
+        tool_call_parser: The tool call parser to use.
+    """
+    proc: Optional[subprocess.Popen[bytes]] = None
+    try:
+        vllm_serve_args = [
+            "--gpu-memory-utilization",
+            str(gpu_memory_utilization),
+            "--port",
+            str(port),
+        ]
+        if auto_tool_choice:
+            vllm_serve_args.append("--enable-auto-tool-choice")
+        if tool_call_parser is not None:
+            vllm_serve_args.append("--tool-call-parser")
+            vllm_serve_args.append(tool_call_parser)
+
+        proc = subprocess.Popen(["vllm", "serve", model_path, *vllm_serve_args])
+
+        # Wait for the server to be ready
+        url = f"http://localhost:{port}/health"
+        start = time.time()
+        client = httpx.Client()
+
+        while True:
+            try:
+                if client.get(url).status_code == 200:
+                    break
+            except Exception:
+                result = proc.poll()
+                if result is not None and result != 0:
+                    raise RuntimeError("Server exited unexpectedly.") from None
+                time.sleep(0.5)
+                if time.time() - start > startup_timeout:
+                    raise RuntimeError(f"Server failed to start in {startup_timeout} seconds.") from None
+
+        yield f"http://localhost:{port}/v1"
+    finally:
+        # Terminate the server
+        if proc is None:
+            return
+        proc.terminate()
+        try:
+            proc.wait(terminate_timeout)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+
+
+if __name__ == "__main__":
+    with vllm_server("Qwen/Qwen2.5-0.5B-Instruct", 8080) as endpoint:
+        client = OpenAI(base_url=endpoint, api_key="dummy")
+        response = client.chat.completions.create(
+            model="Qwen/Qwen2.5-0.5B-Instruct",
+            messages=[{"role": "user", "content": "Hello, what's your name?"}],
+        )
+        console.print(response)
+        assert "qwen" in response.choices[0].message.content.lower()  # type: ignore

examples/minimal/write_metrics.py

@@ -0,0 +1,98 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Demonstrate `MultiMetricsBackend` by emitting metrics to both console and Prometheus.
+
+Usage:
+    python write_metrics.py --duration 10 --prom-port 8000
+
+The script registers a counter and a histogram, pushes events through both the
+`ConsoleMetricsBackend` (for immediate feedback) and the `PrometheusMetricsBackend`
+for scraping via `/metrics`.
+
+Run a Prometheus server (for example via `docker/compose.prometheus-memory-store.yml`)
+and add the host running this script as a scrape target. By default the metrics
+endpoint binds to `0.0.0.0:9105`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import random
+import signal
+import sys
+import time
+from typing import Sequence
+
+from prometheus_client import start_http_server
+
+from agentlightning import setup_logging
+from agentlightning.utils.metrics import (
+    ConsoleMetricsBackend,
+    MetricsBackend,
+    MultiMetricsBackend,
+    PrometheusMetricsBackend,
+)
+
+
+def _register_metrics(backend: MetricsBackend) -> None:
+    backend.register_counter("minimal_requests_total", ["operation", "status"])
+    backend.register_histogram(
+        "minimal_latency_seconds",
+        ["operation"],
+        buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0],
+    )
+
+
+async def _emit_metrics(backend: MetricsBackend, duration: float, operations: Sequence[str]) -> None:
+    statuses = ["200", "404", "500"]
+    end_time = time.time() + duration
+    random.seed(1337)
+    while time.time() < end_time:
+        operation = random.choice(operations)
+        status = random.choices(statuses, weights=[0.9, 0.05, 0.05], k=1)[0]
+        latency = random.lognormvariate(-4.0, 0.5)
+        await backend.inc_counter("minimal_requests_total", labels={"operation": operation, "status": status})
+        await backend.observe_histogram("minimal_latency_seconds", value=latency, labels={"operation": operation})
+        await asyncio.sleep(0.25)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--duration", type=float, default=10.0, help="Seconds to emit metrics before shutting down.")
+    parser.add_argument("--prom-port", type=int, default=9105, help="Port for the /metrics endpoint.")
+    parser.add_argument("--prom-host", default="0.0.0.0", help="Host/IP for the /metrics endpoint.")
+    parser.add_argument("--group-level", type=int, default=2, help="ConsoleMetricsBackend label grouping depth.")
+    args = parser.parse_args()
+
+    setup_logging()
+
+    console_backend = ConsoleMetricsBackend(window_seconds=15.0, log_interval_seconds=2.0, group_level=args.group_level)
+    prom_backend = PrometheusMetricsBackend()
+    backend = MultiMetricsBackend([console_backend, prom_backend])
+    _register_metrics(backend)
+
+    start_http_server(args.prom_port, addr=args.prom_host)
+    print(f"Prometheus metrics exposed on http://{args.prom_host}:{args.prom_port}/metrics")
+    print(f"Emitting demo metrics for {args.duration:.1f}s ...")
+
+    # Handle CTRL+C gracefully
+    interrupted = False
+
+    def _handle_interrupt(signum: int, frame: object | None) -> None:  # pragma: no cover - signal handler
+        nonlocal interrupted
+        print(f"Received signal {signum}, stopping...")
+        interrupted = True
+
+    original_handler = signal.signal(signal.SIGINT, _handle_interrupt)
+    try:
+        asyncio.run(_emit_metrics(backend, duration=args.duration, operations=["search", "summary", "answer"]))
+    finally:
+        signal.signal(signal.SIGINT, original_handler)
+
+    if interrupted:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

examples/minimal/write_traces.py

@@ -0,0 +1,313 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example to write traces to a LightningStore via raw OpenTelemetry or AgentOpsTracer.
+
+The example can be run with or without using a Lightning Store server.
+When running this server, the traces will be written to the server via OTLP endpoint.
+
+Prior to running this example with `--use-client` flag, please start a LightningStore server with OTLP enabled first:
+
+```bash
+agl store --port 45993 --log-level DEBUG
+```
+
+The CLI also ships an `operation` mode showing how to record a synthetic operation span with
+[`operation`][agentlightning.operation], build link attributes via
+[`make_link_attributes`][agentlightning.utils.otel.make_link_attributes], tag the
+follow-up reward with [`make_tag_attributes`][agentlightning.utils.otel.make_tag_attributes],
+emit a reward span tied back to that operation, and then verify the recorded spans by
+extracting rewards, tags, and links from the store using `agentlightning.utils.otel` helpers.
+"""
+
+import argparse
+import asyncio
+import random
+import time
+from typing import Any, Dict, List, Sequence
+from uuid import uuid4
+
+from openai import AsyncOpenAI
+from rich.console import Console
+
+from agentlightning import AgentOpsTracer, LightningStoreClient, OtelTracer, Span, emit_reward, operation, setup_logging
+from agentlightning.semconv import AGL_OPERATION, LightningSpanAttributes
+from agentlightning.store import InMemoryLightningStore
+from agentlightning.utils.otel import (
+    extract_links_from_attributes,
+    extract_tags_from_attributes,
+    filter_and_unflatten_attributes,
+    get_tracer_provider,
+    make_link_attributes,
+    make_tag_attributes,
+    query_linked_spans,
+)
+
+console = Console()
+
+
+async def send_traces_via_otel(use_client: bool = False):
+    tracer = OtelTracer()
+    if not use_client:
+        store = InMemoryLightningStore()
+    else:
+        store = LightningStoreClient("http://localhost:45993")
+    rollout = await store.start_rollout(input={"origin": "write_traces_example"})
+
+    with tracer.lifespan(store):
+        # Initialize the capture of one single trace for one single rollout
+        async with tracer.trace_context(
+            "trace-manual", store=store, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
+        ) as tracer:
+            with tracer.start_as_current_span("grpc-span-1"):
+                time.sleep(0.01)
+
+                # Nested Span
+                with tracer.start_as_current_span("grpc-span-2"):
+                    time.sleep(0.01)
+
+            with tracer.start_as_current_span("grpc-span-3"):
+                time.sleep(0.01)
+
+            # This creates a reward span
+            emit_reward(1.0)
+
+    traces = await store.query_spans(rollout_id=rollout.rollout_id)
+    console.print(traces)
+
+    # Quickly validate the traces
+    assert len(traces) == 4
+    span_names = [span.name for span in traces]
+    assert "grpc-span-1" in span_names
+    assert "grpc-span-2" in span_names
+    assert "grpc-span-3" in span_names
+    assert "agentlightning.annotation" in span_names
+
+    last_span = traces[-1]
+    assert last_span.name == "agentlightning.annotation"
+    # NOTE: Try not to rely on this attribute like this example do. It may change in the future.
+    # Use utils from agentlightning.emitter to get the reward value.
+    assert last_span.attributes["agentlightning.reward.0.value"] == 1.0
+
+    if use_client:
+        # When using client, the resource should have rollout_id and attempt_id set
+        for span in traces:
+            assert "agentlightning.rollout_id" in span.resource.attributes
+            assert "agentlightning.attempt_id" in span.resource.attributes
+
+    if isinstance(store, LightningStoreClient):
+        await store.close()
+
+
+async def send_traces_via_agentops(use_client: bool = False):
+    tracer = AgentOpsTracer()
+    if not use_client:
+        store = InMemoryLightningStore()
+    else:
+        store = LightningStoreClient("http://localhost:45993")
+    rollout = await store.start_rollout(input={"origin": "write_traces_example"})
+
+    # Initialize the tracer lifespan
+    # One lifespan can contain multiple traces
+    with tracer.lifespan(store):
+        # Inspect current tracer provider
+        get_tracer_provider(inspect=True)
+
+        # Initialize the capture of one single trace for one single rollout
+        async with tracer.trace_context(
+            "trace-1", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
+        ):
+            openai_client = AsyncOpenAI()
+            response = await openai_client.chat.completions.create(
+                model="gpt-4.1-mini",
+                messages=[
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": "Hello, what's your name?"},
+                ],
+            )
+            console.print(response)
+            assert response.choices[0].message.content is not None
+            assert "chatgpt" in response.choices[0].message.content.lower()
+
+    traces = await store.query_spans(rollout_id=rollout.rollout_id)
+    console.print(traces)
+    await _verify_agentops_traces(traces, use_client=use_client)
+    if isinstance(store, LightningStoreClient):
+        await store.close()
+
+
+async def _verify_agentops_traces(spans: Sequence[Span], use_client: bool = False):
+    """Expected traces to something like:
+
+    ```python
+    Span(
+        rollout_id='ro-ef9ff8a429d1',
+        attempt_id='at-37cc5f24',
+        sequence_id=1,
+        trace_id='b3a16b603f7805934215d467e717c9e7',
+        span_id='2782d5d750f49b2d',
+        parent_id='2fb97c818363bce3',
+        name='openai.chat.completion',
+        status=TraceStatus(status_code='OK', description=None),
+        attributes={
+            'gen_ai.request.type': 'chat',
+            'gen_ai.system': 'OpenAI',
+            'gen_ai.request.model': 'gpt-4.1-mini',
+            'gen_ai.request.streaming': False,
+            'gen_ai.prompt.0.role': 'system',
+            'gen_ai.prompt.0.content': 'You are a helpful assistant.',
+            'gen_ai.prompt.1.role': 'user',
+            'gen_ai.prompt.1.content': "Hello, what's your name?",
+            'gen_ai.response.id': 'chatcmpl-Cc1osPWiArOwCS8nUkp0kZuZPkpY4',
+            'gen_ai.response.model': 'gpt-4.1-mini-2025-04-14',
+            'gen_ai.completion.0.role': 'assistant',
+            'gen_ai.completion.0.content': "Hello! I'm ChatGPT, your AI assistant. How can I help you today?",
+        },
+        resource=OtelResource(
+            attributes={
+                'agentops.project.id': 'temporary',
+                'agentlightning.rollout_id': 'ro-ef9ff8a429d1',
+                'agentlightning.attempt_id': 'at-37cc5f24'
+            },
+            schema_url=''
+        )
+    )
+    ```
+    """
+    assert len(spans) == 2
+    for span in spans:
+        if span.name == "openai.chat.completion":
+            assert span.attributes["gen_ai.request.model"] == "gpt-4.1-mini"
+            assert span.attributes["gen_ai.request.streaming"] == False
+            assert span.attributes["gen_ai.prompt.0.role"] == "system"
+            assert span.attributes["gen_ai.prompt.0.content"] == "You are a helpful assistant."
+            assert span.attributes["gen_ai.prompt.1.role"] == "user"
+            assert span.attributes["gen_ai.prompt.1.content"] == "Hello, what's your name?"
+            assert "chatgpt" in span.attributes["gen_ai.completion.0.content"].lower()  # type: ignore
+            if use_client:
+                assert "agentlightning.rollout_id" in span.resource.attributes
+                assert "agentlightning.attempt_id" in span.resource.attributes
+        else:
+            assert "trace-1" in span.name
+            assert span.attributes["agentops.span.kind"] == "session"
+
+
+async def send_operation_links(use_client: bool = False) -> None:
+    """Demonstrate operation spans wired to reward annotations and verify the stored spans."""
+
+    tracer = OtelTracer()
+    if not use_client:
+        store = InMemoryLightningStore()
+    else:
+        store = LightningStoreClient("http://localhost:45993")
+    conversation_id = "chat-42"
+    tags: Sequence[str] = ("demo.operation", "reward.positive")
+    reward_value = 0.9
+    operation_id = f"{conversation_id}-{uuid4().hex[:8]}"
+    rollout = await store.start_rollout(input={"origin": "write_traces_operation"})
+
+    with tracer.lifespan(store):
+        async with tracer.trace_context(
+            "operation-demo", store=store, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
+        ):
+            console.print(f"[operation] recording span conversation={conversation_id} operation_id={operation_id}")
+            with operation(conversation_id=conversation_id, operation_id=operation_id) as op_ctx:
+                op_ctx.set_input(
+                    task={"conversation_id": conversation_id},
+                    metadata={"operation_id": operation_id},
+                )
+                synthetic_payload = {
+                    "operation_id": operation_id,
+                    "status": "ok",
+                    "latency_seconds": round(random.uniform(0.05, 0.2), 3),
+                }
+                await asyncio.sleep(0.05)
+                op_ctx.set_output(synthetic_payload)
+
+            link_attrs = make_link_attributes({"conversation_id": conversation_id, "operation_id": operation_id})
+            tag_attrs = make_tag_attributes(list(tags))
+            emit_reward(
+                reward_value,
+                attributes={**link_attrs, **tag_attrs},
+            )
+
+    spans = await store.query_spans(rollout_id=rollout.rollout_id)
+    console.print(spans)
+    _verify_operation_spans(spans, conversation_id, operation_id, tags, reward_value)
+
+    if isinstance(store, LightningStoreClient):
+        await store.close()
+
+
+def _verify_operation_spans(
+    spans: Sequence[Span],
+    conversation_id: str,
+    operation_id: str,
+    tags: Sequence[str],
+    expected_reward: float,
+) -> None:
+    """Verify spans recorded by the operation demo using OTEL helpers."""
+
+    operation_spans = [span for span in spans if span.name == AGL_OPERATION]
+    if not operation_spans:
+        raise RuntimeError("No operation spans recorded.")
+    console.print(f"[verify] found {len(operation_spans)} operation spans")
+
+    reward_span: Span | None = None
+    reward_payload: List[Dict[str, Any]] = []
+    for span in spans:
+        flattened = dict(span.attributes or {})
+        reward_section = filter_and_unflatten_attributes(flattened, LightningSpanAttributes.REWARD.value)
+        if reward_section:
+            reward_span = span
+            if isinstance(reward_section, list):
+                reward_payload = [dict(item) for item in reward_section]  # type: ignore[arg-type]
+            else:
+                reward_payload = [dict(reward_section)]  # type: ignore[arg-type]
+            break
+
+    if reward_span is None or not reward_payload:
+        raise RuntimeError("No reward span recorded for operation demo.")
+
+    primary_reward = reward_payload[0].get("value")
+    console.print(f"[verify] reward dimensions: {reward_payload}")
+    if primary_reward != expected_reward:
+        raise AssertionError(f"Expected reward {expected_reward}, observed {primary_reward}")
+
+    reward_attributes = dict(reward_span.attributes or {})
+    extracted_tags = extract_tags_from_attributes(reward_attributes)
+    console.print(f"[verify] reward tags: {extracted_tags}")
+    for tag in tags:
+        if tag not in extracted_tags:
+            raise AssertionError(f"Missing tag '{tag}' on reward span")
+
+    link_models = extract_links_from_attributes(reward_attributes)
+    matches = query_linked_spans(operation_spans, link_models)
+    if not matches:
+        raise AssertionError("No operation span matched the reward links")
+    console.print(f"[verify] reward links resolved spans: {[span.span_id for span in matches]}")
+
+    linked_attrs = dict(matches[0].attributes or {})
+    if linked_attrs.get("conversation_id") != conversation_id or linked_attrs.get("operation_id") != operation_id:
+        raise AssertionError("Linked operation span attributes do not match expected identifiers")
+    console.print("[verify] linked operation span attributes validated")
+
+
+def main():
+    setup_logging("DEBUG")
+    parser = argparse.ArgumentParser()
+    parser.add_argument("mode", choices=["otel", "agentops", "operation"])
+    parser.add_argument("--use-client", action="store_true")
+    args = parser.parse_args()
+
+    if args.mode == "otel":
+        asyncio.run(send_traces_via_otel(use_client=args.use_client))
+    elif args.mode == "agentops":
+        asyncio.run(send_traces_via_agentops(use_client=args.use_client))
+    elif args.mode == "operation":
+        asyncio.run(send_operation_links(use_client=args.use_client))
+    else:
+        raise ValueError(f"Invalid mode: {args.mode}")
+
+
+if __name__ == "__main__":
+    main()