feat(bedrock): Record request parameters

Author: alexander-alderman-webb

sentry_sdk/integrations/boto3/_services/bedrock.py

@@ -28,12 +28,31 @@ 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",
         }
 
+        inference_config = ctx.params.get("inferenceConfig")
+
+        if not isinstance(inference_config, dict):
+            return attributes
+
+        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
+
+        top_p = inference_config.get("topP")
+        if top_p is not None:
+            attributes[SPANDATA.GEN_AI_REQUEST_TOP_P] = top_p
+
+        return attributes
+
     def get_response_attributes(
         self,
         ctx: "AwsCallContext",

tests/integrations/boto3/test_bedrock.py

@@ -55,6 +55,11 @@ def test_converse(
             client.converse(
                 modelId="anthropic.claude-3-sonnet-20240229-v1:0",
                 messages=[{"role": "user", "content": [{"text": "Hi"}]}],
+                inferenceConfig={
+                    "maxTokens": 256,
+                    "temperature": 0.7,
+                    "topP": 0.9,
+                },
             )
 
         sentry_sdk.flush()
@@ -68,13 +73,22 @@ def test_converse(
         assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock"
         assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
         assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False
+
+        assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 256
+        assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
+        assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9
     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"}]}],
+                inferenceConfig={
+                    "maxTokens": 256,
+                    "temperature": 0.7,
+                    "topP": 0.9,
+                },
             )
 
         (transaction,) = events
@@ -84,6 +98,10 @@ def test_converse(
         assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
         assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False
 
+        assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 256
+        assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
+        assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9
+
 
 @pytest.mark.parametrize("span_streaming", [True, False])
 @pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
@@ -193,6 +211,11 @@ def test_converse_stream(
             client.converse_stream(
                 modelId="anthropic.claude-3-sonnet-20240229-v1:0",
                 messages=[{"role": "user", "content": [{"text": "Hi"}]}],
+                inferenceConfig={
+                    "maxTokens": 256,
+                    "temperature": 0.7,
+                    "topP": 0.9,
+                },
             )
 
         sentry_sdk.flush()
@@ -206,13 +229,22 @@ def test_converse_stream(
         assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock"
         assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
         assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True
+
+        assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 256
+        assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
+        assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9
     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"}]}],
+                inferenceConfig={
+                    "maxTokens": 256,
+                    "temperature": 0.7,
+                    "topP": 0.9,
+                },
             )
 
         (transaction,) = events
@@ -222,6 +254,10 @@ def test_converse_stream(
         assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
         assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True
 
+        assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 256
+        assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
+        assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9
+
 
 @pytest.mark.parametrize("span_streaming", [True, False])
 @pytest.mark.parametrize("stream_gen_ai_spans", [True, False])