Skip to content

feat(bedrock): Set gen_ai.tool.definitions attribute #8233

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
alexander-alderman-webb wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

feat(bedrock): Set gen_ai.tool.definitions attribute #8233

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 57 additions & 2 deletions sentry_sdk/integrations/boto3/_services/bedrock.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import json
from typing import TYPE_CHECKING

from sentry_sdk.consts import OP, SPANDATA
Expand All @@ -6,7 +7,7 @@
if TYPE_CHECKING:
from typing import Any, Optional

from sentry_sdk._types import Attributes
from sentry_sdk._types import Attributes, ToolDefinition
from sentry_sdk.integrations.boto3._context import AwsCallContext


Expand All @@ -28,15 +29,69 @@
if ctx.operation_name != "Converse" and ctx.operation_name != "ConverseStream":
return {}

return {
attributes: "Attributes" = {
SPANDATA.GEN_AI_PROVIDER_NAME: "aws.bedrock",
SPANDATA.GEN_AI_OPERATION_NAME: "chat",
SPANDATA.GEN_AI_RESPONSE_STREAMING: ctx.operation_name == "ConverseStream",
}

tool_config = ctx.params.get("toolConfig")
if not isinstance(tool_config, dict):
return attributes

tools = tool_config.get("tools")
if not isinstance(tools, list):
return attributes

attributes[SPANDATA.GEN_AI_TOOL_DEFINITIONS] = json.dumps(
_transform_tool_definitions(tools)
)

Check warning on line 48 in sentry_sdk/integrations/boto3/_services/bedrock.py

View check run for this annotation

@sentry/warden / warden: code-review

Bedrock tool definitions ignore the GenAI input opt-out 

When a user explicitly sets `data_collection.gen_ai.inputs` to `false`, Bedrock still serializes tool names, descriptions, and parameter schemas into `gen_ai.tool.definitions`. Please honor that opt-out when setting the attribute, while preserving the legacy behavior when data collection was not explicitly configured.
Comment thread
alexander-alderman-webb marked this conversation as resolved.
Resolved

return attributes

def get_response_attributes(
self,
ctx: "AwsCallContext",
response: "Any",
) -> "Attributes":
return {}


def _transform_tool_definitions(
tools: "list[dict[str, Any]]",
) -> "list[ToolDefinition]":
"""
Transform tool definitions to the schema used by the "gen_ai.tool.definitions" attribute.
"""
result: "list[ToolDefinition]" = []
for tool in tools:
if not isinstance(tool, dict):
continue

if "toolSpec" in tool:
tool_definition: "ToolDefinition" = {"type": "function"}
tool_spec = tool["toolSpec"]
if not isinstance(tool_spec, dict):
result.append(tool_definition)
continue

if "name" in tool_spec:
tool_definition["name"] = tool_spec["name"]
if "description" in tool_spec:
tool_definition["description"] = tool_spec["description"]
if "inputSchema" in tool_spec and "json" in tool_spec["inputSchema"]:
tool_definition["parameters"] = tool_spec["inputSchema"]["json"]

result.append(tool_definition)
continue

if "systemTool" not in tool:
continue

system_tool = tool["systemTool"]
if not isinstance(system_tool, dict) or "name" not in system_tool:
continue

result.append({"type": system_tool["name"]})

return result
199 changes: 199 additions & 0 deletions tests/integrations/boto3/test_bedrock.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,19 @@
import json

import boto3
import botocore
import pytest
from botocore.exceptions import ClientError
from botocore.loaders import Loader
from botocore.stub import Stubber
from packaging.version import Version

import sentry_sdk
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.integrations.boto3 import Boto3Integration

BOTOCORE_VERSION = Version(botocore.__version__)

pytestmark = pytest.mark.skipif(
"bedrock-runtime" not in Loader().list_available_services("service-2"),
reason="botocore is too old for bedrock-runtime.",
Expand Down Expand Up @@ -167,6 +173,103 @@
assert span["status"] == "internal_error"


@pytest.mark.parametrize("span_streaming", [True, False])
@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
def test_converse_tool_definitions(
sentry_init,
capture_items,
capture_events,
span_streaming,
stream_gen_ai_spans,
converse_response,
):
sentry_init(
traces_sample_rate=1.0,
integrations=[Boto3Integration()],
stream_gen_ai_spans=stream_gen_ai_spans,
trace_lifecycle="stream" if span_streaming else "static",
)

tools: list = [
{
"toolSpec": {
"name": "get_weather",
"description": "Get the current weather in a given location",
"inputSchema": {
"json": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
}
},
}
},
]
expected_tool_definitions = [
{
"type": "function",
"name": "get_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
},
]

if BOTOCORE_VERSION >= Version("1.40.62"):
tools.append({"systemTool": {"name": "web_search"}})
expected_tool_definitions.append({"type": "web_search"})

session = boto3.Session(
aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1"
)
client = session.client("bedrock-runtime")
stubber = Stubber(client)
stubber.add_response("converse", converse_response)

if span_streaming or stream_gen_ai_spans:
items = capture_items("span")

with stubber, sentry_sdk.start_transaction(name="parent"):
client.converse(
modelId="anthropic.claude-3-sonnet-20240229-v1:0",
messages=[{"role": "user", "content": [{"text": "Hi"}]}],
toolConfig={"tools": tools},
)

sentry_sdk.flush()
spans = [item.payload for item in items]
(span,) = (
span
for span in spans
if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT
)

assert (
json.loads(span["attributes"][SPANDATA.GEN_AI_TOOL_DEFINITIONS])
== expected_tool_definitions
)
else:
events = capture_events()

with stubber, sentry_sdk.start_transaction(name="parent"):
client.converse(
modelId="anthropic.claude-3-sonnet-20240229-v1:0",
messages=[{"role": "user", "content": [{"text": "Hi"}]}],
toolConfig={"tools": tools},
)

(transaction,) = events
(span,) = transaction["spans"]

assert (
json.loads(span["data"][SPANDATA.GEN_AI_TOOL_DEFINITIONS])
== expected_tool_definitions
)


@pytest.mark.parametrize("span_streaming", [True, False])
@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
def test_converse_stream(
Expand Down Expand Up @@ -305,3 +408,99 @@
assert span["data"][SPANDATA.AWS_REQUEST_ID] == "test-id"
assert span["data"][SPANDATA.HTTP_REQUEST_RESEND_COUNT] == 1
assert span["status"] == "internal_error"


@pytest.mark.parametrize("span_streaming", [True, False])
@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
def test_converse_stream_tool_definitions(
sentry_init,
capture_items,
capture_events,
span_streaming,
stream_gen_ai_spans,
):
sentry_init(
traces_sample_rate=1.0,
integrations=[Boto3Integration()],
stream_gen_ai_spans=stream_gen_ai_spans,
trace_lifecycle="stream" if span_streaming else "static",
)

tools: list = [
{
"toolSpec": {
"name": "get_weather",
"description": "Get the current weather in a given location",
"inputSchema": {
"json": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
}
},
}
},
]
expected_tool_definitions = [
{
"type": "function",
"name": "get_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
},
]

if BOTOCORE_VERSION >= Version("1.40.62"):
tools.append({"systemTool": {"name": "web_search"}})
expected_tool_definitions.append({"type": "web_search"})

session = boto3.Session(
aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1"
)
client = session.client("bedrock-runtime")
stubber = Stubber(client)
stubber.add_response("converse_stream", {})

if span_streaming or stream_gen_ai_spans:
items = capture_items("span")

with stubber, sentry_sdk.start_transaction(name="parent"):
client.converse_stream(
modelId="anthropic.claude-3-sonnet-20240229-v1:0",
messages=[{"role": "user", "content": [{"text": "Hi"}]}],
toolConfig={"tools": tools},
)

sentry_sdk.flush()
spans = [item.payload for item in items]
(span,) = (
span
for span in spans
if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT
)

assert (
json.loads(span["attributes"][SPANDATA.GEN_AI_TOOL_DEFINITIONS])
== expected_tool_definitions
)
else:
events = capture_events()

with stubber, sentry_sdk.start_transaction(name="parent"):
client.converse_stream(
modelId="anthropic.claude-3-sonnet-20240229-v1:0",
messages=[{"role": "user", "content": [{"text": "Hi"}]}],
toolConfig={"tools": tools},
)

(transaction,) = events
(span,) = transaction["spans"]

assert (
json.loads(span["data"][SPANDATA.GEN_AI_TOOL_DEFINITIONS])
== expected_tool_definitions

Check warning on line 505 in tests/integrations/boto3/test_bedrock.py

View check run for this annotation

@sentry/warden / warden: code-review

[6AC-HUG] Bedrock tool definitions ignore the GenAI input opt-out (additional location) 

When a user explicitly sets `data_collection.gen_ai.inputs` to `false`, Bedrock still serializes tool names, descriptions, and parameter schemas into `gen_ai.tool.definitions`. Please honor that opt-out when setting the attribute, while preserving the legacy behavior when data collection was not explicitly configured.
)
Loading