diff --git a/sentry_sdk/integrations/boto3/_services/bedrock.py b/sentry_sdk/integrations/boto3/_services/bedrock.py index 1fc11d7581..d99475cedc 100644 --- a/sentry_sdk/integrations/boto3/_services/bedrock.py +++ b/sentry_sdk/integrations/boto3/_services/bedrock.py @@ -1,15 +1,60 @@ +import json from typing import TYPE_CHECKING +import sentry_sdk from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations.boto3._services.base import _ServiceExtension +from sentry_sdk.utils import has_data_collection_enabled 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 +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"] + + input_schema = tool_spec.get("inputSchema") + if isinstance(input_schema, dict) and "json" in input_schema: + tool_definition["parameters"] = input_schema["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 + + class _BedrockExtension(_ServiceExtension): __slots__ = () skip_if_no_active_span = False @@ -40,20 +85,35 @@ def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": inference_config = ctx.params.get("inferenceConfig") - if not isinstance(inference_config, dict): - return attributes + if isinstance(inference_config, dict): + max_tokens = inference_config.get("maxTokens") + if max_tokens is not None: + attributes[SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] = max_tokens + + temperature = inference_config.get("temperature") + if temperature is not None: + attributes[SPANDATA.GEN_AI_REQUEST_TEMPERATURE] = temperature - max_tokens = inference_config.get("maxTokens") - if max_tokens is not None: - attributes[SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] = max_tokens + top_p = inference_config.get("topP") + if top_p is not None: + attributes[SPANDATA.GEN_AI_REQUEST_TOP_P] = top_p - temperature = inference_config.get("temperature") - if temperature is not None: - attributes[SPANDATA.GEN_AI_REQUEST_TEMPERATURE] = temperature + client_options = sentry_sdk.get_client().options + if has_data_collection_enabled(client_options): + if not client_options["data_collection"]["gen_ai"]["inputs"]: + return attributes + + 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 - top_p = inference_config.get("topP") - if top_p is not None: - attributes[SPANDATA.GEN_AI_REQUEST_TOP_P] = top_p + attributes[SPANDATA.GEN_AI_TOOL_DEFINITIONS] = json.dumps( + _transform_tool_definitions(tools) + ) return attributes diff --git a/tests/integrations/boto3/test_bedrock.py b/tests/integrations/boto3/test_bedrock.py index 1040231230..340a3712b6 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.", @@ -219,6 +225,114 @@ def raise_client_error(**kwargs): assert span["status"] == "internal_error" +@pytest.mark.parametrize("data_collection", [True, False]) +@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, + data_collection, + span_streaming, + stream_gen_ai_spans, + converse_response, +): + if data_collection: + sentry_init( + traces_sample_rate=1.0, + integrations=[Boto3Integration()], + data_collection={"gen_ai": {"inputs": True}}, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + else: + 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( @@ -409,3 +523,110 @@ 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("data_collection", [True, False]) +@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, + data_collection, + span_streaming, + stream_gen_ai_spans, +): + if data_collection: + sentry_init( + traces_sample_rate=1.0, + integrations=[Boto3Integration()], + data_collection={"gen_ai": {"inputs": True}}, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + else: + 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 + )