diff --git a/sentry_sdk/integrations/boto3/_services/bedrock.py b/sentry_sdk/integrations/boto3/_services/bedrock.py index e6e1446aa3..d322727f38 100644 --- a/sentry_sdk/integrations/boto3/_services/bedrock.py +++ b/sentry_sdk/integrations/boto3/_services/bedrock.py @@ -1,3 +1,4 @@ +import json from typing import TYPE_CHECKING from sentry_sdk.consts import OP, SPANDATA @@ -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 @@ -28,15 +29,69 @@ def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": 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) + ) + + 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 diff --git a/tests/integrations/boto3/test_bedrock.py b/tests/integrations/boto3/test_bedrock.py index 5219bdf8b2..8594833a76 100644 --- a/tests/integrations/boto3/test_bedrock.py +++ b/tests/integrations/boto3/test_bedrock.py @@ -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.", @@ -167,6 +173,103 @@ def raise_client_error(**kwargs): 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( @@ -305,3 +408,99 @@ def raise_client_error(**kwargs): 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 + )