diff --git a/crates/bedrock/src/models.rs b/crates/bedrock/src/models.rs index 8ead77f9c47e7091da4120f8aed3fe8c796247ad..e7ff1f46d5d8e6a04066ae0953502659b680df9d 100644 --- a/crates/bedrock/src/models.rs +++ b/crates/bedrock/src/models.rs @@ -72,6 +72,9 @@ pub enum Model { MistralMixtral8x7BInstructV0, MistralMistralLarge2402V1, MistralMistralSmall2402V1, + // Writer models + PalmyraWriterX5, + PalmyraWriterX4, #[serde(rename = "custom")] Custom { name: String, @@ -149,6 +152,8 @@ 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::PalmyraWriterX4 => "writer.palmyra-x4-v1:0", + Model::PalmyraWriterX5 => "writer.palmyra-x5-v1:0", Self::Custom { name, .. } => name, } } @@ -195,6 +200,8 @@ impl Model { Self::MistralMixtral8x7BInstructV0 => "Mistral Mixtral 8x7B Instruct V0", Self::MistralMistralLarge2402V1 => "Mistral Large 2402 V1", Self::MistralMistralSmall2402V1 => "Mistral Small 2402 V1", + Self::PalmyraWriterX5 => "Writer Palmyra X5", + Self::PalmyraWriterX4 => "Writer Palmyra X4", Self::Custom { display_name, name, .. } => display_name.as_deref().unwrap_or(name), @@ -208,6 +215,8 @@ impl Model { | Self::Claude3Sonnet | Self::Claude3_5Haiku | Self::Claude3_7Sonnet => 200_000, + Self::PalmyraWriterX5 => 1_000_000, + Self::PalmyraWriterX4 => 128_000, Self::Custom { max_tokens, .. } => *max_tokens, _ => 200_000, } @@ -217,7 +226,7 @@ impl Model { match self { Self::Claude3Opus | Self::Claude3Sonnet | Self::Claude3_5Haiku => 4_096, Self::Claude3_7Sonnet | Self::Claude3_7SonnetThinking => 128_000, - Self::Claude3_5SonnetV2 => 8_192, + Self::Claude3_5SonnetV2 | Self::PalmyraWriterX4 | Self::PalmyraWriterX5 => 8_192, Self::Custom { max_output_tokens, .. } => max_output_tokens.unwrap_or(4_096), @@ -340,6 +349,12 @@ impl Model { 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)) + } + // Any other combination is not supported _ => Ok(self.id().into()), }