add contract tests for cohere command-r model

aws-observability · Nov 12, 2024 · f8c036f · f8c036f
1 parent 2d94b6e
commit f8c036f
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 1 deletion.
diff --git a/contract-tests/images/applications/aws-sdk/server.js b/contract-tests/images/applications/aws-sdk/server.js
@@ -648,6 +648,28 @@ async function handleBedrockRequest(req, res, path) {
               prompt: prompt,
             };
           }
+
+          if (path.includes('cohere.command-r')) {
+            modelId = 'cohere.command-r-v1:0';
+
+            request_body = {
+              message: prompt,
+              max_tokens: 512,
+              temperature: 0.5,
+              p: 0.65,
+            };
+
+            response_body = {
+              finish_reason: 'COMPLETE',
+              text: 'test-generation-text',
+              prompt: prompt,
+              request: {
+                commandInput: {
+                  modelId: modelId,
+                },
+              },
+            }
+          }
 
           if (path.includes('ai21.jamba')) {
             modelId = 'ai21.jamba-1-5-large-v1:0';
@@ -678,7 +700,7 @@ async function handleBedrockRequest(req, res, path) {
             }
           }
 
-          if (path.includes('mistral.mistral')) {
+          if (path.includes('mistral')) {
             modelId = 'mistral.mistral-7b-instruct-v0:2';
 
             request_body = {
@@ -761,3 +783,4 @@ prepareAwsServer().then(() => {
     console.log('Ready');
   });
 });
+
diff --git a/contract-tests/tests/test/amazon/aws-sdk/aws_sdk_test.py b/contract-tests/tests/test/amazon/aws-sdk/aws_sdk_test.py
@@ -493,7 +493,35 @@ def test_bedrock_runtime_invoke_model_meta_llama(self):
                 },
             span_name="BedrockRuntime.InvokeModel"
         )
+
+    def test_bedrock_runtime_invoke_model_cohere_command_r(self):
+        self.do_test_requests(
+            "bedrock/invokemodel/invoke-model/cohere.command-r-v1:0",
+            "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='cohere.command-r-v1:0',
+            request_specific_attributes={
+                _GEN_AI_REQUEST_MODEL: 'cohere.command-r-v1:0',
+                _GEN_AI_REQUEST_MAX_TOKENS: 512,
+                _GEN_AI_REQUEST_TEMPERATURE: 0.5,
+                _GEN_AI_REQUEST_TOP_P: 0.65
+                },
+            response_specific_attributes={
+                _GEN_AI_RESPONSE_FINISH_REASONS: ['COMPLETE'],
+                _GEN_AI_USAGE_INPUT_TOKENS: math.ceil(len("Describe the purpose of a 'hello world' program in one line.") / 6),
+                _GEN_AI_USAGE_OUTPUT_TOKENS: math.ceil(len("test-generation-text") / 6)
+                },
+            span_name="BedrockRuntime.InvokeModel"
+        )
 
+    # Delete once this model is fully deprecated on node
     def test_bedrock_runtime_invoke_model_cohere_command(self):
         self.do_test_requests(
             "bedrock/invokemodel/invoke-model/cohere.command-light-text-v14",