Implement core telemetry API

Author: anbuzin

pyproject.toml

@@ -38,6 +38,9 @@ dependencies = [
 anthropic = ["anthropic>=0.83.0"]
 mcp = ["mcp>=1.18.0"]
 openai = ["openai>=2.14.0"]
+otel = [
+    "opentelemetry-sdk>=1.30",
+]
 vercel = [
     "vercel>=0.5.9",
 ]
@@ -67,6 +70,8 @@ dev = [
     "ruff~=0.8.0",
     "async-solipsism>=0.9",
     "ty~=0.0.37",
+    "opentelemetry-sdk>=1.43.0",
+    "opentelemetry-exporter-otlp-proto-http>=1.43.0",
 ]
 examples = [
     "fastapi>=0.136.1",

src/ai/__init__.py

@@ -1,4 +1,4 @@
-from . import errors, models, providers, util
+from . import errors, models, providers, telemetry, util
 from .agents import (
     Agent,
     AgentTool,
@@ -197,6 +197,7 @@
     "resolve_hook",
     "stream",
     "system_message",
+    "telemetry",
     "text_part",
     "thinking",
     "tool",

src/ai/telemetry/__init__.py

@@ -0,0 +1,52 @@
+"""Telemetry: spans, adapters, and the ambient current span.
+
+See :mod:`.span` for the full story.
+"""
+
+from .span import (
+    FIRST_TOKEN,
+    HOOK_CANCELLED,
+    HOOK_PENDING,
+    HOOK_RESOLVED,
+    RESPONSE_COMPLETE,
+    AiGenerateSpanData,
+    AiStreamSpanData,
+    CustomSpanData,
+    HookSpanData,
+    LoopTurnSpanData,
+    RunSpanData,
+    Span,
+    SpanData,
+    SpanEvent,
+    ToolExecutionSpanData,
+    WrapSpanAdapter,
+    current,
+    register,
+    span,
+    unregister,
+    wrap_span,
+)
+
+__all__ = [
+    "FIRST_TOKEN",
+    "HOOK_CANCELLED",
+    "HOOK_PENDING",
+    "HOOK_RESOLVED",
+    "RESPONSE_COMPLETE",
+    "AiGenerateSpanData",
+    "AiStreamSpanData",
+    "CustomSpanData",
+    "HookSpanData",
+    "LoopTurnSpanData",
+    "RunSpanData",
+    "Span",
+    "SpanData",
+    "SpanEvent",
+    "ToolExecutionSpanData",
+    "WrapSpanAdapter",
+    "current",
+    "register",
+    "span",
+    "unregister",
+    "wrap_span",
+]

src/ai/telemetry/console.py

@@ -0,0 +1,128 @@
+"""Console adapter: prints spans to a terminal as they happen.
+
+Register it like any adapter::
+
+    from ai.telemetry import console
+    ai.telemetry.register(console.ConsoleAdapter())
+
+Prints one ``▸`` line when a span starts (so long runs are visible
+live) and, when a trace's root span ends, the whole tree with
+durations.
+"""
+
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING
+
+from .. import telemetry
+
+if TYPE_CHECKING:
+    from typing import TextIO
+
+
+def _short(text: str, limit: int = 60) -> str:
+    return text if len(text) <= limit else text[: limit - 1] + "…"
+
+
+def _label(sp: telemetry.Span) -> str:
+    match sp.data:
+        case telemetry.AiStreamSpanData() as d:
+            tokens = (
+                f"  in:{d.usage.input_tokens} out:{d.usage.output_tokens} tok"
+                if d.usage is not None
+                else ""
+            )
+            return f"chat {d.model}{tokens}"
+        case telemetry.AiGenerateSpanData() as d:
+            return f"generate {d.model}"
+        case telemetry.ToolExecutionSpanData() as d:
+            args = ", ".join(f"{k}={v!r}" for k, v in (d.args or {}).items())
+            return f"tool {d.tool_name}({_short(args)})"
+        case telemetry.HookSpanData() as d:
+            return f"hook {d.label} {d.hook_type} [{d.status}]"
+        case telemetry.RunSpanData() as d:
+            return f"run {d.agent} ({d.model})"
+        case telemetry.LoopTurnSpanData() as d:
+            return f"turn {d.index}"
+        case telemetry.CustomSpanData() as d:
+            attrs = ", ".join(f"{k}={v!r}" for k, v in d.attributes.items())
+            return sp.name + (f" ({_short(attrs)})" if attrs else "")
+        case _:
+            return sp.name
+
+
+def _line(sp: telemetry.Span) -> str:
+    end = sp.ended_at or sp.started_at
+    # A span's lifetime can extend past the response (tool dispatch
+    # while the stream is open); when the milestone is there, report
+    # the true response latency instead.
+    for ev in sp.span_events:
+        if ev.name == telemetry.RESPONSE_COMPLETE:
+            end = ev.time_ns
+            break
+    duration = (end - sp.started_at) / 1e9
+    replay = "↻ " if sp.replay else ""
+    error = (
+        f"  ✗ {type(sp.error).__name__}: {_short(str(sp.error))}"
+        if sp.error is not None
+        else ""
+    )
+    return f"{replay}{_label(sp)}  {duration:.2f}s{error}"
+
+
+class ConsoleAdapter:
+    """Print spans to ``out`` (default: stdout)."""
+
+    def __init__(self, *, out: TextIO | None = None) -> None:
+        self._out = out if out is not None else sys.stdout
+        self._depth: dict[str, int] = {}
+        self._ended: dict[str, list[telemetry.Span]] = {}
+
+    def on_span_start(self, span: telemetry.Span) -> None:
+        depth = 0
+        if span.parent_id is not None:
+            depth = self._depth.get(span.parent_id, -1) + 1
+        self._depth[span.id] = depth
+        replay = "↻ " if span.replay else ""
+        self._out.write(f"▸ {'  ' * depth}{replay}{_label(span)}\n")
+
+    def on_span_event(
+        self, span: telemetry.Span, event: telemetry.SpanEvent
+    ) -> None:
+        depth = self._depth.get(span.id, 0) + 1
+        offset_ms = (event.time_ns - span.started_at) / 1e6
+        attrs = ", ".join(f"{k}={v!r}" for k, v in event.attributes.items())
+        suffix = f" ({_short(attrs)})" if attrs else ""
+        self._out.write(
+            f"· {'  ' * depth}{event.name} +{offset_ms:.0f}ms{suffix}\n"
+        )
+
+    def on_span_end(self, span: telemetry.Span) -> None:
+        self._ended.setdefault(span.trace_id, []).append(span)
+        if span.parent_id is not None:
+            return
+
+        # Root ended: print the tree and forget the trace.
+        spans = self._ended.pop(span.trace_id)
+        for s in spans:
+            self._depth.pop(s.id, None)
+        children: dict[str | None, list[telemetry.Span]] = {}
+        for s in spans:
+            children.setdefault(s.parent_id, []).append(s)
+
+        lines = [f"trace {span.trace_id}"]
+
+        def render(s: telemetry.Span, prefix: str, kid_prefix: str) -> None:
+            lines.append(prefix + _line(s))
+            kids = sorted(children.get(s.id, []), key=lambda c: c.started_at)
+            for i, kid in enumerate(kids):
+                last = i == len(kids) - 1
+                render(
+                    kid,
+                    kid_prefix + ("└─ " if last else "├─ "),
+                    kid_prefix + ("   " if last else "│  "),
+                )
+
+        render(span, "", "")
+        self._out.write("\n".join(lines) + "\n")

src/ai/telemetry/otel.py

@@ -0,0 +1,169 @@
+"""OpenTelemetry adapter: forwards spans to an otel tracer.
+
+::
+
+    from ai.telemetry import otel
+    otel.install()  # uses the global TracerProvider
+
+Span names and attributes follow the ``gen_ai`` semantic conventions,
+so LLM-aware viewers (Phoenix, Braintrust, Langfuse, Datadog, ...)
+render them natively.
+
+The adapter also attaches each current-setting otel span to the otel
+context in the task that opened it, so raw otel spans a user creates
+inside a tool still parent correctly under ours — and our root spans
+nest under any raw otel span the caller already has open.  Spans
+opened with ``set_as_current=False`` are never attached: they don't
+parent concurrent work in our tree, so they must not do it in otel's.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+
+from .. import errors, telemetry
+
+try:
+    from opentelemetry import context as otel_context
+    from opentelemetry import trace as otel_trace
+except ModuleNotFoundError as exc:  # pragma: no cover
+    raise errors.InstallationError(
+        "could not import `opentelemetry`, which is required for the otel "
+        'telemetry adapter, you can install it with `pip install "ai[otel]"` '
+        'or `uv add "ai[otel]"`'
+    ) from exc
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
+    from ..types import messages as messages_
+
+
+def _messages_json(messages: list[messages_.Message]) -> str:
+    return "[" + ",".join(m.model_dump_json() for m in messages) + "]"
+
+
+def _name(sp: telemetry.Span) -> str:
+    match sp.data:
+        case telemetry.AiStreamSpanData() as d:
+            return f"chat {d.model}"
+        case telemetry.AiGenerateSpanData() as d:
+            return f"generate_content {d.model}"
+        case telemetry.ToolExecutionSpanData() as d:
+            return f"execute_tool {d.tool_name}"
+        case telemetry.RunSpanData() as d:
+            return f"invoke_agent {d.agent}"
+        case telemetry.LoopTurnSpanData() as d:
+            return f"loop_turn {d.index}"
+        case _:
+            return sp.name
+
+
+def _attributes(sp: telemetry.Span) -> dict[str, Any]:
+    attrs: dict[str, Any] = {}
+    if sp.replay:
+        attrs["ai.replay"] = True
+    match sp.data:
+        case telemetry.AiStreamSpanData() as d:
+            attrs["gen_ai.operation.name"] = "chat"
+            attrs["gen_ai.request.model"] = d.model
+            attrs["gen_ai.input.messages"] = _messages_json(d.messages)
+            if d.message is not None:
+                attrs["gen_ai.output.messages"] = _messages_json([d.message])
+            if d.usage is not None:
+                attrs["gen_ai.usage.input_tokens"] = d.usage.input_tokens
+                attrs["gen_ai.usage.output_tokens"] = d.usage.output_tokens
+        case telemetry.AiGenerateSpanData() as d:
+            attrs["gen_ai.operation.name"] = "generate_content"
+            attrs["gen_ai.request.model"] = d.model
+            if d.message is not None:
+                attrs["gen_ai.output.messages"] = _messages_json([d.message])
+        case telemetry.ToolExecutionSpanData() as d:
+            attrs["gen_ai.operation.name"] = "execute_tool"
+            attrs["gen_ai.tool.name"] = d.tool_name
+            attrs["gen_ai.tool.call.id"] = d.tool_call_id
+            if d.args is not None:
+                attrs["gen_ai.tool.call.arguments"] = json.dumps(
+                    d.args, default=str
+                )
+            if d.result is not None:
+                attrs["gen_ai.tool.call.result"] = json.dumps(
+                    d.result, default=str
+                )
+            if d.is_error:
+                attrs["ai.tool.is_error"] = True
+        case telemetry.RunSpanData() as d:
+            attrs["gen_ai.operation.name"] = "invoke_agent"
+            attrs["gen_ai.agent.name"] = d.agent
+            attrs["gen_ai.request.model"] = d.model
+        case telemetry.HookSpanData() as d:
+            attrs["ai.hook.label"] = d.label
+            attrs["ai.hook.type"] = d.hook_type
+            attrs["ai.hook.status"] = d.status
+        case telemetry.LoopTurnSpanData() as d:
+            attrs["ai.loop_turn.index"] = d.index
+        case telemetry.CustomSpanData() as d:
+            for key, value in d.attributes.items():
+                attrs[key] = (
+                    value
+                    if isinstance(value, str | bool | int | float)
+                    else repr(value)
+                )
+    return attrs
+
+
+def install(
+    *, tracer_provider: otel_trace.TracerProvider | None = None
+) -> telemetry.WrapSpanAdapter:
+    """Create the otel adapter, register it, and return it.
+
+    Uses the global tracer provider unless one is passed.
+    """
+    tracer = otel_trace.get_tracer("ai", tracer_provider=tracer_provider)
+
+    @telemetry.wrap_span
+    async def otel_spans(span: telemetry.Span) -> AsyncGenerator[None, Any]:
+        # Parenting is ambient: the current otel context holds the
+        # parent (ours attach below; a raw otel span the caller has
+        # open works the same way), and it mirrors the framework's
+        # parenting because only current-setting spans attach.
+        otel_span = tracer.start_span(_name(span), start_time=span.started_at)
+        # Mirror the framework's currentness: a set_as_current=False
+        # span doesn't parent concurrent work in our tree, so it must
+        # not become current in the otel context either.
+        token = (
+            otel_context.attach(otel_trace.set_span_in_context(otel_span))
+            if span.set_as_current
+            else None
+        )
+        try:
+            # Milestones resume the yield live and are forwarded right
+            # away, so backends that render in-progress spans show them
+            # as they happen; span end resumes with None.
+            while (ev := (yield)) is not None:
+                otel_span.add_event(
+                    ev.name,
+                    {
+                        k: v
+                        if isinstance(v, str | bool | int | float)
+                        else repr(v)
+                        for k, v in ev.attributes.items()
+                    },
+                    timestamp=ev.time_ns,
+                )
+        finally:
+            if token is not None:
+                otel_context.detach(token)
+            for key, value in _attributes(span).items():
+                otel_span.set_attribute(key, value)
+            if span.error is not None:
+                if isinstance(span.error, Exception):
+                    otel_span.record_exception(span.error)
+                otel_span.set_status(
+                    otel_trace.StatusCode.ERROR, str(span.error)
+                )
+            otel_span.end(end_time=span.ended_at)
+
+    telemetry.register(otel_spans)
+    return otel_spans