From 8d7922644553e0cf79298a43c2663b32d146bc53 Mon Sep 17 00:00:00 2001
From: Shardul Vaidya <31039336+5herlocked@users.noreply.github.com>
Date: Mon, 12 May 2025 05:13:37 -0400
Subject: [PATCH] bedrock: Add support for Mistral - Pixtral Large (#28274)

Release Notes:

- AWS Bedrock: Added support for Pixtral Large 25.02 v1

---------

Co-authored-by: Peter Tripp <peter@zed.dev>
Co-authored-by: Marshall Bowers <git@maxdeviant.com>
---
 crates/bedrock/src/models.rs | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/crates/bedrock/src/models.rs b/crates/bedrock/src/models.rs
index 6745b49cd58f9907446882980d39282bbca4dda9..cf42b4ff24216acfb3ab9185d25b1536f4cebeb0 100644
--- a/crates/bedrock/src/models.rs
+++ b/crates/bedrock/src/models.rs
@@ -73,6 +73,7 @@ pub enum Model {
     MistralMixtral8x7BInstructV0,
     MistralMistralLarge2402V1,
     MistralMistralSmall2402V1,
+    MistralPixtralLarge2502V1,
     // Writer models
     PalmyraWriterX5,
     PalmyraWriterX4,
@@ -154,6 +155,7 @@ impl Model {
             Model::MistralMixtral8x7BInstructV0 => "mistral.mixtral-8x7b-instruct-v0:1",
             Model::MistralMistralLarge2402V1 => "mistral.mistral-large-2402-v1:0",
             Model::MistralMistralSmall2402V1 => "mistral.mistral-small-2402-v1:0",
+            Model::MistralPixtralLarge2502V1 => "mistral.pixtral-large-2502-v1:0",
             Model::PalmyraWriterX4 => "writer.palmyra-x4-v1:0",
             Model::PalmyraWriterX5 => "writer.palmyra-x5-v1:0",
             Self::Custom { name, .. } => name,
@@ -203,6 +205,7 @@ impl Model {
             Self::MistralMixtral8x7BInstructV0 => "Mistral Mixtral 8x7B Instruct V0",
             Self::MistralMistralLarge2402V1 => "Mistral Large 2402 V1",
             Self::MistralMistralSmall2402V1 => "Mistral Small 2402 V1",
+            Self::MistralPixtralLarge2502V1 => "Pixtral Large 25.02 V1",
             Self::PalmyraWriterX5 => "Writer Palmyra X5",
             Self::PalmyraWriterX4 => "Writer Palmyra X4",
             Self::Custom {
@@ -222,7 +225,7 @@ impl Model {
             Self::PalmyraWriterX5 => 1_000_000,
             Self::PalmyraWriterX4 => 128_000,
             Self::Custom { max_tokens, .. } => *max_tokens,
-            _ => 200_000,
+            _ => 128_000,
         }
     }
 
@@ -324,7 +327,10 @@ impl Model {
             | (Model::Claude3_5Sonnet, "us")
             | (Model::Claude3_7Sonnet, "us")
             | (Model::Claude3_7SonnetThinking, "us")
-            | (Model::AmazonNovaPremier, "us") => Ok(format!("{}.{}", region_group, model_id)),
+            | (Model::AmazonNovaPremier, "us")
+            | (Model::MistralPixtralLarge2502V1, "us") => {
+                Ok(format!("{}.{}", region_group, model_id))
+            }
 
             // Models available in US, EU, and APAC
             (Model::Claude3_5SonnetV2, "us")