diff --git a/.codex b/.codex new file mode 100644 index 00000000..e69de29b diff --git a/src/backends/aws/mod.rs b/src/backends/aws/mod.rs index 0af04000..82eeae0d 100644 --- a/src/backends/aws/mod.rs +++ b/src/backends/aws/mod.rs @@ -10,9 +10,10 @@ use async_trait::async_trait; use aws_config::BehaviorVersion; use aws_sdk_bedrockruntime::{ types::{ - CachePointBlock, CachePointType, ContentBlock, ContentBlockDelta, ContentBlockStart, ConversationRole, ConverseStreamOutput, - Message, SystemContentBlock, Tool, ToolConfiguration, ToolInputSchema, ToolResultBlock, - ToolResultContentBlock, ToolUseBlock, + CachePointBlock, CachePointType, ContentBlock, ContentBlockDelta, ContentBlockStart, + ConversationRole, ConverseStreamOutput, JsonSchemaDefinition, Message, OutputConfig, + OutputFormat, OutputFormatStructure, OutputFormatType, SystemContentBlock, Tool, + ToolConfiguration, ToolInputSchema, ToolResultBlock, ToolResultContentBlock, ToolUseBlock, }, Client as BedrockClient, }; @@ -27,7 +28,7 @@ use std::sync::Arc; use tokio::sync::OnceCell; use crate::chat::{ - ChatMessage as LlmChatMessage, ChatProvider, StreamChunk as LlmStreamChunk, StreamChoice, + ChatMessage as LlmChatMessage, ChatProvider, StreamChoice, StreamChunk as LlmStreamChunk, StreamDelta, StreamResponse, StructuredOutputFormat, Tool as LlmTool, ToolChoice as LlmToolChoice, }; @@ -39,7 +40,7 @@ use crate::embedding::EmbeddingProvider; use crate::models::ModelsProvider; use crate::stt::SpeechToTextProvider; use crate::tts::TextToSpeechProvider; -use crate::{FunctionCall, ToolCall, LLMProvider}; +use crate::{FunctionCall, LLMProvider, ToolCall}; mod error; mod models; @@ -81,6 +82,7 @@ struct PreparedChatRequest { system: Option, tool_config: Option, inference_config: aws_sdk_bedrockruntime::types::InferenceConfiguration, + output_config: Option, } #[derive(Debug, Default, Clone)] @@ -186,7 +188,7 @@ impl BedrockBackend { Ok(Self { client: Arc::new(OnceCell::new()), region, - model: model.map(BedrockModel::Custom), + model: model.map(BedrockModel::from_id), max_tokens, temperature, timeout_seconds, @@ -268,7 +270,6 @@ impl BedrockBackend { .set_max_tokens(request.max_tokens.or(self.max_tokens).map(|t| t as i32)) .set_temperature(request.temperature.map(|t| t as f32).or(self.temperature)) .set_top_p(request.top_p.map(|p| p as f32).or(self.top_p)) - .set_stop_sequences(request.stop_sequences) .build(), ); @@ -329,6 +330,7 @@ impl BedrockBackend { system, tool_config, inference_config, + output_config, } = self.prepare_chat_request(request)?; let mut converse_request = client @@ -346,6 +348,12 @@ impl BedrockBackend { converse_request = converse_request.tool_config(tool_config); } + // Apply native structured output config for models that support outputConfig.textFormat + // (e.g. Nova). For other models this will be None and the tool-based path is used instead. + if let Some(output_config) = output_config { + converse_request = converse_request.output_config(output_config); + } + // Add inference configuration converse_request = converse_request.inference_config(inference_config); @@ -498,6 +506,7 @@ impl BedrockBackend { system, tool_config, inference_config, + output_config, } = self.prepare_chat_request(request)?; let mut converse_request = client @@ -515,6 +524,11 @@ impl BedrockBackend { converse_request = converse_request.tool_config(tool_config); } + // Apply native structured output config for models that support outputConfig.textFormat + if let Some(output_config) = output_config { + converse_request = converse_request.output_config(output_config); + } + // Add inference configuration converse_request = converse_request.inference_config(inference_config); @@ -590,6 +604,7 @@ impl BedrockBackend { system, tool_config, inference_config, + output_config, } = self.prepare_chat_request(request)?; let mut converse_request = client @@ -605,6 +620,11 @@ impl BedrockBackend { converse_request = converse_request.tool_config(tool_config); } + // Apply native structured output config for models that support outputConfig.textFormat + if let Some(output_config) = output_config { + converse_request = converse_request.output_config(output_config); + } + converse_request = converse_request.inference_config(inference_config); let response = converse_request @@ -669,8 +689,7 @@ impl BedrockBackend { _ => {} }, ConverseStreamOutput::ContentBlockStop(stop) => { - let index = - usize::try_from(stop.content_block_index).unwrap_or(0); + let index = usize::try_from(stop.content_block_index).unwrap_or(0); if let Some(state) = tool_states.remove(&index) { if state.started { pending.push_back(LlmStreamChunk::ToolUseComplete { @@ -798,6 +817,8 @@ impl BedrockBackend { } let mut tool_choice = self.tool_choice.clone(); + let mut output_config: Option = None; + if let Some(response_format) = self.json_schema.as_ref() { let schema = response_format.schema.clone().ok_or_else(|| { BedrockError::InvalidRequest( @@ -805,21 +826,63 @@ impl BedrockBackend { ) })?; - let input_schema = ToolInputSchema::Json(Self::value_to_document(&schema)); + if self.model_supports(&model_id, ModelCapability::NativeStructuredOutput) { + // Nova and other models that advertise NativeStructuredOutput receive the schema + // via outputConfig.textFormat. The response arrives as ContentBlock::Text + // (plain JSON), so no tool-call unwrapping is needed. + let schema_str = serde_json::to_string(&schema).map_err(|e| { + BedrockError::InvalidRequest(format!("Failed to serialize schema: {}", e)) + })?; - let tool_spec = aws_sdk_bedrockruntime::types::ToolSpecification::builder() - .name("json_schema_tool") - .description( - "Generates structured output in JSON format according to the provided schema.", - ) - .input_schema(input_schema) - .build() - .map_err(|e| { - BedrockError::InvalidRequest(format!("Failed to build tool spec: {:?}", e)) + let mut json_schema_builder = JsonSchemaDefinition::builder().schema(schema_str); + if !response_format.name.is_empty() { + json_schema_builder = json_schema_builder.name(&response_format.name); + } + if let Some(desc) = &response_format.description { + json_schema_builder = json_schema_builder.description(desc); + } + + let json_schema_def = json_schema_builder.build().map_err(|e| { + BedrockError::InvalidRequest(format!( + "Failed to build JSON schema definition: {:?}", + e + )) })?; - bedrock_tools.push(Tool::ToolSpec(tool_spec)); - tool_choice = Some(LlmToolChoice::Tool("json_schema_tool".to_string())); + let output_format = OutputFormat::builder() + .r#type(OutputFormatType::JsonSchema) + .structure(OutputFormatStructure::JsonSchema(json_schema_def)) + .build() + .map_err(|e| { + BedrockError::InvalidRequest(format!( + "Failed to build output format: {:?}", + e + )) + })?; + + // OutputConfig::builder().build() is infallible - the SDK builder has no + // required fields beyond what we set via .text_format(). + output_config = Some(OutputConfig::builder().text_format(output_format).build()); + } else { + // Fallback for models without native structured output (e.g. Claude): + // inject a synthetic tool and force the model to call it, then unwrap + // the tool call in convert_chat_response. + let input_schema = ToolInputSchema::Json(Self::value_to_document(&schema)); + + let tool_spec = aws_sdk_bedrockruntime::types::ToolSpecification::builder() + .name("json_schema_tool") + .description( + "Generates structured output in JSON format according to the provided schema.", + ) + .input_schema(input_schema) + .build() + .map_err(|e| { + BedrockError::InvalidRequest(format!("Failed to build tool spec: {:?}", e)) + })?; + + bedrock_tools.push(Tool::ToolSpec(tool_spec)); + tool_choice = Some(LlmToolChoice::Tool("json_schema_tool".to_string())); + } } if let Some(tools) = &self.tools { @@ -905,6 +968,7 @@ impl BedrockBackend { system, tool_config, inference_config, + output_config, }) } @@ -1384,7 +1448,12 @@ impl ChatProvider for BedrockBackend { &self, messages: &[LlmChatMessage], ) -> std::result::Result< - Pin> + Send>>, + Pin< + Box< + dyn Stream> + + Send, + >, + >, crate::error::LLMError, > { let aws_messages: Vec = messages @@ -1434,17 +1503,15 @@ impl ChatProvider for BedrockBackend { }], usage: None, })), - Ok(LlmStreamChunk::ToolUseComplete { tool_call, .. }) => { - Some(Ok(StreamResponse { - choices: vec![StreamChoice { - delta: StreamDelta { - content: None, - tool_calls: Some(vec![tool_call]), - }, - }], - usage: None, - })) - } + Ok(LlmStreamChunk::ToolUseComplete { tool_call, .. }) => Some(Ok(StreamResponse { + choices: vec![StreamChoice { + delta: StreamDelta { + content: None, + tool_calls: Some(vec![tool_call]), + }, + }], + usage: None, + })), Ok(LlmStreamChunk::Done { .. }) => None, Ok(_) => None, Err(e) => Some(Err(crate::error::LLMError::ProviderError(e.to_string()))), @@ -1459,7 +1526,12 @@ impl ChatProvider for BedrockBackend { messages: &[LlmChatMessage], tools: Option<&[LlmTool]>, ) -> std::result::Result< - Pin> + Send>>, + Pin< + Box< + dyn Stream> + + Send, + >, + >, crate::error::LLMError, > { let aws_messages: Vec = messages @@ -1741,4 +1813,123 @@ streaming = true assert_eq!(tool_call.function.arguments, "{}"); } + + #[test] + fn test_prepare_chat_request_nova_uses_output_config_not_tool() { + // Nova models must use outputConfig.textFormat for structured output, + // not the synthetic json_schema_tool workaround. + let backend = BedrockBackend::new( + "us-east-1".to_string(), + Some("amazon.nova-pro-v1:0".to_string()), + None, + None, + None, + None, + None, + None, + None, + None, + None, + Some(crate::chat::StructuredOutputFormat { + name: "result".to_string(), + description: None, + schema: Some(serde_json::json!({"type": "object", "properties": {}})), + strict: None, + }), + ) + .unwrap(); + + let request = ChatRequest::new(vec![ChatMessage::user("hello")]); + let prepared = backend.prepare_chat_request(request).unwrap(); + + // output_config must be set for the native path + assert!(prepared.output_config.is_some()); + // tool_config must be None -- no synthetic json_schema_tool should be injected + assert!(prepared.tool_config.is_none()); + } + + #[test] + fn test_prepare_chat_request_claude_uses_tool_not_output_config() { + // Claude (and other non-Nova models) must use the tool-based workaround. + let backend = BedrockBackend::new( + "us-east-1".to_string(), + Some("us.anthropic.claude-sonnet-4-0-v1:0".to_string()), + None, + None, + None, + None, + None, + None, + None, + None, + None, + Some(crate::chat::StructuredOutputFormat { + name: "result".to_string(), + description: None, + schema: Some(serde_json::json!({"type": "object", "properties": {}})), + strict: None, + }), + ) + .unwrap(); + + let request = ChatRequest::new(vec![ChatMessage::user("hello")]); + let prepared = backend.prepare_chat_request(request).unwrap(); + + // tool_config must contain the synthetic json_schema_tool + let tool_config = prepared + .tool_config + .expect("tool_config should be present for Claude"); + assert!(tool_config + .tools() + .iter() + .any(|t| matches!(t, Tool::ToolSpec(spec) if spec.name() == "json_schema_tool"))); + // output_config must not be set + assert!(prepared.output_config.is_none()); + } + + #[test] + fn test_prepare_chat_request_nova_with_real_tools_and_schema() { + // When Nova has both real user tools AND json_schema set, real tools go into + // tool_config and the schema goes into output_config. Both can coexist because + // the synthetic json_schema_tool is never injected on the native path. + let backend = BedrockBackend::new( + "us-east-1".to_string(), + Some("amazon.nova-pro-v1:0".to_string()), + None, + None, + None, + None, + None, + None, + None, + None, + None, + Some(crate::chat::StructuredOutputFormat { + name: "result".to_string(), + description: None, + schema: Some(serde_json::json!({"type": "object", "properties": {}})), + strict: None, + }), + ) + .unwrap(); + + let tools = vec![ToolDefinition { + name: "get_weather".to_string(), + description: "Get weather".to_string(), + input_schema: serde_json::json!({"type": "object", "properties": {}}), + cache_control: None, + }]; + + let request = ChatRequest::new(vec![ChatMessage::user("hello")]).with_tools(tools); + let prepared = backend.prepare_chat_request(request).unwrap(); + + assert!(prepared.output_config.is_some()); + let tool_config = prepared + .tool_config + .expect("real tools should produce tool_config"); + let tools = tool_config.tools(); + // Only the real tool, no json_schema_tool pollution + assert_eq!(tools.len(), 1); + assert!(matches!(tools[0], Tool::ToolSpec(ref spec) if spec.name() == "get_weather")); + } } diff --git a/src/backends/aws/models.rs b/src/backends/aws/models.rs index 9eb95e48..d0fdfc65 100644 --- a/src/backends/aws/models.rs +++ b/src/backends/aws/models.rs @@ -79,6 +79,16 @@ pub enum DirectModel { #[serde(rename = "amazon.titan-embed-text-v1")] TitanEmbedV1, + // Amazon Nova models + #[serde(rename = "amazon.nova-pro-v1:0")] + NovaProV1, + + #[serde(rename = "amazon.nova-lite-v1:0")] + NovaLiteV1, + + #[serde(rename = "amazon.nova-micro-v1:0")] + NovaMicroV1, + // Cohere models #[serde(rename = "cohere.command-r-plus-v1:0")] CohereCommandRPlus, @@ -122,6 +132,16 @@ pub enum CrossRegionModel { #[serde(rename = "claude-3-haiku-20240307-v1:0")] ClaudeHaiku3, + // Amazon Nova models + #[serde(rename = "nova-pro-v1:0")] + NovaProV1, + + #[serde(rename = "nova-lite-v1:0")] + NovaLiteV1, + + #[serde(rename = "nova-micro-v1:0")] + NovaMicroV1, + // Mistral models #[serde(rename = "pixtral-large-2502-v1:0")] MistralPixtralLarge, @@ -212,11 +232,28 @@ impl BedrockModel { } } - // Try to match against direct models + // Try to match against direct models (exact match first) if let Some(direct) = DirectModel::from_id(&id) { return Self::Direct(direct); } + // Try stripping a geo inference-profile prefix and retrying. + // IDs like "us.amazon.nova-pro-v1:0", "eu.amazon.nova-lite-v1:0", or a future + // "sa.amazon.nova-pro-v1:0" are functionally identical to the unprefixed form + // for capability checks. We detect the prefix by looking for a short all-alpha + // segment before the first dot (geo codes are 2-3 chars; vendor names are longer). + if let Some(dot) = id.find('.') { + let prefix = &id[..dot]; + if !prefix.is_empty() + && prefix.len() <= 3 + && prefix.chars().all(|c| c.is_ascii_alphabetic()) + { + if let Some(direct) = DirectModel::from_id(&id[dot + 1..]) { + return Self::Direct(direct); + } + } + } + // Otherwise treat as custom Self::Custom(id) } @@ -305,6 +342,9 @@ impl DirectModel { Self::TitanTextLite => "amazon.titan-text-lite-v1", Self::TitanEmbedV2 => "amazon.titan-embed-text-v2:0", Self::TitanEmbedV1 => "amazon.titan-embed-text-v1", + Self::NovaProV1 => "amazon.nova-pro-v1:0", + Self::NovaLiteV1 => "amazon.nova-lite-v1:0", + Self::NovaMicroV1 => "amazon.nova-micro-v1:0", Self::CohereCommandRPlus => "cohere.command-r-plus-v1:0", Self::CohereCommandR => "cohere.command-r-v1:0", Self::CohereEmbedV3 => "cohere.embed-english-v3", @@ -333,6 +373,9 @@ impl DirectModel { "amazon.titan-text-lite-v1" => Some(Self::TitanTextLite), "amazon.titan-embed-text-v2:0" => Some(Self::TitanEmbedV2), "amazon.titan-embed-text-v1" => Some(Self::TitanEmbedV1), + "amazon.nova-pro-v1:0" => Some(Self::NovaProV1), + "amazon.nova-lite-v1:0" => Some(Self::NovaLiteV1), + "amazon.nova-micro-v1:0" => Some(Self::NovaMicroV1), "cohere.command-r-plus-v1:0" => Some(Self::CohereCommandRPlus), "cohere.command-r-v1:0" => Some(Self::CohereCommandR), "cohere.embed-english-v3" => Some(Self::CohereEmbedV3), @@ -353,6 +396,9 @@ impl CrossRegionModel { Self::ClaudeOpus3 => "claude-3-opus-20240229-v1:0", Self::ClaudeSonnet35 => "claude-3-5-sonnet-20240620-v1:0", Self::ClaudeHaiku3 => "claude-3-haiku-20240307-v1:0", + Self::NovaProV1 => "nova-pro-v1:0", + Self::NovaLiteV1 => "nova-lite-v1:0", + Self::NovaMicroV1 => "nova-micro-v1:0", Self::MistralPixtralLarge => "pixtral-large-2502-v1:0", Self::CohereEmbedV4 => "embed-v4:0", } @@ -366,6 +412,7 @@ impl CrossRegionModel { | Self::ClaudeOpus3 | Self::ClaudeSonnet35 | Self::ClaudeHaiku3 => "anthropic", + Self::NovaProV1 | Self::NovaLiteV1 | Self::NovaMicroV1 => "amazon", Self::MistralPixtralLarge => "mistral", Self::CohereEmbedV4 => "cohere", } @@ -387,6 +434,9 @@ impl CrossRegionModel { Some(Self::ClaudeSonnet35) } ("anthropic", id) if id.contains("claude-3-haiku") => Some(Self::ClaudeHaiku3), + ("amazon", id) if id.contains("nova-pro") => Some(Self::NovaProV1), + ("amazon", id) if id.contains("nova-lite") => Some(Self::NovaLiteV1), + ("amazon", id) if id.contains("nova-micro") => Some(Self::NovaMicroV1), ("mistral", id) if id.contains("pixtral-large") => Some(Self::MistralPixtralLarge), ("cohere", id) if id.contains("embed-v4") => Some(Self::CohereEmbedV4), _ => None, @@ -413,6 +463,9 @@ impl BedrockModel { ModelCapability::Vision => self.supports_vision_impl(), ModelCapability::ToolUse => self.supports_tools_impl(), ModelCapability::Streaming => self.is_text_model() || self.is_chat_model(), + ModelCapability::NativeStructuredOutput => { + self.supports_native_structured_output_impl() + } } } @@ -454,6 +507,14 @@ impl BedrockModel { } } + fn supports_native_structured_output_impl(&self) -> bool { + match self.inner_model() { + InnerModel::Direct(model) => model.supports_native_structured_output(), + InnerModel::CrossRegion(model) => model.supports_native_structured_output(), + _ => false, + } + } + /// Get the maximum tokens this model can handle in output pub fn max_output_tokens(&self) -> u32 { match self.inner_model() { @@ -505,6 +566,8 @@ impl DirectModel { | Self::ClaudeHaiku3 | Self::Llama32_90B | Self::Llama32_11B + | Self::NovaProV1 + | Self::NovaLiteV1 ) } @@ -520,9 +583,16 @@ impl DirectModel { | Self::CohereCommandRPlus | Self::CohereCommandR | Self::MistralLarge + | Self::NovaProV1 + | Self::NovaLiteV1 + | Self::NovaMicroV1 ) } + fn supports_native_structured_output(&self) -> bool { + matches!(self, Self::NovaProV1 | Self::NovaLiteV1 | Self::NovaMicroV1) + } + fn max_output_tokens(&self) -> u32 { match self { Self::ClaudeSonnet35V2 | Self::ClaudeSonnet4 => 8192, @@ -534,6 +604,7 @@ impl DirectModel { Self::TitanTextLite => 4096, Self::CohereCommandRPlus | Self::CohereCommandR => 4096, Self::MistralLarge | Self::MistralSmall => 8192, + Self::NovaProV1 | Self::NovaLiteV1 | Self::NovaMicroV1 => 5120, _ => 0, } } @@ -551,6 +622,8 @@ impl DirectModel { Self::TitanTextExpress | Self::TitanTextLite => 8_000, Self::CohereCommandRPlus | Self::CohereCommandR => 128_000, Self::MistralLarge | Self::MistralSmall => 128_000, + Self::NovaProV1 | Self::NovaLiteV1 => 300_000, + Self::NovaMicroV1 => 128_000, _ => 0, } } @@ -575,6 +648,8 @@ impl CrossRegionModel { | Self::ClaudeSonnet35 | Self::ClaudeHaiku3 | Self::MistralPixtralLarge + | Self::NovaProV1 + | Self::NovaLiteV1 ) } @@ -588,15 +663,23 @@ impl CrossRegionModel { | Self::ClaudeSonnet35 | Self::ClaudeHaiku3 | Self::MistralPixtralLarge + | Self::NovaProV1 + | Self::NovaLiteV1 + | Self::NovaMicroV1 ) } + fn supports_native_structured_output(&self) -> bool { + matches!(self, Self::NovaProV1 | Self::NovaLiteV1 | Self::NovaMicroV1) + } + fn max_output_tokens(&self) -> u32 { match self { Self::ClaudeSonnet4 | Self::ClaudeSonnet45 | Self::ClaudeSonnet35V2 => 8192, Self::ClaudeOpus3 | Self::ClaudeSonnet35 => 8192, Self::ClaudeHaiku3 => 4096, Self::MistralPixtralLarge => 8192, + Self::NovaProV1 | Self::NovaLiteV1 | Self::NovaMicroV1 => 5120, Self::CohereEmbedV4 => 0, } } @@ -606,6 +689,8 @@ impl CrossRegionModel { Self::ClaudeSonnet4 | Self::ClaudeSonnet45 | Self::ClaudeSonnet35V2 => 200_000, Self::ClaudeOpus3 | Self::ClaudeSonnet35 | Self::ClaudeHaiku3 => 200_000, Self::MistralPixtralLarge => 128_000, + Self::NovaProV1 | Self::NovaLiteV1 => 300_000, + Self::NovaMicroV1 => 128_000, Self::CohereEmbedV4 => 0, } } @@ -649,6 +734,11 @@ pub enum ModelCapability { /// Streaming responses Streaming, + + /// Native structured output via outputConfig.textFormat (Bedrock Converse). + /// Models with this capability receive JSON directly without the synthetic + /// json_schema_tool workaround needed by Claude and other models. + NativeStructuredOutput, } #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -741,6 +831,8 @@ pub struct ModelCapabilityOverride { pub tool_use: Option, #[serde(default)] pub streaming: Option, + #[serde(default)] + pub native_structured_output: Option, } impl ModelCapabilityOverride { @@ -752,6 +844,7 @@ impl ModelCapabilityOverride { ModelCapability::Vision => self.vision, ModelCapability::ToolUse => self.tool_use, ModelCapability::Streaming => self.streaming, + ModelCapability::NativeStructuredOutput => self.native_structured_output, } } } @@ -760,6 +853,167 @@ impl ModelCapabilityOverride { mod tests { use super::*; + #[test] + fn test_nova_model_ids() { + let nova_pro = BedrockModel::Direct(DirectModel::NovaProV1); + assert_eq!(nova_pro.model_id(), "amazon.nova-pro-v1:0"); + + let nova_lite = BedrockModel::Direct(DirectModel::NovaLiteV1); + assert_eq!(nova_lite.model_id(), "amazon.nova-lite-v1:0"); + + let nova_micro = BedrockModel::Direct(DirectModel::NovaMicroV1); + assert_eq!(nova_micro.model_id(), "amazon.nova-micro-v1:0"); + + let nova_pro_eu = BedrockModel::eu(CrossRegionModel::NovaProV1); + assert!(nova_pro_eu.model_id().contains("nova-pro-v1:0")); + assert!(nova_pro_eu.model_id().contains("amazon")); + } + + #[test] + fn test_nova_from_id() { + let model = BedrockModel::from_id("amazon.nova-pro-v1:0"); + assert!(matches!( + model, + BedrockModel::Direct(DirectModel::NovaProV1) + )); + + let model = BedrockModel::from_id("amazon.nova-lite-v1:0"); + assert!(matches!( + model, + BedrockModel::Direct(DirectModel::NovaLiteV1) + )); + + let model = BedrockModel::from_id("amazon.nova-micro-v1:0"); + assert!(matches!( + model, + BedrockModel::Direct(DirectModel::NovaMicroV1) + )); + } + + #[test] + fn test_nova_geo_prefixed_from_id() { + // Geo-prefixed inference profile IDs must resolve to the same Direct variant + // so that capability checks (NativeStructuredOutput, ToolUse, etc.) work correctly. + // Includes a hypothetical future prefix ("sa") to verify we don't rely on an + // enumerated allowlist. + for prefix in ["us.", "eu.", "ap.", "sa."] { + let id = format!("{}amazon.nova-pro-v1:0", prefix); + assert!( + matches!( + BedrockModel::from_id(&id), + BedrockModel::Direct(DirectModel::NovaProV1) + ), + "{id} should resolve to NovaProV1" + ); + + let id = format!("{}amazon.nova-lite-v1:0", prefix); + assert!( + matches!( + BedrockModel::from_id(&id), + BedrockModel::Direct(DirectModel::NovaLiteV1) + ), + "{id} should resolve to NovaLiteV1" + ); + + let id = format!("{}amazon.nova-micro-v1:0", prefix); + assert!( + matches!( + BedrockModel::from_id(&id), + BedrockModel::Direct(DirectModel::NovaMicroV1) + ), + "{id} should resolve to NovaMicroV1" + ); + } + } + + #[test] + fn test_nova_geo_prefixed_native_structured_output_capability() { + // The whole point of the fix: geo-prefixed Nova IDs must report + // NativeStructuredOutput = true so prepare_chat_request uses outputConfig.textFormat. + for id in [ + "us.amazon.nova-pro-v1:0", + "eu.amazon.nova-lite-v1:0", + "ap.amazon.nova-micro-v1:0", + ] { + let model = BedrockModel::from_id(id); + assert!( + model.supports(ModelCapability::NativeStructuredOutput), + "{id} should support NativeStructuredOutput" + ); + } + } + + #[test] + fn test_nova_capabilities() { + let nova_pro = BedrockModel::Direct(DirectModel::NovaProV1); + assert!(nova_pro.supports(ModelCapability::Chat)); + assert!(nova_pro.supports(ModelCapability::Vision)); + assert!(nova_pro.supports(ModelCapability::ToolUse)); + assert!(nova_pro.supports(ModelCapability::NativeStructuredOutput)); + assert!(!nova_pro.supports(ModelCapability::Embeddings)); + + let nova_lite = BedrockModel::Direct(DirectModel::NovaLiteV1); + assert!(nova_lite.supports(ModelCapability::Vision)); + assert!(nova_lite.supports(ModelCapability::NativeStructuredOutput)); + + let nova_micro = BedrockModel::Direct(DirectModel::NovaMicroV1); + assert!(!nova_micro.supports(ModelCapability::Vision)); + assert!(nova_micro.supports(ModelCapability::ToolUse)); + assert!(nova_micro.supports(ModelCapability::NativeStructuredOutput)); + + let nova_pro_eu = BedrockModel::eu(CrossRegionModel::NovaProV1); + assert!(nova_pro_eu.supports(ModelCapability::NativeStructuredOutput)); + assert!(nova_pro_eu.supports(ModelCapability::ToolUse)); + + let nova_micro_eu = BedrockModel::eu(CrossRegionModel::NovaMicroV1); + assert!(nova_micro_eu.supports(ModelCapability::NativeStructuredOutput)); + } + + #[test] + fn test_claude_does_not_support_native_structured_output() { + let claude = BedrockModel::Direct(DirectModel::ClaudeSonnet4); + assert!(!claude.supports(ModelCapability::NativeStructuredOutput)); + + let claude_eu = BedrockModel::eu(CrossRegionModel::ClaudeSonnet4); + assert!(!claude_eu.supports(ModelCapability::NativeStructuredOutput)); + } + + #[test] + fn test_nova_context_windows() { + let nova_pro = BedrockModel::Direct(DirectModel::NovaProV1); + assert_eq!(nova_pro.context_window(), 300_000); + + let nova_lite = BedrockModel::Direct(DirectModel::NovaLiteV1); + assert_eq!(nova_lite.context_window(), 300_000); + + let nova_micro = BedrockModel::Direct(DirectModel::NovaMicroV1); + assert_eq!(nova_micro.context_window(), 128_000); + + let nova_pro_eu = BedrockModel::eu(CrossRegionModel::NovaProV1); + assert_eq!(nova_pro_eu.context_window(), 300_000); + } + + #[test] + fn test_capability_override_native_structured_output() { + let mut override_entry = ModelCapabilityOverride::default(); + assert_eq!( + override_entry.supports(ModelCapability::NativeStructuredOutput), + None + ); + + override_entry.native_structured_output = Some(true); + assert_eq!( + override_entry.supports(ModelCapability::NativeStructuredOutput), + Some(true) + ); + + override_entry.native_structured_output = Some(false); + assert_eq!( + override_entry.supports(ModelCapability::NativeStructuredOutput), + Some(false) + ); + } + #[test] fn test_model_id() { let model = BedrockModel::Direct(DirectModel::ClaudeSonnet4);