bedrock: Fix cross-region inference (#30659)

Shardul Vaidya and Marshall Bowers created 6 months ago

Closes #30535

Release Notes:

- AWS Bedrock: Add support for Meta Llama 4 Scout and Maverick models.
- AWS Bedrock: Fixed cross-region inference for all regions.
- AWS Bedrock: Updated all models available through Cross Region
inference.

---------

Co-authored-by: Marshall Bowers <[e-mail]>

Change summary

crates/bedrock/src/models.rs                   | 181 +++++++++++--------
crates/language_models/src/provider/bedrock.rs |  12 
2 files changed, 107 insertions(+), 86 deletions(-)

Detailed changes

crates/bedrock/src/models.rs

@@ -71,16 +71,20 @@ pub enum Model {
     // DeepSeek
     DeepSeekR1,
     // Meta models
-    MetaLlama38BInstructV1,
-    MetaLlama370BInstructV1,
-    MetaLlama318BInstructV1_128k,
-    MetaLlama318BInstructV1,
-    MetaLlama3170BInstructV1_128k,
-    MetaLlama3170BInstructV1,
-    MetaLlama3211BInstructV1,
-    MetaLlama3290BInstructV1,
-    MetaLlama321BInstructV1,
-    MetaLlama323BInstructV1,
+    MetaLlama3_8BInstruct,
+    MetaLlama3_70BInstruct,
+    MetaLlama31_8BInstruct,
+    MetaLlama31_70BInstruct,
+    MetaLlama31_405BInstruct,
+    MetaLlama32_1BInstruct,
+    MetaLlama32_3BInstruct,
+    MetaLlama32_11BMultiModal,
+    MetaLlama32_90BMultiModal,
+    MetaLlama33_70BInstruct,
+    #[allow(non_camel_case_types)]
+    MetaLlama4Scout_17BInstruct,
+    #[allow(non_camel_case_types)]
+    MetaLlama4Maverick_17BInstruct,
     // Mistral models
     MistralMistral7BInstructV0,
     MistralMixtral8x7BInstructV0,
@@ -145,7 +149,7 @@ impl Model {
             Model::AmazonNovaMicro => "amazon.nova-micro-v1:0",
             Model::AmazonNovaPro => "amazon.nova-pro-v1:0",
             Model::AmazonNovaPremier => "amazon.nova-premier-v1:0",
-            Model::DeepSeekR1 => "us.deepseek.r1-v1:0",
+            Model::DeepSeekR1 => "deepseek.r1-v1:0",
             Model::AI21J2GrandeInstruct => "ai21.j2-grande-instruct",
             Model::AI21J2JumboInstruct => "ai21.j2-jumbo-instruct",
             Model::AI21J2Mid => "ai21.j2-mid",
@@ -160,16 +164,18 @@ impl Model {
             Model::CohereCommandRV1 => "cohere.command-r-v1:0",
             Model::CohereCommandRPlusV1 => "cohere.command-r-plus-v1:0",
             Model::CohereCommandLightTextV14_4k => "cohere.command-light-text-v14:7:4k",
-            Model::MetaLlama38BInstructV1 => "meta.llama3-8b-instruct-v1:0",
-            Model::MetaLlama370BInstructV1 => "meta.llama3-70b-instruct-v1:0",
-            Model::MetaLlama318BInstructV1_128k => "meta.llama3-1-8b-instruct-v1:0:128k",
-            Model::MetaLlama318BInstructV1 => "meta.llama3-1-8b-instruct-v1:0",
-            Model::MetaLlama3170BInstructV1_128k => "meta.llama3-1-70b-instruct-v1:0:128k",
-            Model::MetaLlama3170BInstructV1 => "meta.llama3-1-70b-instruct-v1:0",
-            Model::MetaLlama3211BInstructV1 => "meta.llama3-2-11b-instruct-v1:0",
-            Model::MetaLlama3290BInstructV1 => "meta.llama3-2-90b-instruct-v1:0",
-            Model::MetaLlama321BInstructV1 => "meta.llama3-2-1b-instruct-v1:0",
-            Model::MetaLlama323BInstructV1 => "meta.llama3-2-3b-instruct-v1:0",
+            Model::MetaLlama3_8BInstruct => "meta.llama3-8b-instruct-v1:0",
+            Model::MetaLlama3_70BInstruct => "meta.llama3-70b-instruct-v1:0",
+            Model::MetaLlama31_8BInstruct => "meta.llama3-1-8b-instruct-v1:0",
+            Model::MetaLlama31_70BInstruct => "meta.llama3-1-70b-instruct-v1:0",
+            Model::MetaLlama31_405BInstruct => "meta.llama3-1-405b-instruct-v1:0",
+            Model::MetaLlama32_11BMultiModal => "meta.llama3-2-11b-instruct-v1:0",
+            Model::MetaLlama32_90BMultiModal => "meta.llama3-2-90b-instruct-v1:0",
+            Model::MetaLlama32_1BInstruct => "meta.llama3-2-1b-instruct-v1:0",
+            Model::MetaLlama32_3BInstruct => "meta.llama3-2-3b-instruct-v1:0",
+            Model::MetaLlama33_70BInstruct => "meta.llama3-3-70b-instruct-v1:0",
+            Model::MetaLlama4Scout_17BInstruct => "meta.llama4-scout-17b-instruct-v1:0",
+            Model::MetaLlama4Maverick_17BInstruct => "meta.llama4-maverick-17b-instruct-v1:0",
             Model::MistralMistral7BInstructV0 => "mistral.mistral-7b-instruct-v0:2",
             Model::MistralMixtral8x7BInstructV0 => "mistral.mixtral-8x7b-instruct-v0:1",
             Model::MistralMistralLarge2402V1 => "mistral.mistral-large-2402-v1:0",
@@ -214,16 +220,18 @@ impl Model {
             Self::CohereCommandRV1 => "Cohere Command R V1",
             Self::CohereCommandRPlusV1 => "Cohere Command R Plus V1",
             Self::CohereCommandLightTextV14_4k => "Cohere Command Light Text V14 4K",
-            Self::MetaLlama38BInstructV1 => "Meta Llama 3 8B Instruct V1",
-            Self::MetaLlama370BInstructV1 => "Meta Llama 3 70B Instruct V1",
-            Self::MetaLlama318BInstructV1_128k => "Meta Llama 3 1.8B Instruct V1 128K",
-            Self::MetaLlama318BInstructV1 => "Meta Llama 3 1.8B Instruct V1",
-            Self::MetaLlama3170BInstructV1_128k => "Meta Llama 3 1 70B Instruct V1 128K",
-            Self::MetaLlama3170BInstructV1 => "Meta Llama 3 1 70B Instruct V1",
-            Self::MetaLlama3211BInstructV1 => "Meta Llama 3 2 11B Instruct V1",
-            Self::MetaLlama3290BInstructV1 => "Meta Llama 3 2 90B Instruct V1",
-            Self::MetaLlama321BInstructV1 => "Meta Llama 3 2 1B Instruct V1",
-            Self::MetaLlama323BInstructV1 => "Meta Llama 3 2 3B Instruct V1",
+            Self::MetaLlama3_8BInstruct => "Meta Llama 3 8B Instruct",
+            Self::MetaLlama3_70BInstruct => "Meta Llama 3 70B Instruct",
+            Self::MetaLlama31_8BInstruct => "Meta Llama 3.1 8B Instruct",
+            Self::MetaLlama31_70BInstruct => "Meta Llama 3.1 70B Instruct",
+            Self::MetaLlama31_405BInstruct => "Meta Llama 3.1 405B Instruct",
+            Self::MetaLlama32_11BMultiModal => "Meta Llama 3.2 11B Vision Instruct",
+            Self::MetaLlama32_90BMultiModal => "Meta Llama 3.2 90B Vision Instruct",
+            Self::MetaLlama32_1BInstruct => "Meta Llama 3.2 1B Instruct",
+            Self::MetaLlama32_3BInstruct => "Meta Llama 3.2 3B Instruct",
+            Self::MetaLlama33_70BInstruct => "Meta Llama 3.3 70B Instruct",
+            Self::MetaLlama4Scout_17BInstruct => "Meta Llama 4 Scout 17B Instruct",
+            Self::MetaLlama4Maverick_17BInstruct => "Meta Llama 4 Maverick 17B Instruct",
             Self::MistralMistral7BInstructV0 => "Mistral 7B Instruct V0",
             Self::MistralMixtral8x7BInstructV0 => "Mistral Mixtral 8x7B Instruct V0",
             Self::MistralMistralLarge2402V1 => "Mistral Large 2402 V1",
@@ -365,55 +373,60 @@ impl Model {
                 Ok(format!("{}.{}", region_group, model_id))
             }
 
-            // Models available only in US
-            (Model::Claude3Opus, "us")
-            | (Model::Claude3_5Haiku, "us")
-            | (Model::Claude3_7Sonnet, "us")
-            | (Model::ClaudeSonnet4, "us")
-            | (Model::ClaudeOpus4, "us")
-            | (Model::ClaudeSonnet4Thinking, "us")
-            | (Model::ClaudeOpus4Thinking, "us")
-            | (Model::Claude3_7SonnetThinking, "us")
-            | (Model::AmazonNovaPremier, "us")
-            | (Model::MistralPixtralLarge2502V1, "us") => {
+            // Available everywhere
+            (Model::AmazonNovaLite | Model::AmazonNovaMicro | Model::AmazonNovaPro, _) => {
                 Ok(format!("{}.{}", region_group, model_id))
             }
 
-            // Models available in US, EU, and APAC
-            (Model::Claude3_5SonnetV2, "us")
-            | (Model::Claude3_5SonnetV2, "apac")
-            | (Model::Claude3_5Sonnet, _)
-            | (Model::Claude3Haiku, _)
-            | (Model::Claude3Sonnet, _)
-            | (Model::AmazonNovaLite, _)
-            | (Model::AmazonNovaMicro, _)
-            | (Model::AmazonNovaPro, _) => Ok(format!("{}.{}", region_group, model_id)),
-
-            // Models with limited EU availability
-            (Model::MetaLlama321BInstructV1, "us")
-            | (Model::MetaLlama321BInstructV1, "eu")
-            | (Model::MetaLlama323BInstructV1, "us")
-            | (Model::MetaLlama323BInstructV1, "eu") => {
-                Ok(format!("{}.{}", region_group, model_id))
-            }
-
-            // US-only models (all remaining Meta models)
-            (Model::MetaLlama38BInstructV1, "us")
-            | (Model::MetaLlama370BInstructV1, "us")
-            | (Model::MetaLlama318BInstructV1, "us")
-            | (Model::MetaLlama318BInstructV1_128k, "us")
-            | (Model::MetaLlama3170BInstructV1, "us")
-            | (Model::MetaLlama3170BInstructV1_128k, "us")
-            | (Model::MetaLlama3211BInstructV1, "us")
-            | (Model::MetaLlama3290BInstructV1, "us") => {
-                Ok(format!("{}.{}", region_group, model_id))
-            }
-
-            // Writer models only available in the US
-            (Model::PalmyraWriterX4, "us") | (Model::PalmyraWriterX5, "us") => {
-                // They have some goofiness
-                Ok(format!("{}.{}", region_group, model_id))
-            }
+            // Models in US
+            (
+                Model::AmazonNovaPremier
+                | Model::Claude3_5Haiku
+                | Model::Claude3_5Sonnet
+                | Model::Claude3_5SonnetV2
+                | Model::Claude3_7Sonnet
+                | Model::Claude3_7SonnetThinking
+                | Model::Claude3Haiku
+                | Model::Claude3Opus
+                | Model::Claude3Sonnet
+                | Model::DeepSeekR1
+                | Model::MetaLlama31_405BInstruct
+                | Model::MetaLlama31_70BInstruct
+                | Model::MetaLlama31_8BInstruct
+                | Model::MetaLlama32_11BMultiModal
+                | Model::MetaLlama32_1BInstruct
+                | Model::MetaLlama32_3BInstruct
+                | Model::MetaLlama32_90BMultiModal
+                | Model::MetaLlama33_70BInstruct
+                | Model::MetaLlama4Maverick_17BInstruct
+                | Model::MetaLlama4Scout_17BInstruct
+                | Model::MistralPixtralLarge2502V1
+                | Model::PalmyraWriterX4
+                | Model::PalmyraWriterX5,
+                "us",
+            ) => Ok(format!("{}.{}", region_group, model_id)),
+
+            // Models available in EU
+            (
+                Model::Claude3_5Sonnet
+                | Model::Claude3_7Sonnet
+                | Model::Claude3_7SonnetThinking
+                | Model::Claude3Haiku
+                | Model::Claude3Sonnet
+                | Model::MetaLlama32_1BInstruct
+                | Model::MetaLlama32_3BInstruct
+                | Model::MistralPixtralLarge2502V1,
+                "eu",
+            ) => Ok(format!("{}.{}", region_group, model_id)),
+
+            // Models available in APAC
+            (
+                Model::Claude3_5Sonnet
+                | Model::Claude3_5SonnetV2
+                | Model::Claude3Haiku
+                | Model::Claude3Sonnet,
+                "apac",
+            ) => Ok(format!("{}.{}", region_group, model_id)),
 
             // Any other combination is not supported
             _ => Ok(self.id().into()),
@@ -464,6 +477,10 @@ mod tests {
             Model::Claude3_5SonnetV2.cross_region_inference_id("ap-northeast-1")?,
             "apac.anthropic.claude-3-5-sonnet-20241022-v2:0"
         );
+        assert_eq!(
+            Model::Claude3_5SonnetV2.cross_region_inference_id("ap-southeast-2")?,
+            "apac.anthropic.claude-3-5-sonnet-20241022-v2:0"
+        );
         assert_eq!(
             Model::AmazonNovaLite.cross_region_inference_id("ap-south-1")?,
             "apac.amazon.nova-lite-v1:0"
@@ -489,11 +506,15 @@ mod tests {
     fn test_meta_models_inference_ids() -> anyhow::Result<()> {
         // Test Meta models
         assert_eq!(
-            Model::MetaLlama370BInstructV1.cross_region_inference_id("us-east-1")?,
-            "us.meta.llama3-70b-instruct-v1:0"
+            Model::MetaLlama3_70BInstruct.cross_region_inference_id("us-east-1")?,
+            "meta.llama3-70b-instruct-v1:0"
+        );
+        assert_eq!(
+            Model::MetaLlama31_70BInstruct.cross_region_inference_id("us-east-1")?,
+            "us.meta.llama3-1-70b-instruct-v1:0"
         );
         assert_eq!(
-            Model::MetaLlama321BInstructV1.cross_region_inference_id("eu-west-1")?,
+            Model::MetaLlama32_1BInstruct.cross_region_inference_id("eu-west-1")?,
             "eu.meta.llama3-2-1b-instruct-v1:0"
         );
         Ok(())

crates/language_models/src/provider/bedrock.rs

@@ -531,13 +531,13 @@ impl LanguageModel for BedrockModel {
     > {
         let Ok(region) = cx.read_entity(&self.state, |state, _cx| {
             // Get region - from credentials or directly from settings
-            let region = state
-                .credentials
-                .as_ref()
-                .map(|s| s.region.clone())
-                .unwrap_or(String::from("us-east-1"));
+            let credentials_region = state.credentials.as_ref().map(|s| s.region.clone());
+            let settings_region = state.settings.as_ref().and_then(|s| s.region.clone());
 
-            region
+            // Use credentials region if available, otherwise use settings region, finally fall back to default
+            credentials_region
+                .or(settings_region)
+                .unwrap_or(String::from("us-east-1"))
         }) else {
             return async move {
                 anyhow::bail!("App State Dropped");