Skip to content

Add AWS Bedrock Runtime (boto3) integration for Converse and InvokeModel instrumentation #279

@AbhiPrasad

Description

Summary

The AWS SDK for Python (boto3) provides the bedrock-runtime service client for invoking foundation models hosted on Amazon Bedrock. The Bedrock Runtime client exposes execution APIs via converse(), converse_stream(), invoke_model(), and invoke_model_with_response_stream(). This repository has zero direct instrumentation for any Bedrock Runtime SDK surface — no integration, no wrapper, no patcher, no auto-instrumentation config. Users who call the AWS Bedrock Runtime SDK directly via boto3 get no Braintrust spans.

The Converse API is AWS's unified interface for chat completions across all Bedrock-hosted models (Claude, Titan, Llama, Mistral, Cohere, etc.). It has its own request/response format distinct from any provider's native API. Users cannot use wrap_openai() or any other existing wrapper with boto3.client("bedrock-runtime").

This is the Python equivalent of braintrustdata/braintrust-sdk-javascript#1741.

What needs to be instrumented

The boto3 package exposes these execution surfaces via boto3.client("bedrock-runtime"), none of which are instrumented:

Converse API (highest priority)

SDK Method Description Streaming Return type
client.converse(...) Unified chat completions across all Bedrock models No dict
client.converse_stream(...) Streaming chat completions EventStream iterator dict with "stream" key

Response shape is well-defined: converse() returns a dict with output.message (role + content), usage (inputTokens, outputTokens, totalTokens), stopReason, and metrics (latencyMs). Standard span metrics extraction should be straightforward.

Streaming: converse_stream() returns a dict with a "stream" key containing a botocore EventStream. Events include contentBlockStart, contentBlockDelta, messageStop, and metadata (with usage and latency). The integration must accumulate deltas and finalize the span when the stream completes.

InvokeModel API (lower priority)

SDK Method Description Streaming Return type
client.invoke_model(...) Direct model invocation with provider-specific JSON payloads No dict with "body" as StreamingBody
client.invoke_model_with_response_stream(...) Streaming direct model invocation EventStream of byte chunks dict with "body" as EventStream

Request/response bodies are opaque: The body field is provider-specific JSON (Anthropic Messages API format, Meta Llama format, Cohere format, etc.). Initial implementation should log raw parsed JSON. Provider-specific normalization (extracting messages, token counts) can be added as a follow-up.

Implementation notes

Dynamic client methods: Unlike OpenAI or Anthropic SDKs where classes like Completions exist in importable modules, boto3 generates client methods at runtime from botocore service models. There is no BedrockRuntime.converse class defined in any importable module — methods are created dynamically when boto3.client("bedrock-runtime") is called.

Patching strategy: Because methods are dynamically generated, the recommended approach is:

  1. Auto-instrument path: Patch botocore.client.ClientCreator.create_client so that when a bedrock-runtime client is created, its methods are automatically wrapped. The wrapper checks client._service_model.service_name == "bedrock-runtime" and skips non-Bedrock clients.
  2. Manual wrap path: Provide a wrap_bedrock(client) helper that wraps methods directly on a client instance.

This two-pronged approach matches the OpenAI pattern (auto_instrument() patches classes, wrap_openai(client) patches an instance).

converse() parameters relevant for span metadata: modelId, inferenceConfig (contains temperature, maxTokens, topP, stopSequences), toolConfig, guardrailConfig, system, additionalModelRequestFields.

No async in boto3: The standard boto3 library is synchronous. Async support via aioboto3 / aiobotocore could be added in a follow-up phase.

AWS auth: Bedrock requests use IAM SigV4 signing rather than API keys. VCR cassettes will need auth header sanitization but no API key filtering.

Proposed span shape

converse() / converse_stream()

Span field Content
input messages, system, toolConfig extracted from kwargs
output role, content, stopReason from the response
metadata provider: "aws_bedrock", model (from modelId), inference config params
metrics tokens, prompt_tokens, completion_tokens, time_to_first_token (streaming)

invoke_model() / invoke_model_with_response_stream()

Span field Content
input modelId, parsed JSON body
output Parsed JSON response body
metadata provider: "aws_bedrock", model (from modelId), api: "invoke_model"
metrics Provider-specific (best-effort extraction from known formats)

Phased rollout

Phase 1 (MVP)

  • converse() sync tracing with full span shape
  • converse_stream() sync tracing with stream aggregation
  • wrap_bedrock(client) manual helper
  • auto_instrument() integration via ClientCreator patching
  • VCR-backed tests

Phase 2

  • invoke_model() with raw JSON logging
  • invoke_model_with_response_stream() with stream aggregation
  • Provider-specific body normalization for common model families

Phase 3

  • Async support via aioboto3 / aiobotocore
  • Multimodal attachment handling for image inputs in Converse API

No coverage in any instrumentation layer

  • No integration directory (py/src/braintrust/integrations/boto3_bedrock/)
  • No wrapper function (e.g. wrap_bedrock())
  • No patcher in any existing integration
  • No nox test session (test_boto3_bedrock)
  • No version entry in py/src/braintrust/integrations/versioning.py
  • No mention in py/src/braintrust/integrations/__init__.py

A grep for boto3, bedrock, or botocore across py/src/braintrust/integrations/ returns zero matches. The only repo references to boto3/botocore are in py/src/braintrust/aws.py (AWS credential helpers) and py/src/braintrust/cli/install/api.py (CLI tooling).

Braintrust docs status

unclear — The Braintrust tracing guide lists "AWS Bedrock" as a supported AI provider, but this likely refers to indirect coverage through frameworks (LangChain, LiteLLM) or OpenAI-compatible endpoints. There is no direct boto3 Bedrock wrapper or setup documentation.

Upstream references

Local repo files inspected

  • py/src/braintrust/integrations/ — no boto3_bedrock/ directory exists on main
  • py/src/braintrust/wrappers/ — no Bedrock wrapper
  • py/noxfile.py — no test_boto3_bedrock session
  • py/src/braintrust/integrations/__init__.py — Bedrock not listed in integration registry
  • py/src/braintrust/integrations/versioning.py — no Bedrock version matrix
  • Full repo grep for "boto3", "bedrock", "botocore" in integrations — zero matches

Metadata

Metadata

Assignees

No one assigned

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions