Update models used in bedrock system tests (#55229)

o-nikolas · web-flow · commit ed3d3e188178 · 2025-09-03T19:52:30.000-07:00
The anthropic.claude-3-sonnet-20240229-v1:0 model is being deprecated so
we need to move to a new model. Also updating the input schema that the
newer model expects.
diff --git a/providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py b/providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py
@@ -67,7 +67,7 @@
 #   the Amazon Bedrock console and may take up to 24 hours to apply:
 #######################################################################
 
-CLAUDE_MODEL_ID = "anthropic.claude-3-sonnet-20240229-v1:0"
+CLAUDE_MODEL_ID = "anthropic.claude-3-5-sonnet-20241022-v2:0"
 ANTHROPIC_VERSION = "bedrock-2023-05-31"
 
 # Batch inferences currently require a minimum of 100 prompts per batch.
diff --git a/providers/amazon/tests/system/amazon/aws/example_bedrock_retrieve_and_generate.py b/providers/amazon/tests/system/amazon/aws/example_bedrock_retrieve_and_generate.py
@@ -81,8 +81,9 @@
 #   the Amazon Bedrock console and may take up to 24 hours to apply:
 #######################################################################
 
-CLAUDE_MODEL_ID = "anthropic.claude-v2"
+CLAUDE_MODEL_ID = "anthropic.claude-3-5-sonnet-20241022-v2:0"
 TITAN_MODEL_ID = "amazon.titan-embed-text-v1"
+ANTHROPIC_VERSION = "bedrock-2023-05-31"
 
 # Externally fetched variables:
 ROLE_ARN_KEY = "ROLE_ARN"
@@ -487,7 +488,13 @@ def delete_opensearch_policies(collection_name: str):
     invoke_claude_completions = BedrockInvokeModelOperator(
         task_id="invoke_claude_completions",
         model_id=CLAUDE_MODEL_ID,
-        input_data={"max_tokens_to_sample": 4000, "prompt": f"\n\nHuman: {PROMPT}\n\nAssistant:"},
+        input_data={
+            "max_tokens": 4000,
+            "anthropic_version": ANTHROPIC_VERSION,
+            "messages": [
+                {"role": "user", "content": f"\n\nHuman: {PROMPT}\n\nAssistant:"},
+            ],
+        },
     )
     # [END howto_operator_invoke_claude_model]
 
