add AI21 Jamba and Mistral AI model test cases

Author: ezhang6811

test/contract-tests/images/applications/TestSimpleApp.AWSSDK.Core/BedrockTests.cs

@@ -160,6 +160,83 @@ public object InvokeModelCohereCommandResponse()
         };
     }
 
+    public void InvokeModelAi21Jamba()
+    {
+        bedrockRuntime.InvokeModelAsync(new InvokeModelRequest
+        {
+            ModelId = "ai21.jamba-1-5-large-v1:0",
+            Body = new MemoryStream(Encoding.UTF8.GetBytes(JsonSerializer.Serialize(new
+            {
+                messages = new object[]
+                {
+                    new
+                    {
+                        role = "USER",
+                        content = "sample input text",
+                    },
+                },
+                temperature = 0.123,
+                top_p = 0.456,
+                max_tokens = 123,
+            }))),
+            ContentType = "application/json",
+        });
+        return;
+    }
+
+    public object InvokeModelAi21JambaResponse()
+    {
+        return new
+        {
+            choices = new object[]
+            {
+                new
+                {
+                    finish_reason = "finish_reason",
+                },
+            },
+            usage = new
+            {
+                prompt_tokens = 456,
+                completion_tokens = 789,
+            },
+        };
+    }
+
+    public void InvokeModelMistralAi()
+    {
+        bedrockRuntime.InvokeModelAsync(new InvokeModelRequest
+        {
+            ModelId = "mistral.mistral-7b-instruct-v0:2",
+            Body = new MemoryStream(Encoding.UTF8.GetBytes(JsonSerializer.Serialize(new
+            {
+                // prompt is 72 chars long, input_tokens should be estimated as ceil(72/6) = 12
+                prompt = "sample input text sample input text sample input text sample input text ",
+                temperature = 0.123,
+                top_p = 0.456,
+                max_tokens = 123,
+            }))),
+            ContentType = "application/json",
+        });
+        return;
+    }
+
+    public object InvokeModelMistralAiResponse()
+    {
+        return new
+        {
+            outputs = new object[]
+            {
+                new
+                {
+                    // response is 56 chars long, output_tokens should be estimated as ceil(56/6) = 10
+                    text = "sample output text sample output text sample output text",
+                    stop_reason = "finish_reason",
+                },
+            },
+        };
+    }
+
     public Task<GetAgentResponse> GetAgent()
     {
         return bedrockAgent.GetAgentAsync(new GetAgentRequest

test/contract-tests/images/applications/TestSimpleApp.AWSSDK.Core/Program.cs

@@ -141,6 +141,14 @@
     .WithName("invoke-model-command")
     .WithOpenApi();
 
+app.MapGet("bedrock/invokemodel/invoke-model-jamba", (BedrockTests bedrock) => bedrock.InvokeModelAi21Jamba())
+    .WithName("invoke-model-jamba")
+    .WithOpenApi();
+
+app.MapGet("bedrock/invokemodel/invoke-model-mistral", (BedrockTests bedrock) => bedrock.InvokeModelMistralAi())
+    .WithName("invoke-model-mistral")
+    .WithOpenApi();
+
 app.MapGet("bedrock/getagent/get-agent", (BedrockTests bedrock) => bedrock.GetAgent())
     .WithName("get-agent")
     .WithOpenApi();
@@ -168,6 +176,8 @@
 app.MapPost("model/anthropic.claude-v2:1/invoke", (BedrockTests bedrock) => bedrock.InvokeModelAnthropicClaudeResponse());
 app.MapPost("model/meta.llama3-8b-instruct-v1:0/invoke", (BedrockTests bedrock) => bedrock.InvokeModelMetaLlamaResponse());
 app.MapPost("model/cohere.command-r-v1:0/invoke", (BedrockTests bedrock) => bedrock.InvokeModelCohereCommandResponse());
+app.MapPost("model/ai21.jamba-1-5-large-v1:0/invoke", (BedrockTests bedrock) => bedrock.InvokeModelAi21JambaResponse());
+app.MapPost("model/mistral.mistral-7b-instruct-v0:2/invoke", (BedrockTests bedrock) => bedrock.InvokeModelMistralAiResponse());
 app.MapGet("agents/test-agent", (BedrockTests bedrock) => bedrock.GetAgentResponse());
 app.MapGet("knowledgebases/test-knowledge-base", (BedrockTests bedrock) => bedrock.GetKnowledgeBaseResponse());
 app.MapGet("knowledgebases/test-knowledge-base/datasources/test-data-source", (BedrockTests bedrock) => bedrock.GetDataSourceResponse());

test/contract-tests/tests/test/amazon/awssdk/awssdk_test.py

@@ -418,6 +418,54 @@ def test_bedrock_runtime_invoke_model_command(self):
             },
             span_name="Bedrock Runtime.InvokeModel",
         )
+    
+    def test_bedrock_runtime_invoke_model_jamba(self):
+        self.do_test_requests(
+            "bedrock/invokemodel/invoke-model-jamba",
+            "GET",
+            200,
+            0,
+            0,
+            rpc_service="Bedrock Runtime",
+            remote_service="AWS::BedrockRuntime",
+            remote_operation="InvokeModel",
+            remote_resource_type="AWS::Bedrock::Model",
+            remote_resource_identifier="ai21.jamba-1-5-large-v1:0",
+            request_specific_attributes={
+                _GEN_AI_REQUEST_MODEL: "ai21.jamba-1-5-large-v1:0",
+                _GEN_AI_REQUEST_TEMPERATURE: 0.123,
+                _GEN_AI_REQUEST_TOP_P: 0.456,
+                _GEN_AI_REQUEST_MAX_TOKENS: 123,
+                _GEN_AI_USAGE_INPUT_TOKENS: 456,
+                _GEN_AI_USAGE_OUTPUT_TOKENS: 789,
+                _GEN_AI_RESPONSE_FINISH_REASONS: ["finish_reason"],
+            },
+            span_name="Bedrock Runtime.InvokeModel",
+        )
+
+    def test_bedrock_runtime_invoke_model_mistral(self):
+        self.do_test_requests(
+            "bedrock/invokemodel/invoke-model-mistral",
+            "GET",
+            200,
+            0,
+            0,
+            rpc_service="Bedrock Runtime",
+            remote_service="AWS::BedrockRuntime",
+            remote_operation="InvokeModel",
+            remote_resource_type="AWS::Bedrock::Model",
+            remote_resource_identifier="mistral.mistral-7b-instruct-v0:2",
+            request_specific_attributes={
+                _GEN_AI_REQUEST_MODEL: "mistral.mistral-7b-instruct-v0:2",
+                _GEN_AI_REQUEST_TEMPERATURE: 0.123,
+                _GEN_AI_REQUEST_TOP_P: 0.456,
+                _GEN_AI_REQUEST_MAX_TOKENS: 123,
+                _GEN_AI_USAGE_INPUT_TOKENS: 12,
+                _GEN_AI_USAGE_OUTPUT_TOKENS: 10,
+                _GEN_AI_RESPONSE_FINISH_REASONS: ["finish_reason"],
+            },
+            span_name="Bedrock Runtime.InvokeModel",
+        )
 
     def test_bedrock_agent_runtime_invoke_agent(self):
         self.do_test_requests(
@@ -687,6 +735,8 @@ def _filter_bedrock_metrics(self, target_metrics: List[Metric]):
             "POST model/anthropic.claude-v2:1/invoke",
             "POST model/meta.llama3-8b-instruct-v1:0/invoke",
             "POST model/cohere.command-r-v1:0/invoke",
+            "POST model/ai21.jamba-1-5-large-v1:0/invoke",
+            "POST model/mistral.mistral-7b-instruct-v0:2/invoke",
             "POST knowledgebases/test-knowledge-base/retrieve"
         }
         for metric in target_metrics: