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

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
 

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
+        )