Skip to content

Add Vercel AI SDK for Python integration for agent, model, and tool tracing #633

Description

Summary

Add a first-class Braintrust integration for Vercel's AI SDK for Python (PyPI package ai). The SDK provides provider-agnostic model streaming, structured output, agent loops, tools, hooks, MCP, multimodal content, and direct AI Gateway support.

There is currently no ai integration, setup helper, or auto_instrument() wiring in this repository. Existing OpenAI/Anthropic integrations may observe some direct provider calls, but they do not capture the AI SDK's framework-level agent, loop-turn, tool, hook, and gateway execution hierarchy.

The JavaScript SDK already has mature Vercel AI SDK support through wrapAISDK and the AI SDK instrumentation plugin. The Python integration should aim for equivalent span semantics where the Python API exposes the same concepts.

Recommended integration surface

AI SDK for Python v0.4.0 introduced built-in, provider-neutral telemetry under ai.experimental_telemetry. It emits typed spans and supports vendor adapters via register(), making a Braintrust adapter preferable to patching provider internals:

AI SDK telemetry Braintrust behavior
RunSpanData (Agent.run) Parent task span for the agent run
LoopTurnSpanData Preserve turn grouping without obscuring the model/tool children
AiStreamSpanData (ai.stream) llm span with input/output, model/provider metadata, finish reason, response IDs, and token metrics
AiGenerateSpanData Generation span for image/video and other non-streaming generation
ToolExecutionSpanData Child tool span with tool name, arguments, result, and error
HookSpanData Trace hook suspension/resolution/cancellation while preserving tool linkage
CustomSpanData Preserve user-created span names and attributes

The adapter should map the upstream first_token event to time_to_first_token and map Usage fields to Braintrust metrics:

  • input_tokensprompt_tokens
  • output_tokenscompletion_tokens
  • total_tokenstokens
  • reasoning_tokenscompletion_reasoning_tokens
  • cache_read_tokensprompt_cached_tokens
  • cache_write_tokensprompt_cache_creation_tokens

Proposed scope

  • Add an integration under py/src/braintrust/integrations/ with an idempotent setup API (for example, setup_ai_sdk()) that registers/unregisters a Braintrust telemetry adapter.
  • Wire it into auto_instrument() when the ai package is installed.
  • Preserve the upstream trace_id / parent_id hierarchy so an agent run contains model calls and tool executions in execution order.
  • Normalize AI SDK messages, reasoning, tool calls/results, structured output, and provider metadata into Braintrust's canonical span shape.
  • Materialize inline FilePart image/document/audio/video data as Braintrust attachments while preserving remote URLs.
  • Propagate provider errors and preserve async context-manager, async-iterator, early-close, cancellation, and replay/durable-execution behavior.
  • Avoid duplicate model spans and token accounting when Braintrust's underlying OpenAI or Anthropic integrations are also enabled.
  • Use context.span_origin.instrumentation.name = "ai-sdk" to match the JavaScript SDK.
  • Add a pinned ai version matrix entry and dedicated nox session. The package requires Python 3.12+, and its telemetry API is experimental as of v0.4.0, so CI/version routing must account for both constraints.
  • Add real provider or AI Gateway cassette-backed coverage for streaming, agent tool loops, token/cache/reasoning metrics, structured output, multimodal attachments, errors, and auto_instrument() import ordering.

Upstream references

Current repository gap

  • No py/src/braintrust/integrations/ai_sdk/ integration
  • No setup_ai_sdk() or equivalent public API
  • No AI SDK registration in py/src/braintrust/auto.py
  • No ai provider matrix entry or dedicated nox session
  • No AI SDK integration tests or cassettes

Metadata

Metadata

Assignees

No one assigned

    Labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions