Skip to content

Commit 27f96bc

Browse files
committed
feat(bedrock): trace bedrock runtime calls
Add a Python Bedrock Runtime integration that auto-instruments boto3/botocore client creation. The public auto instrumentation API is: ```python import braintrust braintrust.auto_instrument(bedrock=True) ``` The integration requires botocore >= 1.34.116, the first supported version whose Bedrock Runtime service model exposes both Converse and ConverseStream. It traces Converse, ConverseStream, InvokeModel, and InvokeModelWithResponseStream with Bedrock-shaped LLM spans. Manual wrapping is also exposed for existing clients: ```python from braintrust.integrations.bedrock_runtime import wrap_bedrock client = wrap_bedrock(client) ```
1 parent 28fae42 commit 27f96bc

25 files changed

Lines changed: 1612 additions & 0 deletions

examples/README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ Unless noted otherwise, every example below uses `braintrust.auto_instrument()`.
6262
| `agno/` | Agno agents and teams, sync + async, streaming + non-streaming |
6363
| `anthropic/` | Sync and async Anthropic clients |
6464
| `autogen/` | AutoGen `AssistantAgent` backed by `OpenAIChatCompletionClient` |
65+
| `bedrock_runtime/` | boto3 Bedrock Runtime `converse()` call against Amazon Nova Lite |
6566
| `claude_agent_sdk/` | Claude Agent SDK subprocess query |
6667
| `cohere/` | Cohere `ClientV2` chat call |
6768
| `crewai/` | CrewAI `Agent` + `Task` + `Crew` end-to-end |

examples/bedrock_runtime/README.md

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
# AWS Bedrock + Braintrust
2+
3+
Calls `braintrust.auto_instrument()` to patch boto3's Bedrock Runtime client factory, then runs a single `client.converse(...)` call against Amazon Nova Lite. The trace captures a parent `answer_question` span plus a child `bedrock.converse` LLM span with input, output, metadata, and token metrics.
4+
5+
## Run
6+
7+
```bash
8+
export BRAINTRUST_API_KEY=...
9+
export AWS_BEARER_TOKEN_BEDROCK=...
10+
export AWS_REGION=us-east-1
11+
export AWS_DEFAULT_REGION=us-east-1
12+
13+
uv sync
14+
uv run python example.py
15+
```
16+
17+
If you use IAM credentials instead of a Bedrock API key, set the usual AWS variables instead:
18+
19+
```bash
20+
export AWS_ACCESS_KEY_ID=...
21+
export AWS_SECRET_ACCESS_KEY=...
22+
export AWS_SESSION_TOKEN=... # if using temporary credentials
23+
```
24+
25+
The default model is `us.amazon.nova-lite-v1:0`. Override it with `BRAINTRUST_BEDROCK_MODEL` if needed.
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
#!/usr/bin/env python
2+
"""boto3 Bedrock Runtime client traced via braintrust.auto_instrument()."""
3+
4+
import os
5+
6+
import braintrust
7+
8+
9+
braintrust.auto_instrument()
10+
braintrust.init_logger(project="example-bedrock")
11+
12+
import boto3 # noqa: E402
13+
14+
15+
MODEL = os.getenv("BRAINTRUST_BEDROCK_MODEL", "us.amazon.nova-lite-v1:0")
16+
REGION = os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION") or "us-east-1"
17+
18+
19+
@braintrust.traced
20+
def answer_question(question: str) -> str:
21+
client = boto3.client("bedrock-runtime", region_name=REGION)
22+
response = client.converse(
23+
modelId=MODEL,
24+
messages=[{"role": "user", "content": [{"text": question}]}],
25+
inferenceConfig={"maxTokens": 64, "temperature": 0},
26+
)
27+
message = response["output"]["message"]
28+
return "".join(block.get("text", "") for block in message.get("content", []))
29+
30+
31+
def main() -> None:
32+
answer = answer_question("What is the capital of Australia? Reply in one sentence.")
33+
print(answer)
34+
braintrust.flush()
35+
36+
37+
if __name__ == "__main__":
38+
main()
Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
[project]
2+
name = "braintrust-bedrock-example"
3+
version = "0.1.0"
4+
description = "boto3 Bedrock Runtime client traced with Braintrust"
5+
requires-python = ">=3.10"
6+
dependencies = [
7+
"braintrust",
8+
"boto3",
9+
]
10+
11+
[tool.uv.sources]
12+
braintrust = { path = "../../py", editable = true }

py/noxfile.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,19 @@ def test_cohere(session, version):
230230
_run_tests(session, f"{INTEGRATION_DIR}/cohere/test_cohere.py", version=version)
231231

232232

233+
BOTO3_VERSIONS = _get_matrix_versions("boto3")
234+
235+
236+
@nox.session()
237+
@nox.parametrize("version", BOTO3_VERSIONS, ids=BOTO3_VERSIONS)
238+
def test_bedrock_runtime(session, version):
239+
"""Test the boto3 Bedrock Runtime integration."""
240+
_install_test_deps(session)
241+
_install_matrix_dep(session, "boto3", version)
242+
_install_matrix_dep(session, "botocore", version)
243+
_run_tests(session, f"{INTEGRATION_DIR}/bedrock_runtime/test_bedrock_runtime.py", version=version)
244+
245+
233246
INSTRUCTOR_VERSIONS = _get_matrix_versions("instructor")
234247

235248

py/pyproject.toml

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -438,6 +438,14 @@ latest = "pytest==9.0.3"
438438
[tool.braintrust.matrix.braintrust-core]
439439
latest = "braintrust-core==0.0.59"
440440

441+
[tool.braintrust.matrix.boto3]
442+
latest = "boto3==1.43.29"
443+
"1.34.116" = "boto3==1.34.116"
444+
445+
[tool.braintrust.matrix.botocore]
446+
latest = "botocore==1.43.29"
447+
"1.34.116" = "botocore==1.34.116"
448+
441449
# ---------------------------------------------------------------------------
442450
# Vendor packages — optional third-party packages the SDK can work without.
443451
# Keys are matrix keys; values are Python import names. The noxfile uses this
@@ -461,6 +469,7 @@ agentscope = ["agentscope"]
461469
agno = ["agno"]
462470
autogen = ["autogen-agentchat"]
463471
anthropic = ["anthropic"]
472+
bedrock_runtime = ["boto3", "botocore"]
464473
cohere = ["cohere"]
465474
claude_agent_sdk = ["claude-agent-sdk"]
466475
crewai = ["crewai"]
@@ -488,6 +497,8 @@ anthropic = "anthropic"
488497
cohere = "cohere"
489498
autoevals = "autoevals"
490499
braintrust-core = "braintrust_core"
500+
boto3 = "boto3"
501+
botocore = "botocore"
491502
crewai = "crewai"
492503
dspy = "dspy"
493504
google-adk = "google.adk"

py/src/braintrust/auto.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
AgnoIntegration,
1414
AnthropicIntegration,
1515
AutoGenIntegration,
16+
BedrockRuntimeIntegration,
1617
ClaudeAgentSDKIntegration,
1718
CohereIntegration,
1819
CrewAIIntegration,
@@ -72,6 +73,7 @@ def auto_instrument(
7273
openai_agents: bool = True,
7374
cohere: bool = True,
7475
autogen: bool = True,
76+
bedrock: bool = True,
7577
crewai: bool = True,
7678
strands: bool = True,
7779
temporal: bool = True,
@@ -106,6 +108,7 @@ def auto_instrument(
106108
openai_agents: Enable OpenAI Agents SDK instrumentation (default: True)
107109
cohere: Enable Cohere instrumentation (default: True)
108110
autogen: Enable AutoGen instrumentation (default: True)
111+
bedrock: Enable boto3 Bedrock Runtime instrumentation (default: True)
109112
crewai: Enable CrewAI instrumentation (default: True)
110113
strands: Enable Strands Agents instrumentation (default: True)
111114
temporal: Enable Temporal instrumentation (default: True)
@@ -195,6 +198,8 @@ def auto_instrument(
195198
results["cohere"] = _instrument_integration(CohereIntegration)
196199
if autogen:
197200
results["autogen"] = _instrument_integration(AutoGenIntegration)
201+
if bedrock:
202+
results["bedrock_runtime"] = _instrument_integration(BedrockRuntimeIntegration)
198203
if crewai:
199204
results["crewai"] = _instrument_integration(CrewAIIntegration)
200205
if strands:

py/src/braintrust/conftest.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,12 @@ def get_vcr_config():
247247
"openai-api-key",
248248
"x-goog-api-key",
249249
"x-bt-auth-token",
250+
"x-amz-security-token",
251+
"x-amz-date",
252+
"x-amz-content-sha256",
253+
"amz-sdk-invocation-id",
254+
"amz-sdk-request",
255+
"x-amzn-bedrock-api-key",
250256
],
251257
}
252258

py/src/braintrust/integrations/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from .agno import AgnoIntegration
44
from .anthropic import AnthropicIntegration
55
from .autogen import AutoGenIntegration
6+
from .bedrock_runtime import BedrockRuntimeIntegration
67
from .claude_agent_sdk import ClaudeAgentSDKIntegration
78
from .cohere import CohereIntegration
89
from .crewai import CrewAIIntegration
@@ -29,6 +30,7 @@
2930
"AgnoIntegration",
3031
"AnthropicIntegration",
3132
"AutoGenIntegration",
33+
"BedrockRuntimeIntegration",
3234
"ClaudeAgentSDKIntegration",
3335
"CohereIntegration",
3436
"CrewAIIntegration",
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
"""Test auto_instrument for boto3 Bedrock Runtime."""
2+
3+
import os
4+
5+
6+
os.environ.setdefault("AWS_EC2_METADATA_DISABLED", "true")
7+
if not os.environ.get("AWS_PROFILE") and not os.environ.get("AWS_BEARER_TOKEN_BEDROCK"):
8+
os.environ.setdefault("AWS_ACCESS_KEY_ID", "testing")
9+
os.environ.setdefault("AWS_SECRET_ACCESS_KEY", "testing")
10+
os.environ.setdefault("AWS_SESSION_TOKEN", "testing")
11+
12+
import boto3
13+
from braintrust.auto import auto_instrument
14+
from braintrust.integrations.test_utils import autoinstrument_test_context
15+
16+
17+
MODEL = os.getenv("BRAINTRUST_BEDROCK_CONVERSE_MODEL", "us.amazon.nova-lite-v1:0")
18+
REGION = os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION") or "us-east-1"
19+
20+
results = auto_instrument()
21+
assert results.get("bedrock_runtime") is True
22+
23+
results2 = auto_instrument()
24+
assert results2.get("bedrock_runtime") is True
25+
26+
with autoinstrument_test_context("test_auto_bedrock_runtime", integration="bedrock_runtime") as memory_logger:
27+
client = boto3.client("bedrock-runtime", region_name=REGION)
28+
response = client.converse(
29+
modelId=MODEL,
30+
messages=[{"role": "user", "content": [{"text": "Say hello in one word."}]}],
31+
inferenceConfig={"maxTokens": 20, "temperature": 0},
32+
)
33+
assert response["output"]["message"]["role"] == "assistant"
34+
35+
spans = memory_logger.pop()
36+
assert len(spans) == 1, f"Expected 1 span, got {len(spans)}"
37+
span = spans[0]
38+
assert span["metadata"]["provider"] == "bedrock"
39+
assert span["metadata"]["model"] == MODEL
40+
assert span["metadata"]["endpoint"] == "converse"
41+
assert span["span_attributes"]["name"] == "bedrock.converse"
42+
43+
print("SUCCESS")

0 commit comments

Comments
 (0)