From 2280594408d30f425f9c950f9fe74a56a03b27d5 Mon Sep 17 00:00:00 2001 From: Shardul Vaidya <31039336+5herlocked@users.noreply.github.com> Date: Wed, 4 Jun 2025 02:00:41 -0400 Subject: [PATCH] bedrock: Allow users to pick Thinking vs. Non-Thinking models (#31600) Release Notes: - bedrock: Added ability to pick between Thinking and Non-Thinking models --- crates/bedrock/src/models.rs | 95 ++++++++++++++++++++++++++++++++++-- 1 file changed, 92 insertions(+), 3 deletions(-) diff --git a/crates/bedrock/src/models.rs b/crates/bedrock/src/models.rs index 3a2a6d9db5..d46828cb40 100644 --- a/crates/bedrock/src/models.rs +++ b/crates/bedrock/src/models.rs @@ -129,6 +129,60 @@ impl Model { } pub fn id(&self) -> &str { + match self { + Model::ClaudeSonnet4 => "claude-4-sonnet", + Model::ClaudeSonnet4Thinking => "claude-4-sonnet-thinking", + Model::ClaudeOpus4 => "claude-4-opus", + Model::ClaudeOpus4Thinking => "claude-4-opus-thinking", + Model::Claude3_5SonnetV2 => "claude-3-5-sonnet-v2", + Model::Claude3_5Sonnet => "claude-3-5-sonnet", + Model::Claude3Opus => "claude-3-opus", + Model::Claude3Sonnet => "claude-3-sonnet", + Model::Claude3Haiku => "claude-3-haiku", + Model::Claude3_5Haiku => "claude-3-5-haiku", + Model::Claude3_7Sonnet => "claude-3-7-sonnet", + Model::Claude3_7SonnetThinking => "claude-3-7-sonnet-thinking", + Model::AmazonNovaLite => "amazon-nova-lite", + Model::AmazonNovaMicro => "amazon-nova-micro", + Model::AmazonNovaPro => "amazon-nova-pro", + Model::AmazonNovaPremier => "amazon-nova-premier", + Model::DeepSeekR1 => "deepseek-r1", + Model::AI21J2GrandeInstruct => "ai21-j2-grande-instruct", + Model::AI21J2JumboInstruct => "ai21-j2-jumbo-instruct", + Model::AI21J2Mid => "ai21-j2-mid", + Model::AI21J2MidV1 => "ai21-j2-mid-v1", + Model::AI21J2Ultra => "ai21-j2-ultra", + Model::AI21J2UltraV1_8k => "ai21-j2-ultra-v1-8k", + Model::AI21J2UltraV1 => "ai21-j2-ultra-v1", + Model::AI21JambaInstructV1 => "ai21-jamba-instruct-v1", + Model::AI21Jamba15LargeV1 => "ai21-jamba-1-5-large-v1", + Model::AI21Jamba15MiniV1 => "ai21-jamba-1-5-mini-v1", + Model::CohereCommandTextV14_4k => "cohere-command-text-v14-4k", + Model::CohereCommandRV1 => "cohere-command-r-v1", + Model::CohereCommandRPlusV1 => "cohere-command-r-plus-v1", + Model::CohereCommandLightTextV14_4k => "cohere-command-light-text-v14-4k", + Model::MetaLlama38BInstructV1 => "meta-llama3-8b-instruct-v1", + Model::MetaLlama370BInstructV1 => "meta-llama3-70b-instruct-v1", + Model::MetaLlama318BInstructV1_128k => "meta-llama3-1-8b-instruct-v1-128k", + Model::MetaLlama318BInstructV1 => "meta-llama3-1-8b-instruct-v1", + Model::MetaLlama3170BInstructV1_128k => "meta-llama3-1-70b-instruct-v1-128k", + Model::MetaLlama3170BInstructV1 => "meta-llama3-1-70b-instruct-v1", + Model::MetaLlama3211BInstructV1 => "meta-llama3-2-11b-instruct-v1", + Model::MetaLlama3290BInstructV1 => "meta-llama3-2-90b-instruct-v1", + Model::MetaLlama321BInstructV1 => "meta-llama3-2-1b-instruct-v1", + Model::MetaLlama323BInstructV1 => "meta-llama3-2-3b-instruct-v1", + Model::MistralMistral7BInstructV0 => "mistral-7b-instruct-v0", + Model::MistralMixtral8x7BInstructV0 => "mistral-mixtral-8x7b-instruct-v0", + Model::MistralMistralLarge2402V1 => "mistral-large-2402-v1", + Model::MistralMistralSmall2402V1 => "mistral-small-2402-v1", + Model::MistralPixtralLarge2502V1 => "mistral-pixtral-large-2502-v1", + Model::PalmyraWriterX4 => "palmyra-writer-x4", + Model::PalmyraWriterX5 => "palmyra-writer-x5", + Self::Custom { name, .. } => name, + } + } + + pub fn request_id(&self) -> &str { match self { Model::ClaudeSonnet4 | Model::ClaudeSonnet4Thinking => { "anthropic.claude-sonnet-4-20250514-v1:0" @@ -364,11 +418,11 @@ impl Model { anyhow::bail!("Unsupported Region {region}"); }; - let model_id = self.id(); + let model_id = self.request_id(); match (self, region_group) { // Custom models can't have CRI IDs - (Model::Custom { .. }, _) => Ok(self.id().into()), + (Model::Custom { .. }, _) => Ok(self.request_id().into()), // Models with US Gov only (Model::Claude3_5Sonnet, "us-gov") | (Model::Claude3Haiku, "us-gov") => { @@ -431,7 +485,7 @@ impl Model { ) => Ok(format!("{}.{}", region_group, model_id)), // Any other combination is not supported - _ => Ok(self.id().into()), + _ => Ok(self.request_id().into()), } } } @@ -586,4 +640,39 @@ mod tests { Ok(()) } + + #[test] + fn test_friendly_id_vs_request_id() { + // Test that id() returns friendly identifiers + assert_eq!(Model::Claude3_5SonnetV2.id(), "claude-3-5-sonnet-v2"); + assert_eq!(Model::AmazonNovaLite.id(), "amazon-nova-lite"); + assert_eq!(Model::DeepSeekR1.id(), "deepseek-r1"); + assert_eq!( + Model::MetaLlama38BInstructV1.id(), + "meta-llama3-8b-instruct-v1" + ); + + // Test that request_id() returns actual backend model IDs + assert_eq!( + Model::Claude3_5SonnetV2.request_id(), + "anthropic.claude-3-5-sonnet-20241022-v2:0" + ); + assert_eq!(Model::AmazonNovaLite.request_id(), "amazon.nova-lite-v1:0"); + assert_eq!(Model::DeepSeekR1.request_id(), "us.deepseek.r1-v1:0"); + assert_eq!( + Model::MetaLlama38BInstructV1.request_id(), + "meta.llama3-8b-instruct-v1:0" + ); + + // Test thinking models have different friendly IDs but same request IDs + assert_eq!(Model::ClaudeSonnet4.id(), "claude-4-sonnet"); + assert_eq!( + Model::ClaudeSonnet4Thinking.id(), + "claude-4-sonnet-thinking" + ); + assert_eq!( + Model::ClaudeSonnet4.request_id(), + Model::ClaudeSonnet4Thinking.request_id() + ); + } }