Skip to content

Commit 3fc87ab

Browse files
alexander-alderman-webbalexander-alderman-webb
committed
feat(bedrock): Add extension skeleton
1 parent d565fed commit 3fc87ab

7 files changed

Lines changed: 248 additions & 169 deletions

File tree

‎sentry_sdk/integrations/boto3/_client.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,7 @@
1414
_set_span_attributes,
1515
_start_client_span,
1616
)
17-
from sentry_sdk.integrations.boto3._services.registry import (
18-
_resolve_service,
19-
)
17+
from sentry_sdk.integrations.boto3._services.bedrock import _BedrockExtension
2018
from sentry_sdk.integrations.boto3.consts import IDENTIFIER
2119
from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan
2220
from sentry_sdk.utils import capture_internal_exceptions
@@ -104,7 +102,11 @@ def sentry_patched_make_api_call(
104102
with capture_internal_exceptions():
105103
# add optional metadata to the context, e.g. service-name, region-name, etc.
106104
ctx.add_metadata(self)
107-
service_ext = _resolve_service(ctx.service_name)
105+
106+
service_ext = None
107+
if ctx.service_name == "bedrock-runtime":
108+
service_ext = _BedrockExtension()
109+
108110
span = _start_client_span(ctx, service_ext)
109111

110112
if span is None:

‎sentry_sdk/integrations/boto3/_instrumentation.py‎

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -173,30 +173,28 @@ def _start_client_span(
173173
# https://opentelemetry.io/docs/specs/semconv/cloud-providers/aws-sdk/#aws-sdk-spans
174174
SPANDATA.SENTRY_KIND: "client",
175175
}
176-
with capture_internal_exceptions():
177-
attributes.update(_get_client_attributes(ctx))
176+
attributes.update(_get_client_attributes(ctx))
178177

179178
# `sentry.span_op` and `sentry.span_origin` are set to generic defaults;
180179
# a service extension can override them with `get_span_op()` and `get_span_origin()`.
181180
span_op = OP.HTTP_CLIENT
182181
span_origin = ORIGIN
183182

184183
if service_ext is not None:
185-
with capture_internal_exceptions():
186-
service_op = service_ext.get_span_op(ctx)
187-
if service_op is not None:
188-
span_op = service_op
184+
service_op = service_ext.get_span_op(ctx)
185+
if service_op is not None:
186+
span_op = service_op
189187

190-
with capture_internal_exceptions():
191-
service_origin = service_ext.get_span_origin(ctx)
192-
if service_origin is not None:
193-
span_origin = service_origin
188+
service_origin = service_ext.get_span_origin(ctx)
189+
if service_origin is not None:
190+
span_origin = service_origin
194191

195-
with capture_internal_exceptions():
196-
attributes.update(service_ext.get_request_attributes(ctx))
192+
attributes.update(service_ext.get_request_attributes(ctx))
197193

198194
if has_span_streaming_enabled(client.options):
199-
if sentry_sdk.traces.get_current_span() is None:
195+
if (
196+
service_ext is None or service_ext.skip_if_no_active_span
197+
) and sentry_sdk.traces.get_current_span() is None:
200198
return None
201199

202200
# `start_span()` evaluates `ignore_spans` against the initial attributes.

‎sentry_sdk/integrations/boto3/_services/base.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ class _ServiceExtension:
1616
"""
1717

1818
__slots__ = ()
19+
skip_if_no_active_span = True
1920

2021
def get_span_op(self, ctx: "AwsCallContext") -> "Optional[str]":
2122
"""Return an optional `sentry.op` override, or `None` to keep the default."""

‎sentry_sdk/integrations/boto3/_services/bedrock.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
from typing import TYPE_CHECKING
2+
3+
from sentry_sdk.consts import OP, SPANDATA
4+
from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
5+
6+
if TYPE_CHECKING:
7+
from typing import Any, Optional
8+
9+
from sentry_sdk._types import Attributes
10+
from sentry_sdk.integrations.boto3._context import AwsCallContext
11+
12+
13+
class _BedrockExtension(_ServiceExtension):
14+
__slots__ = ()
15+
skip_if_no_active_span = False
16+
17+
def get_span_op(self, ctx: "AwsCallContext") -> "Optional[str]":
18+
if ctx.operation_name == "Converse" and ctx.operation_name == "ConverseStream":
19+
return None
20+
return OP.GEN_AI_CHAT
21+
22+
def get_span_origin(self, ctx: "AwsCallContext") -> "Optional[str]":
23+
if ctx.operation_name != "Converse" and ctx.operation_name != "ConverseStream":
24+
return None
25+
return "auto.ai.boto3"
26+
27+
def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
28+
if ctx.operation_name != "Converse" and ctx.operation_name != "ConverseStream":
29+
return {}
30+
31+
return {
32+
SPANDATA.GEN_AI_PROVIDER_NAME: "aws.bedrock",
33+
SPANDATA.GEN_AI_OPERATION_NAME: "chat",
34+
SPANDATA.GEN_AI_RESPONSE_STREAMING: ctx.operation_name == "ConverseStream",
35+
}
36+
37+
def get_response_attributes(
38+
self,
39+
ctx: "AwsCallContext",
40+
response: "Any",
41+
) -> "Attributes":
42+
return {}

‎sentry_sdk/integrations/boto3/_services/registry.py‎

Lines changed: 0 additions & 22 deletions
This file was deleted.

‎tests/integrations/boto3/test_bedrock.py‎

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
import boto3
2+
import pytest
3+
from botocore.exceptions import ClientError
4+
from botocore.loaders import Loader
5+
from botocore.stub import Stubber
6+
7+
import sentry_sdk
8+
from sentry_sdk.consts import SPANDATA
9+
from sentry_sdk.integrations.boto3 import Boto3Integration
10+
11+
pytestmark = pytest.mark.skipif(
12+
"bedrock-runtime" not in Loader().list_available_services("service-2"),
13+
reason="botocore is too old for bedrock-runtime.",
14+
)
15+
16+
17+
@pytest.fixture
18+
def converse_response():
19+
return {
20+
"output": {"message": {"role": "assistant", "content": [{"text": "Hello"}]}},
21+
"stopReason": "end_turn",
22+
"usage": {"inputTokens": 50, "outputTokens": 30, "totalTokens": 80},
23+
"metrics": {"latencyMs": 500},
24+
}
25+
26+
27+
@pytest.mark.parametrize("span_streaming", [True, False])
28+
def test_converse(sentry_init, capture_items, span_streaming, converse_response):
29+
sentry_init(
30+
traces_sample_rate=1.0,
31+
integrations=[Boto3Integration()],
32+
trace_lifecycle="stream" if span_streaming else "static",
33+
)
34+
35+
session = boto3.Session(
36+
aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1"
37+
)
38+
client = session.client("bedrock-runtime")
39+
stubber = Stubber(client)
40+
stubber.add_response("converse", converse_response)
41+
42+
items = capture_items()
43+
44+
with stubber, sentry_sdk.start_transaction(name="parent"):
45+
client.converse(
46+
modelId="anthropic.claude-3-sonnet-20240229-v1:0",
47+
messages=[{"role": "user", "content": [{"text": "Hi"}]}],
48+
)
49+
50+
sentry_sdk.flush()
51+
(span,) = (item.payload for item in items if item.type == "span")
52+
53+
assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock"
54+
assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False
55+
56+
57+
@pytest.mark.parametrize("span_streaming", [True, False])
58+
def test_converse_error(sentry_init, capture_items, span_streaming):
59+
sentry_init(
60+
traces_sample_rate=1.0,
61+
integrations=[Boto3Integration()],
62+
trace_lifecycle="stream" if span_streaming else "static",
63+
)
64+
65+
session = boto3.Session(
66+
aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1"
67+
)
68+
client = session.client("bedrock-runtime")
69+
original_exception = ClientError(
70+
error_response={
71+
"Error": {"Code": "AccessDeniedException", "Message": "Access denied"},
72+
"ResponseMetadata": {
73+
"RequestId": "test-id",
74+
"HostId": "",
75+
"HTTPStatusCode": 403,
76+
"HTTPHeaders": {},
77+
"RetryAttempts": 1,
78+
},
79+
},
80+
operation_name="Converse",
81+
)
82+
83+
def raise_client_error(**kwargs):
84+
raise original_exception
85+
86+
client.meta.events.register("before-parameter-build", raise_client_error)
87+
88+
items = capture_items()
89+
90+
with sentry_sdk.start_transaction(name="parent"), pytest.raises(ClientError):
91+
client.converse(
92+
modelId="anthropic.claude-3-sonnet-20240229-v1:0",
93+
messages=[{"role": "user", "content": [{"text": "Hi"}]}],
94+
)
95+
96+
sentry_sdk.flush()
97+
(span,) = (item.payload for item in items if item.type == "span")
98+
99+
assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock"
100+
assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
101+
102+
assert span["attributes"][SPANDATA.ERROR_TYPE] == "AccessDeniedException"
103+
assert span["attributes"][SPANDATA.HTTP_STATUS_CODE] == 403
104+
assert span["attributes"][SPANDATA.AWS_REQUEST_ID] == "test-id"
105+
assert span["attributes"][SPANDATA.HTTP_REQUEST_RESEND_COUNT] == 1
106+
assert span["status"] == "error"
107+
108+
109+
@pytest.mark.parametrize("span_streaming", [True, False])
110+
def test_converse_stream(sentry_init, capture_items, span_streaming):
111+
sentry_init(
112+
traces_sample_rate=1.0,
113+
integrations=[Boto3Integration()],
114+
trace_lifecycle="stream" if span_streaming else "static",
115+
)
116+
117+
session = boto3.Session(
118+
aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1"
119+
)
120+
client = session.client("bedrock-runtime")
121+
stubber = Stubber(client)
122+
stubber.add_response("converse_stream", {})
123+
124+
items = capture_items()
125+
126+
with stubber, sentry_sdk.start_transaction(name="parent"):
127+
client.converse_stream(
128+
modelId="anthropic.claude-3-sonnet-20240229-v1:0",
129+
messages=[{"role": "user", "content": [{"text": "Hi"}]}],
130+
)
131+
132+
sentry_sdk.flush()
133+
(span,) = (item.payload for item in items if item.type == "span")
134+
135+
assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock"
136+
assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True
137+
138+
139+
@pytest.mark.parametrize("span_streaming", [True, False])
140+
def test_converse_stream_error(sentry_init, capture_items, span_streaming):
141+
sentry_init(
142+
traces_sample_rate=1.0,
143+
integrations=[Boto3Integration()],
144+
trace_lifecycle="stream" if span_streaming else "static",
145+
)
146+
147+
session = boto3.Session(
148+
aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1"
149+
)
150+
client = session.client("bedrock-runtime")
151+
original_exception = ClientError(
152+
error_response={
153+
"Error": {"Code": "AccessDeniedException", "Message": "Access denied"},
154+
"ResponseMetadata": {
155+
"RequestId": "test-id",
156+
"HostId": "",
157+
"HTTPStatusCode": 403,
158+
"HTTPHeaders": {},
159+
"RetryAttempts": 1,
160+
},
161+
},
162+
operation_name="ConverseStream",
163+
)
164+
165+
def raise_client_error(**kwargs):
166+
raise original_exception
167+
168+
client.meta.events.register("before-parameter-build", raise_client_error)
169+
170+
items = capture_items()
171+
172+
with sentry_sdk.start_transaction(name="parent"), pytest.raises(ClientError):
173+
client.converse_stream(
174+
modelId="anthropic.claude-3-sonnet-20240229-v1:0",
175+
messages=[{"role": "user", "content": [{"text": "Hi"}]}],
176+
)
177+
178+
sentry_sdk.flush()
179+
(span,) = (item.payload for item in items if item.type == "span")
180+
181+
assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock"
182+
assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True
183+
184+
assert span["attributes"][SPANDATA.ERROR_TYPE] == "AccessDeniedException"
185+
assert span["attributes"][SPANDATA.HTTP_STATUS_CODE] == 403
186+
assert span["attributes"][SPANDATA.AWS_REQUEST_ID] == "test-id"
187+
assert span["attributes"][SPANDATA.HTTP_REQUEST_RESEND_COUNT] == 1
188+
assert span["status"] == "error"

0 commit comments

Comments
 (0)