- added support for additional gen ai attributes

- added contract tests for amazon titan and anthropic claude models
- added support to compare float attributes

Author: liustve

contract-tests/images/applications/aws-sdk/server.js

@@ -554,19 +554,39 @@ async function handleBedrockRequest(req, res, path) {
       res.statusCode = 200;
     } else if (path.includes('invokemodel/invoke-model')) {
       await withInjected200Success(bedrockRuntimeClient, ['InvokeModelCommand'], {}, async () => {
-        const modelId = 'amazon.titan-text-premier-v1:0';
+        let modelId = ''
+        let body = {}
         const userMessage = "Describe the purpose of a 'hello world' program in one line.";
         const prompt = `<s>[INST] ${userMessage} [/INST]`;
 
-        const body = JSON.stringify({
-          inputText: prompt,
-          textGenerationConfig: {
-            maxTokenCount: 3072,
-            stopSequences: [],
-            temperature: 0.7,
-            topP: 0.9,
-          },
-        });
+        if (path.includes('amazon.titan')) {
+            modelId = 'amazon.titan-text-premier-v1:0';
+            body = JSON.stringify({
+              inputText: prompt,
+              textGenerationConfig: {
+                maxTokenCount: 3072,
+                stopSequences: [],
+                temperature: 0.7,
+                topP: 0.9,
+              },
+            });
+        }
+
+        if (path.includes('anthropic.claude')) {
+            modelId = 'anthropic.claude-v2:1';
+            body = JSON.stringify({
+              anthropic_version: 'bedrock-2023-05-31',
+              max_tokens: 1000,
+              temperature: 1.1,
+              top_p: 1,
+              messages: [
+                {
+                  role: 'user',
+                  content: [{ type: 'text', text: prompt }],
+                },
+              ],
+            });
+        }      
 
         await bedrockRuntimeClient.send(
           new InvokeModelCommand({
@@ -577,6 +597,7 @@ async function handleBedrockRequest(req, res, path) {
           })
         );
       });
+      
       res.statusCode = 200;
     } else {
       res.statusCode = 404;

contract-tests/tests/test/amazon/aws-sdk/aws_sdk_test.py

@@ -34,6 +34,9 @@
 _AWS_BEDROCK_KNOWLEDGE_BASE_ID: str = "aws.bedrock.knowledge_base.id"
 _AWS_BEDROCK_DATA_SOURCE_ID: str = "aws.bedrock.data_source.id"
 _GEN_AI_REQUEST_MODEL: str = "gen_ai.request.model"
+_GEN_AI_REQUEST_TEMPERATURE: str = "gen_ai.request.temperature"
+_GEN_AI_REQUEST_TOP_P: str = "gen_ai.request.top_p"
+_GEN_AI_REQUEST_MAX_TOKENS: str = "gen_ai.request.max_tokens"
 
 
 # pylint: disable=too-many-public-methods
@@ -406,9 +409,9 @@ def test_kinesis_fault(self):
             span_name="Kinesis.PutRecord",
         )
 
-    def test_bedrock_runtime_invoke_model(self):
+    def test_bedrock_runtime_invoke_model_amazon_titan(self):
         self.do_test_requests(
-            "bedrock/invokemodel/invoke-model",
+            "bedrock/invokemodel/invoke-model/amazon.titan-text-premier-v1:0",
             "GET",
             200,
             0,
@@ -418,11 +421,36 @@ def test_bedrock_runtime_invoke_model(self):
             remote_service="AWS::BedrockRuntime",
             remote_operation="InvokeModel",
             remote_resource_type="AWS::Bedrock::Model",
-            remote_resource_identifier="amazon.titan-text-premier-v1:0",
+            remote_resource_identifier='amazon.titan-text-premier-v1:0',
             request_specific_attributes={
-                _GEN_AI_REQUEST_MODEL: "amazon.titan-text-premier-v1:0",
-            },
-            span_name="BedrockRuntime.InvokeModel",
+                _GEN_AI_REQUEST_MODEL: 'amazon.titan-text-premier-v1:0',
+                _GEN_AI_REQUEST_MAX_TOKENS: 3072,
+                _GEN_AI_REQUEST_TEMPERATURE: 0.7,
+                _GEN_AI_REQUEST_TOP_P: 0.9
+                },
+            span_name="BedrockRuntime.InvokeModel"
+        )
+
+    def test_bedrock_runtime_invoke_model_anthropic_claude(self):
+        self.do_test_requests(
+            "bedrock/invokemodel/invoke-model/anthropic.claude-v2:1",
+            "GET",
+            200,
+            0,
+            0,
+            local_operation="GET /bedrock",
+            rpc_service="BedrockRuntime",
+            remote_service="AWS::BedrockRuntime",
+            remote_operation="InvokeModel",
+            remote_resource_type="AWS::Bedrock::Model",
+            remote_resource_identifier='anthropic.claude-v2:1',
+            request_specific_attributes={
+                _GEN_AI_REQUEST_MODEL: 'anthropic.claude-v2:1',
+                _GEN_AI_REQUEST_MAX_TOKENS: 1000,
+                _GEN_AI_REQUEST_TEMPERATURE: 1.1,
+                _GEN_AI_REQUEST_TOP_P: 1
+                },
+            span_name="BedrockRuntime.InvokeModel"
         )
 
     def test_bedrock_get_guardrail(self):
@@ -538,6 +566,9 @@ def test_bedrock_agent_get_data_source(self):
             },
             span_name="BedrockAgent.GetDataSource",
         )
+    
+    # def test_bedrock_agent_runtime_invoke_agent(self):
+    #     return None
 
     @override
     def _assert_aws_span_attributes(self, resource_scope_spans: List[ResourceScopeSpan], path: str, **kwargs) -> None:
@@ -591,6 +622,7 @@ def _assert_aws_attributes(
     def _assert_semantic_conventions_span_attributes(
         self, resource_scope_spans: List[ResourceScopeSpan], method: str, path: str, status_code: int, **kwargs
     ) -> None:
+
         target_spans: List[Span] = []
         for resource_scope_span in resource_scope_spans:
             # pylint: disable=no-member
@@ -629,6 +661,8 @@ def _assert_semantic_conventions_attributes(
                 self._assert_str_attribute(attributes_dict, key, value)
             elif isinstance(value, int):
                 self._assert_int_attribute(attributes_dict, key, value)
+            elif isinstance(value, float):
+                self._assert_float_attribute(attributes_dict, key, value)
             else:
                 self._assert_array_value_ddb_table_name(attributes_dict, key, value)
 

contract-tests/tests/test/amazon/base/contract_test_base.py

@@ -230,6 +230,12 @@ def _assert_int_attribute(self, attributes_dict: Dict[str, AnyValue], key: str,
         self.assertIsNotNone(actual_value)
         self.assertEqual(expected_value, actual_value.int_value)
 
+    def _assert_float_attribute(self, attributes_dict: Dict[str, AnyValue], key: str, expected_value: float) -> None:
+        self.assertIn(key, attributes_dict)
+        actual_value: AnyValue = attributes_dict[key]
+        self.assertIsNotNone(actual_value)
+        self.assertEqual(expected_value, actual_value.double_value)
+
     def check_sum(self, metric_name: str, actual_sum: float, expected_sum: float) -> None:
         if metric_name is LATENCY_METRIC:
             self.assertTrue(0 < actual_sum < expected_sum)