diff --git a/CHANGELOG.md b/CHANGELOG.md index e994e26a9..0e5686518 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Breaking +- The prompt a provider receives is modelled by role (`aimux-core`), as the AI + SDK's `LanguageModelV4Message`. `LanguageModelPrompt` is + `Vec`: `System { content: String }`, + `User { content: Vec }`, `Assistant { content: Vec }`, + `Tool { content: Vec }`. The parts are `TextPart`, `FilePart`, + `ReasoningPart`, `ToolCallPart`, `ToolResultPart`. A provider sees one file + part, `FilePart { data: FileData, media_type, filename }`; the five + user-facing file variants are folded into it by + `convert_to_language_model_prompt`, which now returns `Result` and rejects a + part its role does not allow. `LanguageModelPromptMessage` is removed. The + user-facing `ModelMessage` and `ContentPart` are unchanged. Recordings store + the new prompt JSON: a system message's `content` is a string and every file + input is a `file` part. + - Provider options and provider metadata have one typed shape, the AI SDK's `Record`: `SharedProviderOptions` and `SharedProviderMetadata` (`ProviderMetadata` is an alias) are diff --git a/aimux-core/src/composite.rs b/aimux-core/src/composite.rs index 022a95340..d4f3b5505 100644 --- a/aimux-core/src/composite.rs +++ b/aimux-core/src/composite.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use crate::language_model::LanguageModel; -use crate::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; +use crate::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use crate::result::GenerateContent; use crate::types::{InputTokenUsage, OutputTokenUsage, Usage}; @@ -117,10 +117,6 @@ pub fn build_aggregator_prompt( body.push('\n'); } - out.push(LanguageModelPromptMessage { - role: crate::message::Role::User, - content: vec![crate::content::ContentPart::text(body)], - provider_options: None, - }); + out.push(LanguageModelMessage::user_text(body)); out } diff --git a/aimux-core/src/generate.rs b/aimux-core/src/generate.rs index 70fd5ad89..fb614db66 100644 --- a/aimux-core/src/generate.rs +++ b/aimux-core/src/generate.rs @@ -527,12 +527,22 @@ pub async fn generate_text( prompt: impl Into, options: GenerateTextOptions, ) -> Result { - // 1. Convert user prompt to provider-facing prompt. + let (messages, instructions) = split_prompt(prompt.into(), options.instructions.as_deref()); + let lm_prompt = convert_to_language_model_prompt(&messages, instructions)?; + generate_text_from_language_model_prompt(model, lm_prompt, messages, options).await +} + +// Replay already has a provider prompt. Keep it intact instead of converting +// through user content, which cannot represent every provider part. +pub(crate) async fn generate_text_from_language_model_prompt( + model: &dyn LanguageModel, + lm_prompt: crate::language_model_message::LanguageModelPrompt, + messages: Vec, + options: GenerateTextOptions, +) -> Result { let repair_tool_call = options.repair_tool_call.clone(); let tools = options.tools.clone(); let operation_instructions = options.instructions.clone(); - let (messages, instructions) = split_prompt(prompt.into(), operation_instructions.as_deref()); - let lm_prompt = convert_to_language_model_prompt(&messages, instructions); // 2. Build CallOptions. let mut call_options = options.into_call_options(lm_prompt); @@ -874,7 +884,7 @@ pub async fn stream_text( let tools = options.tools.clone(); let operation_instructions = options.instructions.clone(); let (messages, instructions) = split_prompt(prompt.into(), operation_instructions.as_deref()); - let lm_prompt = convert_to_language_model_prompt(&messages, instructions); + let lm_prompt = convert_to_language_model_prompt(&messages, instructions)?; // 2. Build CallOptions. let mut call_options = options.into_call_options(lm_prompt); diff --git a/aimux-core/src/language_model_message.rs b/aimux-core/src/language_model_message.rs index ab198e193..2f8f920ee 100644 --- a/aimux-core/src/language_model_message.rs +++ b/aimux-core/src/language_model_message.rs @@ -1,69 +1,501 @@ -//! Provider-facing message types. +//! Provider-facing prompt types. //! -//! `LanguageModelPrompt` is the standardized prompt format that providers receive -//! in `CallOptions.prompt`. It is converted from `ModelMessage` (user-facing) -//! by the `generate_text` / `stream_text` functions. +//! Aligned with Vercel AI SDK `LanguageModelV4Prompt` / `LanguageModelV4Message` +//! (`provider/src/language-model/v4/language-model-v4-prompt.ts`): a prompt is +//! a list of messages, a message is a union by role, and there is exactly one +//! file part. `convert_to_language_model_prompt` is the single conversion from +//! the user-facing `ModelMessage`s. use crate::content::ContentPart; -use crate::message::Role; -use crate::shared::SharedProviderOptions; +use crate::error::AiMuxError; +use crate::message::{MessageContent, ModelMessage, Role}; +use crate::shared::{FileBytes, FileData, GeneratedFileData, SharedProviderOptions}; +use std::borrow::Cow; +use std::collections::HashMap; + use serde::{Deserialize, Serialize}; +use serde_json::Value; use ts_rs::TS; -/// A single provider-facing message. +/// The standardized prompt passed to `LanguageModel::do_generate` / `do_stream`. +pub type LanguageModelPrompt = Vec; + +/// A single provider-facing message, a union by role. /// -/// Unlike `ModelMessage` (which can have string content), this always has -/// `content: Vec` — strings are normalized to `ContentPart::Text`. -#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, TS)] +/// `provider_options` is message-level (e.g. `anthropic.cacheControl`), +/// mirroring `LanguageModelV4Message.providerOptions`. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[serde(tag = "role", rename_all = "snake_case")] +#[ts(export)] +pub enum LanguageModelMessage { + System { + content: String, + provider_options: Option, + }, + User { + content: Vec, + provider_options: Option, + }, + Assistant { + content: Vec, + provider_options: Option, + }, + Tool { + content: Vec, + provider_options: Option, + }, +} + +/// Text content part. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] #[ts(export)] -pub struct LanguageModelPromptMessage { - pub role: Role, - pub content: Vec, - /// Message-level provider-specific options (e.g. - /// `anthropic.cacheControl`), applied to the last part of the message when - /// converting to a provider prompt. Mirrors the Vercel AI SDK - /// `LanguageModelV4Message.providerOptions`. +pub struct TextPart { + pub text: String, + #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_options: Option, } -/// The standardized prompt passed to `LanguageModel::do_generate` / `do_stream`. -pub type LanguageModelPrompt = Vec; +/// File content part (the only file part; `data` is a tagged `FileData`). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[ts(export)] +pub struct FilePart { + pub data: FileData, + pub media_type: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub filename: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_options: Option, +} + +/// Reasoning / thinking content part. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[ts(export)] +#[ts(rename = "LanguageModelReasoningPart")] +pub struct ReasoningPart { + pub text: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_options: Option, +} + +/// Tool call content part. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[ts(export)] +pub struct ToolCallPart { + pub tool_call_id: String, + pub tool_name: String, + pub input: Value, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_executed: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_options: Option, +} + +/// Tool result content part. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[ts(export)] +pub struct ToolResultPart { + pub tool_call_id: String, + pub tool_name: String, + pub output: ToolResultOutput, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_options: Option, +} + +/// The provider-facing result of a tool call. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[serde(tag = "type", rename_all = "kebab-case")] +#[ts(export)] +pub enum ToolResultOutput { + Text { + value: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + provider_options: Option, + }, + Json { + value: Value, + #[serde(default, skip_serializing_if = "Option::is_none")] + provider_options: Option, + }, + ExecutionDenied { + #[serde(default, skip_serializing_if = "Option::is_none")] + reason: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + provider_options: Option, + }, + ErrorText { + value: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + provider_options: Option, + }, + ErrorJson { + value: Value, + #[serde(default, skip_serializing_if = "Option::is_none")] + provider_options: Option, + }, + Content { + value: Vec, + }, +} + +/// Content within a tool result. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[serde(tag = "type", rename_all = "kebab-case")] +#[ts(export)] +pub enum ToolResultContent { + Text(TextPart), + File(FilePart), + Custom { + #[serde(default, skip_serializing_if = "Option::is_none")] + provider_options: Option, + }, +} + +/// A file generated as part of reasoning. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[ts(export)] +pub struct ReasoningFilePart { + pub data: GeneratedFileData, + pub media_type: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_options: Option, +} + +/// Provider-specific content identified by its provider and kind. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[ts(export)] +pub struct CustomPart { + pub kind: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_options: Option, +} + +/// The user's decision for a provider-executed tool approval request. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[ts(export)] +pub struct ToolApprovalResponsePart { + pub approval_id: String, + pub approved: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reason: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_options: Option, +} + +/// Parts allowed in a user message. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[serde(tag = "type", rename_all = "kebab-case")] +#[ts(export)] +pub enum UserPart { + Text(TextPart), + File(FilePart), +} + +/// Parts allowed in an assistant message. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[serde(tag = "type", rename_all = "kebab-case")] +#[ts(export)] +pub enum AssistantPart { + Text(TextPart), + File(FilePart), + Reasoning(ReasoningPart), + ReasoningFile(ReasoningFilePart), + Custom(CustomPart), + ToolCall(ToolCallPart), + ToolResult(ToolResultPart), +} + +/// Parts allowed in a tool message. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[serde(tag = "type", rename_all = "kebab-case")] +#[ts(export)] +pub enum ToolPart { + ToolResult(ToolResultPart), + ToolApprovalResponse(ToolApprovalResponsePart), +} + +impl LanguageModelMessage { + /// User message with a single text part. + #[must_use] + pub fn user_text(text: impl Into) -> Self { + Self::User { + content: vec![UserPart::Text(TextPart { + text: text.into(), + provider_options: None, + })], + provider_options: None, + } + } +} + +/// Map a user-facing part onto the widest provider part (the assistant union); +/// the role then narrows it. +fn to_assistant_part( + part: &ContentPart, + tool_names: &HashMap, +) -> Result { + let file = |data, media_type: &String, filename: &Option, po: &Option<_>| { + AssistantPart::File(FilePart { + data, + media_type: media_type.clone(), + filename: filename.clone(), + provider_options: po.clone(), + }) + }; + Ok(match part { + ContentPart::Text { + text, + provider_options, + } => AssistantPart::Text(TextPart { + text: text.clone(), + provider_options: provider_options.clone(), + }), + ContentPart::Image { + image, + media_type, + provider_options, + } => file( + FileData::Data { + data: FileBytes::Binary(image.clone()), + }, + media_type, + &None, + provider_options, + ), + ContentPart::File { + data, + media_type, + filename, + provider_options, + } => file( + FileData::Data { + data: FileBytes::Binary(data.clone()), + }, + media_type, + filename, + provider_options, + ), + ContentPart::FileBase64 { + data, + media_type, + filename, + provider_options, + } => file( + FileData::Data { + data: FileBytes::Base64(data.clone()), + }, + media_type, + filename, + provider_options, + ), + ContentPart::FileUrl { + url, + media_type, + provider_options, + } => file( + FileData::Url { + url: url.clone(), + original_url: None, + }, + media_type, + &None, + provider_options, + ), + ContentPart::FileReference { + media_type, + reference, + filename, + provider_options, + } => file( + FileData::Reference { + reference: serde_json::from_value(reference.clone()).map_err(|e| { + AiMuxError::InvalidPrompt(format!("invalid file reference: {e}")) + })?, + }, + media_type, + filename, + provider_options, + ), + ContentPart::Reasoning { + text, + provider_options, + .. + } => AssistantPart::Reasoning(ReasoningPart { + text: text.clone(), + provider_options: provider_options.clone(), + }), + ContentPart::ToolCall { + tool_call_id, + tool_name, + input, + provider_executed, + provider_options, + } => AssistantPart::ToolCall(ToolCallPart { + tool_call_id: tool_call_id.clone(), + tool_name: tool_name.clone(), + input: input.clone(), + provider_executed: *provider_executed, + provider_options: provider_options.clone(), + }), + ContentPart::ToolResult { + tool_call_id, + result, + tool_name, + is_error, + provider_options, + .. + } => AssistantPart::ToolResult(ToolResultPart { + tool_call_id: tool_call_id.clone(), + tool_name: tool_name + .as_ref() + .or_else(|| tool_names.get(tool_call_id)) + .cloned() + .ok_or_else(|| { + AiMuxError::InvalidPrompt(format!( + "tool result {tool_call_id} has no tool name or preceding tool call" + )) + })?, + output: if *is_error == Some(true) { + ToolResultOutput::ErrorText { + value: match result { + Value::Null => "unknown error".into(), + Value::String(value) => value.clone(), + _ => result.to_string(), + }, + provider_options: None, + } + } else if let Value::String(value) = result { + ToolResultOutput::Text { + value: value.clone(), + provider_options: None, + } + } else { + ToolResultOutput::Json { + value: result.clone(), + provider_options: None, + } + }, + provider_options: provider_options.clone(), + }), + ContentPart::Custom { + kind, + provider_options, + } => AssistantPart::Custom(CustomPart { + kind: kind.clone(), + provider_options: provider_options.clone(), + }), + ContentPart::ReasoningFile { + data, + media_type, + provider_options, + } => AssistantPart::ReasoningFile(ReasoningFilePart { + data: data.clone(), + media_type: media_type.clone(), + provider_options: provider_options.clone(), + }), + ContentPart::ToolApprovalRequest { .. } => { + return Err(AiMuxError::InvalidPrompt( + "tool approval requests are only allowed in assistant messages".into(), + )); + } + }) +} /// Convert user-facing `ModelMessage`s into a `LanguageModelPrompt`. /// -/// - String content is normalized to a single `ContentPart::Text`. +/// - String content is normalized to a single text part. +/// - The five user-facing file variants become one `FilePart`. /// - If `instructions` is provided, it is prepended as a system message. -#[must_use] +/// +/// # Errors +/// +/// `InvalidPrompt` when a part is not allowed for its role (user: text/file; +/// tool: tool-result), a system message is not plain text, or a file +/// reference is not a `{ provider: id }` map. pub fn convert_to_language_model_prompt( - messages: &[crate::message::ModelMessage], + messages: &[ModelMessage], instructions: Option<&str>, -) -> LanguageModelPrompt { +) -> Result { let mut result = Vec::new(); - - // Prepend system instructions if provided. if let Some(instructions) = instructions { - result.push(LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text(instructions)], + result.push(LanguageModelMessage::System { + content: instructions.to_string(), provider_options: None, }); } + let mut tool_names = HashMap::new(); for msg in messages { - let parts = match &msg.content { - crate::message::MessageContent::Text(text) => vec![ContentPart::text(text)], - crate::message::MessageContent::Parts(parts) => parts - .iter() - .filter(|part| !matches!(part, ContentPart::ToolApprovalRequest { .. })) - .cloned() - .collect(), + let bad = |what: &str| AiMuxError::InvalidPrompt(format!("{:?} message: {what}", msg.role)); + if msg.role == Role::System { + match &msg.content { + MessageContent::Text(text) => result.push(LanguageModelMessage::System { + content: text.clone(), + provider_options: None, + }), + _ => return Err(bad("system content must be plain text")), + } + continue; + } + let parts: Cow<'_, [ContentPart]> = match &msg.content { + MessageContent::Text(text) => Cow::Owned(vec![ContentPart::text(text)]), + MessageContent::Parts(parts) => Cow::Borrowed(parts), }; - result.push(LanguageModelPromptMessage { - role: msg.role, - content: parts, - provider_options: None, + if msg.role == Role::Assistant { + for part in parts.iter() { + if let ContentPart::ToolCall { + tool_call_id, + tool_name, + .. + } = part + { + tool_names.insert(tool_call_id.clone(), tool_name.clone()); + } + } + } + let parts = parts + .iter() + .filter(|part| { + msg.role != Role::Assistant + || !matches!(part, ContentPart::ToolApprovalRequest { .. }) + }) + .filter(|part| { + !matches!(&msg.content, MessageContent::Parts(_)) + || !matches!(part, ContentPart::Text { text, provider_options } + if text.is_empty() && (msg.role == Role::User + || (msg.role == Role::Assistant && provider_options.is_none()))) + }) + .map(|part| to_assistant_part(part, &tool_names)) + .collect::, _>>()?; + let provider_options = None; + result.push(match msg.role { + Role::System => unreachable!(), + Role::Assistant => LanguageModelMessage::Assistant { + content: parts, + provider_options, + }, + Role::User => LanguageModelMessage::User { + content: parts + .into_iter() + .map(|p| match p { + AssistantPart::Text(t) => Ok(UserPart::Text(t)), + AssistantPart::File(f) => Ok(UserPart::File(f)), + _ => Err(bad("only text and file parts are allowed")), + }) + .collect::>()?, + provider_options, + }, + Role::Tool => LanguageModelMessage::Tool { + content: parts + .into_iter() + .map(|p| match p { + AssistantPart::ToolResult(r) => Ok(ToolPart::ToolResult(r)), + _ => Err(bad("only tool-result parts are allowed")), + }) + .collect::>()?, + provider_options, + }, }); } - result + Ok(result) } diff --git a/aimux-core/src/recording.rs b/aimux-core/src/recording.rs index d73c3be16..25ad2fc6f 100644 --- a/aimux-core/src/recording.rs +++ b/aimux-core/src/recording.rs @@ -1884,7 +1884,8 @@ mod tests { crate::language_model_message::convert_to_language_model_prompt( &[crate::message::ModelMessage::user("ping")], None, - ), + ) + .unwrap(), ) } diff --git a/aimux-core/src/replay.rs b/aimux-core/src/replay.rs index f56ff9d9f..8d641d7a0 100644 --- a/aimux-core/src/replay.rs +++ b/aimux-core/src/replay.rs @@ -16,10 +16,13 @@ use async_trait::async_trait; use std::sync::Arc; use crate::error::AiMuxError; -use crate::generate::{GenerateTextOptions, GenerateTextResult, generate_text}; +use crate::generate::{ + GenerateTextOptions, GenerateTextResult, generate_text_from_language_model_prompt, +}; use crate::language_model::LanguageModel; -use crate::language_model_message::LanguageModelPrompt; -use crate::message::{MessageContent, ModelMessage, ModelPrompt}; +use crate::language_model_message::{ + AssistantPart, LanguageModelMessage, LanguageModelPrompt, UserPart, +}; use crate::options::CallOptions; use crate::recording::Recording; use crate::result::{GenerateContent, GenerateResult, StreamResult}; @@ -249,10 +252,42 @@ impl ReplayMatcher for ScoreMatcher { } /// 消息的首个文本内容(非文本消息回退到序列化)。 -fn message_text(m: &crate::language_model_message::LanguageModelPromptMessage) -> String { - match m.content.first() { - Some(crate::content::ContentPart::Text { text, .. }) => text.clone(), - _ => serde_json::to_string(&m.content).unwrap_or_default(), +fn message_text(m: &LanguageModelMessage) -> String { + match m { + LanguageModelMessage::System { content, .. } => content.clone(), + LanguageModelMessage::User { content, .. } => match content.first() { + Some(UserPart::Text(part)) => part.text.clone(), + _ => serde_json::to_string(content).unwrap_or_default(), + }, + LanguageModelMessage::Assistant { content, .. } => match content.first() { + Some(AssistantPart::Text(part)) => part.text.clone(), + _ => serde_json::to_string(content).unwrap_or_default(), + }, + LanguageModelMessage::Tool { content, .. } => { + serde_json::to_string(content).unwrap_or_default() + } + } +} + +fn same_message_content(a: &LanguageModelMessage, b: &LanguageModelMessage) -> bool { + match (a, b) { + ( + LanguageModelMessage::System { content: a, .. }, + LanguageModelMessage::System { content: b, .. }, + ) => a == b, + ( + LanguageModelMessage::User { content: a, .. }, + LanguageModelMessage::User { content: b, .. }, + ) => a == b, + ( + LanguageModelMessage::Assistant { content: a, .. }, + LanguageModelMessage::Assistant { content: b, .. }, + ) => a == b, + ( + LanguageModelMessage::Tool { content: a, .. }, + LanguageModelMessage::Tool { content: b, .. }, + ) => a == b, + _ => false, } } @@ -268,7 +303,7 @@ fn match_score(options: &CallOptions, rec: &Recording) -> u64 { .prompt .iter() .zip(rec.input.prompt.iter()) - .take_while(|(a, b)| a.role == b.role && a.content == b.content) + .take_while(|(a, b)| same_message_content(a, b)) .count(); // 字符级 LCP:第一个不同消息(或双方最后一个消息),在文本内容上计算 // (避免 JSON 结构前缀的伪匹配)。任一 prompt 为空时无消息可比,LCP=0 @@ -343,7 +378,7 @@ impl ReplayMatcher for PrefixMatcher { .prompt .iter() .zip(r.input.prompt.iter()) - .take_while(|(a, b)| a.role == b.role && a.content == b.content) + .take_while(|(a, b)| same_message_content(a, b)) .count(); if common < r.input.prompt.len() { continue; // 录制未被输入完整包含,不命中。 @@ -915,8 +950,8 @@ pub struct ReplayOverrides { /// # Errors /// /// Returns `AiMuxError::JsonParse` when the recorded call options cannot be -/// deserialized, and propagates any error from the re-sent `generate_text` -/// call. +/// deserialized, `AiMuxError::InvalidPrompt` for inline text file data without +/// a user-facing representation, and propagates errors from `generate_text`. pub async fn replay_with_model( recording: &Recording, model: &dyn LanguageModel, @@ -944,26 +979,11 @@ pub async fn replay_with_model( } } - // 3. 转回用户侧类型,经 generate_text 重发。 - let prompt = model_prompt_from_lm(&call_options.prompt); + // Recorded options already contain the complete provider-facing prompt. + // Repair callbacks are runtime-only and are not restored from recordings. + let prompt = std::mem::take(&mut call_options.prompt); let options = generate_options_from_call_options(call_options); - generate_text(model, prompt, options).await -} - -/// `LanguageModelPrompt`(provider 侧)→ `ModelPrompt`(用户侧)。 -/// -/// 逐消息 `provider_options` 丢弃(用户侧无对应字段;语义不丢,仅 cacheControl -/// 之类的 provider 提示不参与重放)。 -fn model_prompt_from_lm(prompt: &LanguageModelPrompt) -> ModelPrompt { - ModelPrompt::Messages( - prompt - .iter() - .map(|m| ModelMessage { - role: m.role, - content: MessageContent::Parts(m.content.clone()), - }) - .collect(), - ) + generate_text_from_language_model_prompt(model, prompt, Vec::new(), options).await } /// `CallOptions` → `GenerateTextOptions`(record 侧到用户侧的反向映射)。 @@ -1009,14 +1029,7 @@ mod tests { use futures::StreamExt; fn sample_options(text: &str, temperature: Option) -> CallOptions { - let prompt = vec![crate::language_model_message::LanguageModelPromptMessage { - role: crate::message::Role::User, - content: vec![crate::content::ContentPart::Text { - text: text.to_string(), - provider_options: None, - }], - provider_options: None, - }]; + let prompt = vec![LanguageModelMessage::user_text(text.to_string())]; GenerateTextOptions { temperature, ..Default::default() @@ -1027,14 +1040,7 @@ mod tests { /// 构造 OpenAI 格式录制。 fn openai_recording(trace_id: &str, prompt_text: &str, reply: &str, finish: &str) -> Recording { let input_prompt: crate::language_model_message::LanguageModelPrompt = - vec![crate::language_model_message::LanguageModelPromptMessage { - role: crate::message::Role::User, - content: vec![crate::content::ContentPart::Text { - text: prompt_text.to_string(), - provider_options: None, - }], - provider_options: None, - }]; + vec![LanguageModelMessage::user_text(prompt_text.to_string())]; let body = serde_json::json!({ "id": "chatcmpl-mock", "object": "chat.completion", @@ -1335,22 +1341,8 @@ mod tests { let matcher = PrefixMatcher::new("openai", "gpt-4o"); // 输入两条消息 ["hello", "world"]:rec 是消息级前缀 → 命中。 let prompt = vec![ - crate::language_model_message::LanguageModelPromptMessage { - role: crate::message::Role::User, - content: vec![crate::content::ContentPart::Text { - text: "hello".into(), - provider_options: None, - }], - provider_options: None, - }, - crate::language_model_message::LanguageModelPromptMessage { - role: crate::message::Role::User, - content: vec![crate::content::ContentPart::Text { - text: "world".into(), - provider_options: None, - }], - provider_options: None, - }, + LanguageModelMessage::user_text("hello"), + LanguageModelMessage::user_text("world"), ]; let opts = GenerateTextOptions::default().into_call_options(prompt); let recs = [rec.clone()]; @@ -1358,16 +1350,8 @@ mod tests { assert_eq!(hit.call_id, "t1"); // 首消息不同 → 不命中(整条消息级前缀,非字符级)。 - let opts2 = GenerateTextOptions::default().into_call_options(vec![ - crate::language_model_message::LanguageModelPromptMessage { - role: crate::message::Role::User, - content: vec![crate::content::ContentPart::Text { - text: "hi".into(), - provider_options: None, - }], - provider_options: None, - }, - ]); + let opts2 = GenerateTextOptions::default() + .into_call_options(vec![LanguageModelMessage::user_text("hi")]); assert!(matcher.r#match(&opts2, &[rec]).is_err()); } @@ -1376,14 +1360,7 @@ mod tests { // 多消息录制:rec-a = ["hello"],rec-b = ["hello", "world"]。 let mut rec_a = openai_recording("ta", "hello", "a", "stop"); let mut rec_b = openai_recording("tb", "hello world", "b", "stop"); - let mk = |text: &str| crate::language_model_message::LanguageModelPromptMessage { - role: crate::message::Role::User, - content: vec![crate::content::ContentPart::Text { - text: text.into(), - provider_options: None, - }], - provider_options: None, - }; + let mk = |text: &str| LanguageModelMessage::user_text(text); rec_a.input.prompt = vec![mk("hello")]; rec_b.input.prompt = vec![mk("hello"), mk("world")]; diff --git a/aimux-core/src/session.rs b/aimux-core/src/session.rs index ead68fb6f..315c00c0e 100644 --- a/aimux-core/src/session.rs +++ b/aimux-core/src/session.rs @@ -452,25 +452,6 @@ fn civil_from_days(z: i64) -> (i64, u32, u32) { #[cfg(test)] mod tests { use super::*; - use crate::content::ContentPart; - use crate::language_model_message::LanguageModelPromptMessage; - use crate::message::Role; - - fn msg(text: &str) -> LanguageModelPromptMessage { - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::Text { - text: text.to_string(), - provider_options: None, - }], - provider_options: None, - } - } - - fn prompt(messages: Vec) -> LanguageModelPrompt { - messages - } - /// `recorded_at` must look like RFC 3339 UTC: `YYYY-MM-DDTHH:MM:SS.mmmZ`. fn is_rfc3339(s: &str) -> bool { let b = s.as_bytes(); @@ -606,105 +587,6 @@ mod tests { assert!(store.list_sessions().is_empty()); } - // ── SessionInferer ────────────────────────────────────────────────────── - - #[test] - fn inferer_disabled_returns_none() { - let mut inf = SessionInferer::new(false); - assert!(inf.resolve(None, &prompt(vec![msg("hi")])).is_none()); - } - - #[test] - fn inferer_explicit_wins_and_is_remembered() { - let mut inf = SessionInferer::new(true); - let p1 = prompt(vec![msg("hi")]); - let (id, source) = inf.resolve(Some("sess-1"), &p1).unwrap(); - assert_eq!(id, "sess-1"); - assert_eq!(source, SessionSource::Explicit); - - // Continuation without an explicit id joins the remembered session. - let p2 = prompt(vec![msg("hi"), msg("hello")]); - let (id2, source2) = inf.resolve(None, &p2).unwrap(); - assert_eq!(id2, "sess-1"); - assert_eq!(source2, SessionSource::Inferred); - } - - #[test] - fn inferer_strong_prefix_continuation_groups_calls() { - let mut inf = SessionInferer::new(true); - let p1 = prompt(vec![msg("u1")]); - let (id1, source1) = inf.resolve(None, &p1).unwrap(); - assert!(id1.starts_with("auto-")); - assert_eq!(source1, SessionSource::Inferred); - - // Full-prefix continuation → same session. - let p2 = prompt(vec![msg("u1"), msg("a1"), msg("u2")]); - let (id2, _) = inf.resolve(None, &p2).unwrap(); - assert_eq!(id2, id1); - - // Equality (retry of the same step) → same session. - let (id3, _) = inf.resolve(None, &p1).unwrap(); - assert_eq!(id3, id1); - } - - #[test] - fn inferer_distinct_prompts_start_new_sessions() { - let mut inf = SessionInferer::new(true); - let (id1, _) = inf - .resolve(None, &prompt(vec![msg("what is rust")])) - .unwrap(); - let (id2, _) = inf - .resolve(None, &prompt(vec![msg("tell me a joke")])) - .unwrap(); - assert_ne!(id1, id2); - } - - #[test] - fn inferer_prefix_is_directional() { - let mut inf = SessionInferer::new(true); - let (id1, _) = inf.resolve(None, &prompt(vec![msg("abc")])).unwrap(); - // Message-level prefix continuation: leading messages identical and - // the new prompt is longer → same session. - let (id2, _) = inf - .resolve(None, &prompt(vec![msg("abc"), msg("a1"), msg("u2")])) - .unwrap(); - assert_eq!(id2, id1); - // A prompt that diverges on the first message → new session. - let (id3, _) = inf.resolve(None, &prompt(vec![msg("abd")])).unwrap(); - assert_ne!(id3, id1, "abd is not a prefix continuation"); - } - - #[test] - fn inferer_chain_continuation_stays_in_session() { - let mut inf = SessionInferer::new(true); - let p1 = prompt(vec![msg("common")]); - let p2 = prompt(vec![msg("common"), msg("b")]); - let p3 = prompt(vec![msg("common"), msg("b"), msg("c")]); - let (id1, _) = inf.resolve(None, &p1).unwrap(); - let (id2, _) = inf.resolve(None, &p2).unwrap(); - let (id3, _) = inf.resolve(None, &p3).unwrap(); - assert_eq!(id2, id1); - assert_eq!(id3, id1); - } - - #[test] - fn inferer_capacity_evicts_oldest_prompt() { - let mut inf = SessionInferer::with_capacity(true, 2); - let p1 = prompt(vec![msg("one")]); - let (id1, _) = inf.resolve(None, &p1).unwrap(); - let p2 = prompt(vec![msg("two")]); - let (id2, _) = inf.resolve(None, &p2).unwrap(); - let p3 = prompt(vec![msg("three")]); - let (id3, _) = inf.resolve(None, &p3).unwrap(); - assert_ne!(id1, id2); - assert_ne!(id2, id3); - - // p1 was evicted; a new call extending it starts a NEW session. - let p1_ext = prompt(vec![msg("one"), msg("more")]); - let (id4, _) = inf.resolve(None, &p1_ext).unwrap(); - assert_ne!(id4, id1, "evicted prompt no longer matches"); - } - // ── helpers ───────────────────────────────────────────────────────────── #[test] diff --git a/aimux-provider-utils/tests/http_recording_test.rs b/aimux-provider-utils/tests/http_recording_test.rs index 58dac8160..09ba83e15 100644 --- a/aimux-provider-utils/tests/http_recording_test.rs +++ b/aimux-provider-utils/tests/http_recording_test.rs @@ -62,18 +62,12 @@ fn request(url: String, call_id: &str, recorder: Arc) -> HttpReque } fn finish_call(recorder: &dyn Recorder, call_id: &str) { - use aimux_core::content::ContentPart; - use aimux_core::language_model_message::LanguageModelPromptMessage; - use aimux_core::message::Role; + use aimux_core::language_model_message::LanguageModelMessage; use aimux_core::options::CallOptions; recorder.record_input( call_id, - &CallOptions::new(vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("hi")], - ..Default::default() - }]), + &CallOptions::new(vec![LanguageModelMessage::user_text("hi")]), "openai", "gpt-test", ); diff --git a/aimux-providers/src/anthropic/convert.rs b/aimux-providers/src/anthropic/convert.rs index 4698795c5..af679e5f5 100644 --- a/aimux-providers/src/anthropic/convert.rs +++ b/aimux-providers/src/anthropic/convert.rs @@ -20,12 +20,13 @@ use std::collections::{BTreeSet, HashSet}; -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; -use aimux_core::language_model_message::LanguageModelPrompt; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, +}; use aimux_core::options::{CallOptions, Tool}; -use aimux_core::shared::SharedProviderOptions; +use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; use aimux_core::types::{FinishReason, FinishReasonUnified, ReasoningEffort, Warning}; use serde_json::{Map, Value, json}; @@ -133,129 +134,108 @@ pub fn convert_prompt_to_anthropic_full_with_tools( } for msg in prompt { - match msg.role { - Role::System => { - // A system message interrupts any in-flight user/assistant group. + let role = match msg { + LanguageModelMessage::System { .. } => None, + LanguageModelMessage::User { .. } | LanguageModelMessage::Tool { .. } => { + Some(Eff::User) + } + LanguageModelMessage::Assistant { .. } => Some(Eff::Assistant), + }; + if let Some(role) = role { + seen_non_system = true; + if last != Some(role) { flush(&mut messages, &mut last, &mut acc); - let n = msg.content.len(); - let mut blocks: Vec = Vec::new(); - for (idx, p) in msg.content.iter().enumerate() { - if let ContentPart::Text { - text, - provider_options, - } = p - { - let is_last = idx + 1 == n; - // cache_control: part-level first, then message-level - // for the last text part (system content is a string in - // the TS model, so message-level is the primary path). - let cache_control = match validator.get_cache_control( - provider_options.as_ref().and_then(|o| o.get("anthropic")), - "system message part", - true, - ) { - Some(v) => Some(v), - None => { - if is_last { - validator.get_cache_control( - msg.provider_options - .as_ref() - .and_then(|o| o.get("anthropic")), - "system message", - true, - ) - } else { - None - } - } - }; - let mut block = json!({ "type": "text", "text": text }); - if let Some(cc) = cache_control { - block["cache_control"] = cc; - } - blocks.push(block); - } - // Non-text parts on a system message are ignored. - } - + } + last = Some(role); + } + match msg { + LanguageModelMessage::System { + content, + provider_options, + } => { + flush(&mut messages, &mut last, &mut acc); + let cc = validator.get_cache_control( + provider_options.as_ref().and_then(|o| o.get("anthropic")), + "system message", + true, + ); + let blocks = vec![apply_cc(json!({ "type": "text", "text": content }), cc)]; if !seen_non_system { - // Initial system block(s) become the top-level system prompt. system.extend(blocks); } else { - // A mid-conversation system message is emitted inline. messages.push(json!({ "role": "system", "content": blocks })); betas.insert(BETA_MID_CONVERSATION_SYSTEM.to_string()); } } - Role::User | Role::Tool => { - seen_non_system = true; - if last != Some(Eff::User) { - flush(&mut messages, &mut last, &mut acc); + LanguageModelMessage::User { + content, + provider_options, + } => { + for (idx, part) in content.iter().enumerate() { + acc.push(match part { + UserPart::Text(TextPart { + text, + provider_options: part_options, + }) => { + let cc = resolve_cache_control( + &mut validator, + part_options.as_ref(), + provider_options.as_ref(), + idx + 1 == content.len(), + "user message part", + "user message", + ); + apply_cc(json!({ "type": "text", "text": text }), cc) + } + UserPart::File(file) => convert_file_part( + file, + &mut betas, + &mut validator, + provider_options.as_ref(), + idx + 1 == content.len(), + "user message part", + "user message", + )?, + }); + } + } + LanguageModelMessage::Tool { + content, + provider_options, + } => { + for (idx, part) in content.iter().enumerate() { + let ToolPart::ToolResult(part) = part else { + continue; + }; + acc.push(convert_tool_result( + part, + &mut validator, + &mut betas, + &mut warnings, + idx + 1 == content.len(), + provider_options.as_ref(), + )?); } - let (part_ctx, msg_ctx) = match msg.role { - Role::Tool => ("tool result part", "tool result message"), - _ => ("user message part", "user message"), - }; - let n = msg.content.len(); - for (idx, p) in msg.content.iter().enumerate() { - let is_last_part = idx + 1 == n; - if let Some(block) = convert_part_to_anthropic( - p, + } + LanguageModelMessage::Assistant { + content, + provider_options, + } => { + for (idx, part) in content.iter().enumerate() { + if let Some(block) = convert_assistant_part( + part, send_reasoning, tool_names, + &mcp_tool_use_ids, &mut betas, &mut warnings, &mut validator, - is_last_part, - msg.provider_options.as_ref(), - part_ctx, - msg_ctx, + idx + 1 == content.len(), + provider_options.as_ref(), )? { acc.push(block); } } - last = Some(Eff::User); - } - Role::Assistant => { - seen_non_system = true; - if last != Some(Eff::Assistant) { - flush(&mut messages, &mut last, &mut acc); - } - let n = msg.content.len(); - for (idx, p) in msg.content.iter().enumerate() { - let is_last_part = idx + 1 == n; - // Assistant-role tool results are provider-executed results; - // Anthropic only accepts a bare `tool_result` in a *user* - // message, so they take a dedicated path. - let block = if let ContentPart::ToolResult { .. } = p { - convert_assistant_tool_result( - p, - tool_names, - &mcp_tool_use_ids, - &mut warnings, - &mut validator, - is_last_part, - msg.provider_options.as_ref(), - ) - } else { - convert_part_to_anthropic( - p, - send_reasoning, - tool_names, - &mut betas, - &mut warnings, - &mut validator, - is_last_part, - msg.provider_options.as_ref(), - "assistant message part", - "assistant message", - )? - }; - if let Some(block) = block { - acc.push(block); - } - } - last = Some(Eff::Assistant); } } } @@ -341,169 +321,248 @@ fn top_level_media_type(media_type: &str) -> &str { media_type.split('/').next().unwrap_or("") } -/// Convert a single content part into an Anthropic content block. -/// -/// Returns `None` when the part should be omitted entirely (e.g. a reasoning -/// part that is dropped). `betas` and `warnings` are appended to as side -/// effects. -/// -/// `cache_control` is resolved from the part's `provider_options.anthropic -/// .cacheControl`, falling back to the message-level `provider_options` for the -/// last part of a message (mirroring the TS SDK). Thinking blocks reject -/// cache_control (validated with a warning). -#[allow(clippy::too_many_arguments)] -fn convert_part_to_anthropic( - part: &ContentPart, - send_reasoning: bool, - tool_names: &ToolNameMapping, - betas: &mut BTreeSet, - warnings: &mut Vec, +fn resolve_cache_control( validator: &mut CacheControlValidator, + part_options: Option<&SharedProviderOptions>, + message_options: Option<&SharedProviderOptions>, is_last_part: bool, - message_provider_options: Option<&SharedProviderOptions>, - part_context_type: &str, - message_context_type: &str, -) -> Result, AiMuxError> { - // Resolve cache_control = part-level ?? (is_last_part ? message-level). - let resolve_cc = |validator: &mut CacheControlValidator, - part_opts: Option<&SharedProviderOptions>| { - match validator.get_cache_control( - part_opts.and_then(|o| o.get("anthropic")), - part_context_type, + part_context: &str, + message_context: &str, +) -> Option { + validator + .get_cache_control( + part_options.and_then(|o| o.get("anthropic")), + part_context, true, - ) { - Some(v) => Some(v), - None => { - if is_last_part { - validator.get_cache_control( - message_provider_options.and_then(|o| o.get("anthropic")), - message_context_type, - true, - ) - } else { - None - } - } - } - }; - - let apply_cc = |block: Value, cc: Option| -> Value { - match cc { - Some(c) => { - let mut b = block; - b["cache_control"] = c; - b + ) + .or_else(|| { + if is_last_part { + validator.get_cache_control( + message_options.and_then(|o| o.get("anthropic")), + message_context, + true, + ) + } else { + None } - None => block, - } - }; - - Ok(Some(match part { - ContentPart::Text { - text, - provider_options, - } => { - let cc = resolve_cc(validator, provider_options.as_ref()); - apply_cc(json!({ "type": "text", "text": text }), cc) - } + }) +} - ContentPart::Image { - image, - media_type, - provider_options, - } => { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(image); - let cc = resolve_cc(validator, provider_options.as_ref()); - apply_cc( - json!({ - "type": "image", - "source": { - "type": "base64", - "media_type": media_type, - "data": b64, - } - }), - cc, - ) - } +fn apply_cc(mut block: Value, cc: Option) -> Value { + if let Some(cc) = cc { + block["cache_control"] = cc; + } + block +} - ContentPart::File { - data, - media_type, - filename, - provider_options, +fn convert_file_part( + file: &FilePart, + betas: &mut BTreeSet, + validator: &mut CacheControlValidator, + message_options: Option<&SharedProviderOptions>, + is_last_part: bool, + part_context: &str, + message_context: &str, +) -> Result { + let FilePart { + data, + media_type, + filename, + provider_options, + } = file; + let block = match data { + FileData::Data { + data: FileBytes::Binary(bytes), } => { - let full = resolve_full_media_type(media_type, data)?; - let block = route_file_bytes(&full, data, filename.as_deref(), betas)?; - let cc = resolve_cc(validator, provider_options.as_ref()); - apply_cc(block, cc) + let full = resolve_full_media_type(media_type, bytes)?; + route_file_bytes(&full, bytes, filename.as_deref(), betas)? } - - ContentPart::FileBase64 { - data, - media_type, - filename, - provider_options, + FileData::Data { + data: FileBytes::Base64(data), } => { use base64::Engine; let bytes = base64::engine::general_purpose::STANDARD .decode(data) .unwrap_or_default(); let full = resolve_full_media_type(media_type, &bytes)?; - let block = route_file_base64(&full, data, &bytes, filename.as_deref(), betas)?; - let cc = resolve_cc(validator, provider_options.as_ref()); - apply_cc(block, cc) - } - - ContentPart::FileUrl { - url, - media_type, - provider_options, - } => { - let block = route_file_url(media_type, url, betas)?; - let cc = resolve_cc(validator, provider_options.as_ref()); - apply_cc(block, cc) + route_file_base64(&full, data, &bytes, filename.as_deref(), betas)? } - - ContentPart::FileReference { - media_type, - reference, - provider_options, - .. - } => { + FileData::Url { url, .. } => route_file_url(media_type, url, betas)?, + FileData::Reference { reference } => { let file_id = resolve_anthropic_reference(reference)?; betas.insert(BETA_FILES_API.to_string()); let container_upload = provider_options .as_ref() .and_then(|o| o.get("anthropic")) .and_then(|a| a.get("containerUpload")) - .and_then(serde_json::Value::as_bool) + .and_then(Value::as_bool) .unwrap_or(false); - let block = if container_upload { - json!({ "type": "container_upload", "file_id": file_id }) + if container_upload { + return Ok(json!({ "type": "container_upload", "file_id": file_id })); } else if top_level_media_type(media_type) == "image" { json!({ "type": "image", "source": { "type": "file", "file_id": file_id } }) } else { json!({ "type": "document", "source": { "type": "file", "file_id": file_id } }) - }; - // container_upload blocks do not carry cache_control in the TS SDK. - let cc = if container_upload { - None - } else { - resolve_cc(validator, provider_options.as_ref()) - }; - apply_cc(block, cc) + } + } + FileData::Text { text } => { + let options = provider_options.as_ref().and_then(|o| o.get("anthropic")); + let mut block = json!({ "type": "document", "source": { + "type": "text", "media_type": "text/plain", "data": text, + }}); + if let Some(title) = options + .and_then(|o| o.get("title")) + .and_then(Value::as_str) + .or(filename.as_deref()) + { + block["title"] = json!(title); + } + if let Some(context) = options + .and_then(|o| o.get("context")) + .and_then(Value::as_str) + .filter(|s| !s.is_empty()) + { + block["context"] = json!(context); + } + if options + .and_then(|o| o.get("citations")) + .and_then(|c| c.get("enabled")) + .and_then(Value::as_bool) + == Some(true) + { + block["citations"] = json!({ "enabled": true }); + } + block + } + }; + let cc = resolve_cache_control( + validator, + provider_options.as_ref(), + message_options, + is_last_part, + part_context, + message_context, + ); + Ok(apply_cc(block, cc)) +} + +fn convert_tool_result( + part: &ToolResultPart, + validator: &mut CacheControlValidator, + betas: &mut BTreeSet, + warnings: &mut Vec, + is_last_part: bool, + message_provider_options: Option<&SharedProviderOptions>, +) -> Result { + let ToolResultPart { + tool_call_id, + output, + provider_options, + .. + } = part; + + if let ToolResultOutput::Content { value } = output { + for item in value { + if let ToolResultContent::File(file) = item + && let FileData::Data { data } = &file.data + { + use base64::Engine; + let bytes = match data { + FileBytes::Binary(bytes) => bytes.clone(), + FileBytes::Base64(data) => base64::engine::general_purpose::STANDARD + .decode(data) + .unwrap_or_default(), + }; + resolve_full_media_type(&file.media_type, &bytes)?; + } + } + } + let (content, is_error) = resolve_tool_result_output(output, betas, warnings); + let mut block = json!({ + "type": "tool_result", + "tool_use_id": tool_call_id, + "content": content, + }); + if is_error { + block["is_error"] = json!(true); + } + // cache_control: part ?? output ?? (is_last_part ? message). + let cc = match validator.get_cache_control( + provider_options.as_ref().and_then(|o| o.get("anthropic")), + "tool result part", + true, + ) { + Some(v) => Some(v), + None => match extract_tool_result_output_provider_options(output) { + Some(out_opts) => { + validator.get_cache_control(out_opts.get("anthropic"), "tool result output", true) + } + None => { + if is_last_part { + validator.get_cache_control( + message_provider_options.and_then(|o| o.get("anthropic")), + "tool result message", + true, + ) + } else { + None + } + } + }, + }; + Ok(apply_cc(block, cc)) +} + +/// Convert an assistant part, preserving reasoning and provider-executed tool results. +/// Part-level cache control falls back to message options on the final part. +#[allow(clippy::too_many_arguments)] +fn convert_assistant_part( + part: &AssistantPart, + send_reasoning: bool, + tool_names: &ToolNameMapping, + mcp_tool_use_ids: &HashSet<&str>, + betas: &mut BTreeSet, + warnings: &mut Vec, + validator: &mut CacheControlValidator, + is_last_part: bool, + message_provider_options: Option<&SharedProviderOptions>, +) -> Result, AiMuxError> { + let resolve_cc = |validator: &mut CacheControlValidator, part_opts| { + resolve_cache_control( + validator, + part_opts, + message_provider_options, + is_last_part, + "assistant message part", + "assistant message", + ) + }; + + Ok(Some(match part { + AssistantPart::Text(TextPart { + text, + provider_options, + }) => { + let cc = resolve_cc(validator, provider_options.as_ref()); + apply_cc(json!({ "type": "text", "text": text }), cc) } - ContentPart::Reasoning { + AssistantPart::File(file) => convert_file_part( + file, + betas, + validator, + message_provider_options, + is_last_part, + "assistant message part", + "assistant message", + )?, + + AssistantPart::Reasoning(ReasoningPart { text, - signature, provider_options, - } => { + }) => { return Ok(convert_reasoning_part( text, - signature.as_deref(), provider_options.as_ref(), send_reasoning, warnings, @@ -511,14 +570,14 @@ fn convert_part_to_anthropic( )); } - ContentPart::ToolCall { + AssistantPart::ToolCall(ToolCallPart { tool_call_id, tool_name, input, provider_executed, provider_options, .. - } => { + }) => { let cc = resolve_cc(validator, provider_options.as_ref()); if *provider_executed == Some(true) { @@ -635,52 +694,18 @@ fn convert_part_to_anthropic( apply_cc(block, cc) } - ContentPart::ToolResult { - tool_call_id, - result, - provider_options, - .. - } => { - let (content, is_error) = resolve_tool_result_output(result); - let mut block = json!({ - "type": "tool_result", - "tool_use_id": tool_call_id, - "content": content, - }); - if is_error { - block["is_error"] = json!(true); - } - // cache_control: part ?? output ?? (is_last_part ? message). - let cc = match validator.get_cache_control( - provider_options.as_ref().and_then(|o| o.get("anthropic")), - part_context_type, - true, - ) { - Some(v) => Some(v), - None => match extract_tool_result_output_provider_options(result) { - Some(out_opts) => validator.get_cache_control( - out_opts.get("anthropic").and_then(Value::as_object), - "tool result output", - true, - ), - None => { - if is_last_part { - validator.get_cache_control( - message_provider_options.and_then(|o| o.get("anthropic")), - message_context_type, - true, - ) - } else { - None - } - } - }, - }; - apply_cc(block, cc) + AssistantPart::ReasoningFile(_) | AssistantPart::Custom(_) => return Ok(None), + AssistantPart::ToolResult(part) => { + return Ok(convert_assistant_tool_result( + part, + tool_names, + mcp_tool_use_ids, + warnings, + validator, + is_last_part, + message_provider_options, + )); } - ContentPart::Custom { .. } - | ContentPart::ReasoningFile { .. } - | ContentPart::ToolApprovalRequest { .. } => return Ok(None), })) } @@ -695,15 +720,15 @@ fn convert_part_to_anthropic( fn collect_mcp_tool_use_ids(prompt: &LanguageModelPrompt) -> HashSet<&str> { let mut ids = HashSet::new(); for msg in prompt { - if msg.role != Role::Assistant { + let LanguageModelMessage::Assistant { content, .. } = msg else { continue; - } - for part in &msg.content { - if let ContentPart::ToolCall { + }; + for part in content { + if let AssistantPart::ToolCall(ToolCallPart { tool_call_id, provider_options: Some(opts), .. - } = part + }) = part && opts .get("anthropic") .and_then(|a| a.get("type")) @@ -736,7 +761,7 @@ fn result_error_code(value: &Value, fallback: &str) -> String { .to_string() } -/// Convert an assistant-role `ContentPart::ToolResult` into the matching +/// Convert an assistant-role `ToolResultPart` into the matching /// Anthropic provider-executed result block. /// /// Anthropic only accepts a bare `tool_result` block inside a **user** message; @@ -748,7 +773,7 @@ fn result_error_code(value: &Value, fallback: &str) -> String { /// /// Returns `None` when the part is skipped. fn convert_assistant_tool_result( - part: &ContentPart, + part: &ToolResultPart, tool_names: &ToolNameMapping, mcp_tool_use_ids: &HashSet<&str>, warnings: &mut Vec, @@ -756,17 +781,13 @@ fn convert_assistant_tool_result( is_last_part: bool, message_provider_options: Option<&SharedProviderOptions>, ) -> Option { - let ContentPart::ToolResult { + let ToolResultPart { tool_call_id, tool_name, - result, - is_error, + output, provider_options, .. - } = part - else { - return None; - }; + } = part; // cache_control: part ?? output ?? (is_last_part ? message) — the same // resolution order the bare `tool_result` path uses. @@ -776,12 +797,10 @@ fn convert_assistant_tool_result( true, ) { Some(v) => Some(v), - None => match extract_tool_result_output_provider_options(result) { - Some(out_opts) => validator.get_cache_control( - out_opts.get("anthropic").and_then(Value::as_object), - "tool result output", - true, - ), + None => match extract_tool_result_output_provider_options(output) { + Some(out_opts) => { + validator.get_cache_control(out_opts.get("anthropic"), "tool result output", true) + } None => { if is_last_part { validator.get_cache_control( @@ -796,14 +815,50 @@ fn convert_assistant_tool_result( }, }; - // The payload is carried bare, with the error flag alongside it. Upstream - // reads both off a `{ type, value }` envelope instead; reconstructing that - // envelope here would mean guessing, and a legitimate payload that happens - // to look like one (`{"type":"error","value":...}`) would be unwrapped and - // have its `is_error` overwritten. The envelope belongs on the type. - let value = result; - let is_error = is_error.unwrap_or(false); - let tool_name = tool_name.as_deref().unwrap_or_default(); + let raw_value = match output { + ToolResultOutput::Json { value, .. } | ToolResultOutput::ErrorJson { value, .. } => { + value.clone() + } + ToolResultOutput::Text { value, .. } | ToolResultOutput::ErrorText { value, .. } => { + json!(value) + } + ToolResultOutput::ExecutionDenied { reason, .. } => { + json!(reason.as_deref().unwrap_or("Tool call execution denied.")) + } + ToolResultOutput::Content { value } => json!(value), + }; + let provider_tool_name = tool_names.to_provider_tool_name(tool_name); + let allowed = match output { + ToolResultOutput::Json { .. } => true, + ToolResultOutput::ErrorJson { .. } => { + mcp_tool_use_ids.contains(tool_call_id.as_str()) + || matches!( + provider_tool_name, + "code_execution" | "web_fetch" | "web_search" | "advisor" + ) + } + ToolResultOutput::ErrorText { .. } => { + provider_tool_name == "code_execution" + && !mcp_tool_use_ids.contains(tool_call_id.as_str()) + } + _ => false, + }; + if !allowed { + warnings.push(Warning::Other { + message: format!( + "provider executed tool result output for tool {tool_name} is not supported" + ), + }); + return None; + } + let value = &raw_value; + let is_error = matches!( + output, + ToolResultOutput::ErrorText { .. } + | ToolResultOutput::ErrorJson { .. } + | ToolResultOutput::ExecutionDenied { .. } + ); + let tool_name = tool_name.as_str(); let payload_type = value.get("type").and_then(|t| t.as_str()); let mut block = if mcp_tool_use_ids.contains(tool_call_id.as_str()) { @@ -815,6 +870,23 @@ fn convert_assistant_tool_result( }) } else { let (block_type, content) = match tool_names.to_provider_tool_name(tool_name) { + "code_execution" if is_error => { + let error = if let Value::String(raw) = value { + serde_json::from_str(raw).unwrap_or(json!({})) + } else { + value.clone() + }; + let code_error = error.get("type").and_then(Value::as_str) + == Some("code_execution_tool_result_error"); + ( + if code_error { + "code_execution_tool_result" + } else { + "bash_code_execution_tool_result" + }, + json!({ "type": if code_error { "code_execution_tool_result_error" } else { "bash_code_execution_tool_result_error" }, "error_code": error.get("errorCode").and_then(Value::as_str).unwrap_or("unknown") }), + ) + } "code_execution" => match assistant_code_execution_content(value, payload_type) { Some(v) => v, None => { @@ -1034,21 +1106,31 @@ fn assistant_advisor_content(value: &Value, payload_type: Option<&str>) -> Value /// returned (e.g. `{ type: 'text', value, providerOptions }`). /// - When the output is a `content` output whose `value` is an array of content /// parts, the `providerOptions` of the first part that has one is returned. -fn extract_tool_result_output_provider_options(output: &Value) -> Option<&Value> { - if let Some(opts) = output.get("providerOptions") { - return Some(opts); - } - let output_type = output.get("type").and_then(|t| t.as_str()); - if output_type == Some("content") - && let Some(value) = output.get("value").and_then(|v| v.as_array()) - { - for part in value { - if let Some(opts) = part.get("providerOptions") { - return Some(opts); - } +fn extract_tool_result_output_provider_options( + output: &ToolResultOutput, +) -> Option<&SharedProviderOptions> { + match output { + ToolResultOutput::Text { + provider_options, .. + } + | ToolResultOutput::Json { + provider_options, .. } + | ToolResultOutput::ErrorText { + provider_options, .. + } + | ToolResultOutput::ErrorJson { + provider_options, .. + } + | ToolResultOutput::ExecutionDenied { + provider_options, .. + } => provider_options.as_ref(), + ToolResultOutput::Content { value } => value.iter().find_map(|part| match part { + ToolResultContent::Text(part) => part.provider_options.as_ref(), + ToolResultContent::File(part) => part.provider_options.as_ref(), + ToolResultContent::Custom { provider_options } => provider_options.as_ref(), + }), } - None } /// Route an inline-bytes file part. `full_media_type` is the resolved @@ -1220,7 +1302,10 @@ fn detect_document_media_type(bytes: &[u8]) -> Option<&'static str> { /// Resolve a file part's media type to a full `type/subtype` form, mirroring /// the TS `resolveFullMediaType`. When the media type is already a full type it /// is returned as-is; otherwise the subtype is sniffed from the inline bytes. -fn resolve_full_media_type(media_type: &str, bytes: &[u8]) -> Result { +pub(crate) fn resolve_full_media_type( + media_type: &str, + bytes: &[u8], +) -> Result { if is_full_media_type(media_type) { return Ok(media_type.to_string()); } @@ -1242,7 +1327,6 @@ fn resolve_full_media_type(media_type: &str, bytes: &[u8]) -> Result, provider_options: Option<&SharedProviderOptions>, send_reasoning: bool, warnings: &mut Vec, @@ -1262,14 +1346,10 @@ fn convert_reasoning_part( .and_then(|a| a.get("redactedData")) .and_then(|v| v.as_str()); - // #6: Fall back to `providerOptions.anthropic.signature` when the - // `signature` field is None (upstream convert-to-anthropic-prompt.ts:669-690). - let effective_signature = signature.or_else(|| { - provider_options - .and_then(|o| o.get("anthropic")) - .and_then(|a| a.get("signature")) - .and_then(|v| v.as_str()) - }); + let effective_signature = provider_options + .and_then(|o| o.get("anthropic")) + .and_then(|a| a.get("signature")) + .and_then(Value::as_str); if let Some(sig) = effective_signature { // Thinking blocks cannot carry cache_control directly — they are cached @@ -1305,23 +1385,22 @@ fn convert_reasoning_part( } /// Resolve the Anthropic file id from a provider-reference object, mirroring the -/// TS `resolveProviderReference`. Panics when no `anthropic` key is present, -/// matching the TS `UnsupportedFunctionalityError`. -fn resolve_anthropic_reference(reference: &Value) -> Result { - if let Some(id) = reference.get("anthropic").and_then(|v| v.as_str()) { +/// TS `resolveProviderReference`. Returns an error when the key is missing. +fn resolve_anthropic_reference( + reference: &std::collections::HashMap, +) -> Result { + if let Some(id) = reference.get("anthropic") { return Ok(id.to_string()); } - let providers: Vec<&str> = reference - .as_object() - .map(|o| o.keys().map(String::as_str).collect()) - .unwrap_or_default(); + let mut providers: Vec<&str> = reference.keys().map(String::as_str).collect(); + providers.sort_unstable(); Err(AiMuxError::InvalidArgument(format!( "No provider reference found for provider 'anthropic'. Available providers: {}", providers.join(", ") ))) } -/// Resolve a `ContentPart::ToolResult` `output` value into the Anthropic +/// Resolve a `ToolResultPart` `output` value into the Anthropic /// `tool_result.content` (and whether it is an error), mirroring the TS SDK. /// /// The V4 `tool-result` `output` is a discriminated `{ type, value }` object: @@ -1330,20 +1409,45 @@ fn resolve_anthropic_reference(reference: &Value) -> Result /// - `error` → `value` as-is with `is_error: true`, /// - `content` → `value` as-is (an array of content blocks), /// - anything else → the `output` is passed through unchanged. -fn resolve_tool_result_output(output: &Value) -> (Value, bool) { - if let (Some(t), Some(v)) = ( - output.get("type").and_then(|x| x.as_str()), - output.get("value"), - ) { - match t { - "json" => (Value::String(v.to_string()), false), - "text" => (v.clone(), false), - "error" => (v.clone(), true), - "content" => (v.clone(), false), - _ => (output.clone(), false), - } - } else { - (output.clone(), false) +fn resolve_tool_result_output( + output: &ToolResultOutput, + betas: &mut BTreeSet, + warnings: &mut Vec, +) -> (Value, bool) { + match output { + ToolResultOutput::Text { value, .. } => (json!(value), false), + ToolResultOutput::ErrorText { value, .. } => (json!(value), true), + ToolResultOutput::Json { value, .. } => (json!(value.to_string()), false), + ToolResultOutput::ErrorJson { value, .. } => (json!(value.to_string()), true), + ToolResultOutput::ExecutionDenied { reason, .. } => (json!(reason.as_deref().unwrap_or("Tool call execution denied.")), false), + ToolResultOutput::Content { value } => (json!(value.iter().filter_map(|part| match part { + ToolResultContent::Text(part) => Some(json!({ "type": "text", "text": part.text })), + ToolResultContent::File(part) => { + use base64::Engine; + let top_level = part.media_type.split('/').next(); + let source = match &part.data { + FileData::Url { url, .. } => json!({ "type": "url", "url": url }), + FileData::Data { data } => { + let (data, bytes) = match data { + FileBytes::Binary(bytes) => (base64::engine::general_purpose::STANDARD.encode(bytes), bytes.clone()), + FileBytes::Base64(data) => (data.clone(), base64::engine::general_purpose::STANDARD.decode(data).unwrap_or_default()), + }; + let full = resolve_full_media_type(&part.media_type, &bytes).ok()?; + if top_level != Some("image") && full != "application/pdf" { warnings.push(Warning::Other { message: "unsupported tool content file media type".to_string() }); return None; } + if full == "application/pdf" { betas.insert("pdfs-2024-09-25".to_string()); } + json!({ "type": "base64", "media_type": full, "data": data }) + }, + _ => { warnings.push(Warning::Other { message: "unsupported tool content file part".to_string() }); return None; }, + }; + Some(json!({ "type": if top_level == Some("image") { "image" } else { "document" }, "source": source })) + }, + ToolResultContent::Custom { provider_options } => { + if let Some(options) = provider_options.as_ref().and_then(|options| options.get("anthropic")) + && options.get("type").and_then(Value::as_str) == Some("tool-reference") { + Some(json!({ "type": "tool_reference", "tool_name": options.get("toolName") })) + } else { warnings.push(Warning::Other { message: "unsupported custom tool content part".to_string() }); None } + }, + }).collect::>()), false), } } diff --git a/aimux-providers/src/anthropic/mod.rs b/aimux-providers/src/anthropic/mod.rs index 5f2c5e39c..404c74750 100644 --- a/aimux-providers/src/anthropic/mod.rs +++ b/aimux-providers/src/anthropic/mod.rs @@ -12,6 +12,7 @@ pub mod types; pub mod usage; use std::collections::HashMap; +use std::sync::Arc; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; @@ -38,7 +39,7 @@ pub(crate) fn anthropic_failed_response_handler() const ANTHROPIC_API_URL: &str = "https://api.anthropic.com"; /// Configuration for the Anthropic provider. -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct AnthropicConfig { /// API key sent via the `x-api-key` header. Empty when authenticating with /// an `auth_token` instead. @@ -64,6 +65,25 @@ pub struct AnthropicConfig { /// (config_snapshot 记为 "explicit");`Some("env:VAR")` = 来自环境变量; /// `Some("none")` = 本地/匿名占位。不存明文之外的信息。 pub api_key_source: Option, + /// Generator for source IDs. Defaults to a 16-character alphanumeric ID. + pub generate_id: Option String + Send + Sync>>, +} + +impl std::fmt::Debug for AnthropicConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AnthropicConfig") + .field("api_key", &"") + .field("auth_token", &self.auth_token.is_some()) + .field("base_url", &self.base_url) + .field("api_version", &self.api_version) + .field("name", &self.name) + .field("headers", &self.headers) + .field("retry_config", &self.retry_config) + .field("body_overrides", &self.body_overrides) + .field("api_key_source", &self.api_key_source) + .field("generate_id", &self.generate_id.is_some()) + .finish() + } } impl AnthropicConfig { @@ -80,9 +100,20 @@ impl AnthropicConfig { retry_config: RetryConfig::default(), body_overrides: None, api_key_source: None, + generate_id: None, } } + /// Set the generator used for source IDs. + #[must_use] + pub fn with_generate_id(mut self, generate_id: F) -> Self + where + F: Fn() -> String + Send + Sync + 'static, + { + self.generate_id = Some(Arc::new(generate_id)); + self + } + /// Start a builder for more involved configurations (e.g. `auth_token`, /// custom `name`, extra `headers`). The builder validates that `api_key` /// and `auth_token` are not both provided. @@ -166,7 +197,7 @@ impl AnthropicConfig { /// Builder for [`AnthropicConfig`] supporting both `api_key` and `auth_token` /// authentication, with conflict validation mirroring the TS `createAnthropic`. -#[derive(Debug, Default, Clone)] +#[derive(Default, Clone)] pub struct AnthropicConfigBuilder { api_key: Option, auth_token: Option, @@ -174,6 +205,21 @@ pub struct AnthropicConfigBuilder { name: Option, headers: Option>, body_overrides: Option, + generate_id: Option String + Send + Sync>>, +} + +impl std::fmt::Debug for AnthropicConfigBuilder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AnthropicConfigBuilder") + .field("api_key", &self.api_key.is_some()) + .field("auth_token", &self.auth_token.is_some()) + .field("base_url", &self.base_url) + .field("name", &self.name) + .field("headers", &self.headers) + .field("body_overrides", &self.body_overrides) + .field("generate_id", &self.generate_id.is_some()) + .finish() + } } impl AnthropicConfigBuilder { @@ -213,6 +259,15 @@ impl AnthropicConfigBuilder { self } + #[must_use] + pub fn generate_id(mut self, generate_id: F) -> Self + where + F: Fn() -> String + Send + Sync + 'static, + { + self.generate_id = Some(Arc::new(generate_id)); + self + } + /// Build the config, validating that `api_key` and `auth_token` are not /// both set and that a provided `base_url` is non-empty. /// @@ -243,6 +298,7 @@ impl AnthropicConfigBuilder { retry_config: RetryConfig::default(), body_overrides: self.body_overrides, api_key_source: None, + generate_id: self.generate_id, }) } } diff --git a/aimux-providers/src/anthropic/model.rs b/aimux-providers/src/anthropic/model.rs index f09e83560..ef9155db7 100644 --- a/aimux-providers/src/anthropic/model.rs +++ b/aimux-providers/src/anthropic/model.rs @@ -22,12 +22,21 @@ use super::tool_name_mapping::ToolNameMapping; pub struct AnthropicModel { model_id: String, config: AnthropicConfig, + generate_id: std::sync::Arc String + Send + Sync>, } impl AnthropicModel { #[must_use] pub fn new(model_id: String, config: AnthropicConfig) -> Self { - Self { model_id, config } + let generate_id = config + .generate_id + .clone() + .unwrap_or_else(|| std::sync::Arc::new(super::stream::generate_id)); + Self { + model_id, + config, + generate_id, + } } fn endpoint(&self) -> String { @@ -133,6 +142,7 @@ impl LanguageModel for AnthropicModel { options.abort_signal.clone(), options.recording_context.clone(), &ToolNameMapping::new(options.tools.as_deref()), + self.generate_id.as_ref(), ) .await } @@ -151,6 +161,7 @@ impl LanguageModel for AnthropicModel { options.abort_signal.clone(), options.recording_context.clone(), ToolNameMapping::new(options.tools.as_deref()), + self.generate_id.clone(), ) .await } diff --git a/aimux-providers/src/anthropic/stream.rs b/aimux-providers/src/anthropic/stream.rs index 8d8dee1dc..92604e3dd 100644 --- a/aimux-providers/src/anthropic/stream.rs +++ b/aimux-providers/src/anthropic/stream.rs @@ -18,6 +18,7 @@ use aimux_core::tool::RawToolCall; use aimux_core::tool::ToolResult; use std::collections::HashMap; +use std::hash::{BuildHasher, RandomState}; use std::sync::atomic::{AtomicU64, Ordering}; use futures::StreamExt; @@ -126,16 +127,19 @@ fn str_field(v: &Value, key: &str) -> Option { v.get(key).and_then(|x| x.as_str()).map(str::to_string) } -/// Process-wide counter backing [`generate_source_id`]. -static SOURCE_ID_COUNTER: AtomicU64 = AtomicU64::new(0); - -/// Unique id for a `Source` derived from a web-search result. Upstream calls -/// `this.generateId()` at the same points. -fn generate_source_id() -> String { - format!( - "anthropic-source-{}", - SOURCE_ID_COUNTER.fetch_add(1, Ordering::Relaxed) - ) +/// Generate a non-cryptographic, 16-character alphanumeric ID, like `generateId`. +#[must_use] +pub(crate) fn generate_id() -> String { + const ALPHABET: &[u8] = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; + static COUNTER: AtomicU64 = AtomicU64::new(0); + let counter = COUNTER.fetch_add(1, Ordering::Relaxed); + let random = RandomState::new(); + (0..16) + .map(|index| { + let value = random.hash_one((counter, index)); + char::from(ALPHABET[(value % ALPHABET.len() as u64) as usize]) + }) + .collect() } /// `web_search_result` → the camel-cased shape upstream exposes (:1243-1249). @@ -404,6 +408,7 @@ pub(crate) fn stream_parts_for_result_block( names: &ToolNameMapping, mcp_tool_calls: &HashMap, server_tool_calls: &HashMap, + generate_id: &(dyn Fn() -> String + Send + Sync), ) -> Vec { let tool_result = |tool_name: String, (result, is_error): (Value, Option), tool_use_id: &str| { @@ -450,7 +455,7 @@ pub(crate) fn stream_parts_for_result_block( )]; parts.extend(results.iter().map(|result| { StreamPart::Source(Source::Url { - id: generate_source_id(), + id: generate_id(), url: str_field(result, "url").unwrap_or_default(), title: str_field(result, "title"), provider_metadata: Some( @@ -553,6 +558,7 @@ pub(crate) fn stream_parts_for_result_block( pub(crate) fn parse_anthropic_content( blocks: &[ContentBlock], names: &ToolNameMapping, + generate_id: &(dyn Fn() -> String + Send + Sync), ) -> Vec { let mut content = Vec::new(); // Result blocks inherit information from their matching calls. Index the @@ -686,7 +692,7 @@ pub(crate) fn parse_anthropic_content( // URLs and titles reach `result.sources`. for result in results { content.push(GenerateContent::Source(Source::Url { - id: generate_source_id(), + id: generate_id(), url: str_field(result, "url").unwrap_or_default(), title: str_field(result, "title"), provider_metadata: Some( @@ -850,6 +856,7 @@ pub(crate) async fn anthropic_generate_core( abort_signal: Option, recording_context: Option, tool_names: &ToolNameMapping, + generate_id: &(dyn Fn() -> String + Send + Sync), ) -> Result { let (request, request_body) = build_anthropic_request( endpoint, @@ -870,7 +877,7 @@ pub(crate) async fn anthropic_generate_core( let response_body = resp.raw_value; let data: AnthropicResponse = resp.value; - let content = parse_anthropic_content(&data.content, tool_names); + let content = parse_anthropic_content(&data.content, tool_names, generate_id); let finish_reason = data .stop_reason @@ -952,6 +959,7 @@ pub(crate) async fn anthropic_stream_core( abort_signal: Option, recording_context: Option, tool_names: ToolNameMapping, + generate_id: std::sync::Arc String + Send + Sync>, ) -> Result { let (request, request_body) = build_anthropic_request( endpoint, @@ -1153,6 +1161,7 @@ pub(crate) async fn anthropic_stream_core( &tool_names, &mcp_tool_calls, &server_tool_calls, + generate_id.as_ref(), ) { yield Ok(part); } diff --git a/aimux-providers/src/anthropic_aws/model.rs b/aimux-providers/src/anthropic_aws/model.rs index ec2d4d431..a1831e853 100644 --- a/aimux-providers/src/anthropic_aws/model.rs +++ b/aimux-providers/src/anthropic_aws/model.rs @@ -173,6 +173,7 @@ impl LanguageModel for AnthropicAwsModel { options.abort_signal.clone(), options.recording_context.clone(), &ToolNameMapping::new(options.tools.as_deref()), + &crate::anthropic::stream::generate_id, ) .await } @@ -190,6 +191,7 @@ impl LanguageModel for AnthropicAwsModel { options.abort_signal.clone(), options.recording_context.clone(), ToolNameMapping::new(options.tools.as_deref()), + std::sync::Arc::new(crate::anthropic::stream::generate_id), ) .await } diff --git a/aimux-providers/src/azure/responses.rs b/aimux-providers/src/azure/responses.rs index 31af23892..800cd428a 100644 --- a/aimux-providers/src/azure/responses.rs +++ b/aimux-providers/src/azure/responses.rs @@ -263,7 +263,7 @@ impl LanguageModel for AzureResponsesModel { async fn do_generate(&self, options: &CallOptions) -> Result { let headers = self.build_headers(options.headers.as_ref()).await?; - let request_result = build_responses_request_body(&self.deployment, options, false); + let request_result = build_responses_request_body(&self.deployment, options, false)?; let mut body = request_result.body; // Apply Azure assistant- file ID prefix passthrough. @@ -300,7 +300,7 @@ impl LanguageModel for AzureResponsesModel { async fn do_stream(&self, options: &CallOptions) -> Result { let headers = self.build_headers(options.headers.as_ref()).await?; - let request_result = build_responses_request_body(&self.deployment, options, true); + let request_result = build_responses_request_body(&self.deployment, options, true)?; let mut body = request_result.body; let warnings = request_result.warnings; let provider_key = provider_key().to_string(); diff --git a/aimux-providers/src/bedrock/convert.rs b/aimux-providers/src/bedrock/convert.rs index 9590540c3..a911d34c7 100644 --- a/aimux-providers/src/bedrock/convert.rs +++ b/aimux-providers/src/bedrock/convert.rs @@ -14,11 +14,12 @@ //! - Consecutive same-role messages are merged into a single message (matching //! the TS `groupIntoBlocks` behaviour). -use aimux_core::content::ContentPart; -use aimux_core::language_model_message::LanguageModelPrompt; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, UserPart, +}; use aimux_core::options::{CallOptions, ToolChoice}; -use aimux_core::shared::SharedProviderOptions; +use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified}; use base64::Engine; @@ -35,8 +36,14 @@ use serde_json::{Value, json}; /// messages into a single `assistant` block. Assistant text is trimmed when it /// is the last content part of the last message of the last block; empty /// assistant text is dropped unless the message also carries reasoning. -#[must_use] -pub fn convert_prompt_to_bedrock(prompt: &LanguageModelPrompt) -> (Vec, Vec) { +/// +/// # Errors +/// +/// Returns an unsupported-functionality error for unsupported tool result content. +pub fn convert_prompt_to_bedrock( + prompt: &LanguageModelPrompt, +) -> Result<(Vec, Vec), aimux_core::error::AiMuxError> { + validate_tool_result_content(prompt)?; #[derive(Clone, Copy, PartialEq, Eq)] enum Blk { System, @@ -47,10 +54,10 @@ pub fn convert_prompt_to_bedrock(prompt: &LanguageModelPrompt) -> (Vec, V // Group consecutive same-block messages (user+tool fold into User). let mut blocks: Vec<(Blk, Vec)> = Vec::new(); for (i, msg) in prompt.iter().enumerate() { - let b = match msg.role { - Role::System => Blk::System, - Role::User | Role::Tool => Blk::User, - Role::Assistant => Blk::Assistant, + let b = match msg { + LanguageModelMessage::System { .. } => Blk::System, + LanguageModelMessage::User { .. } | LanguageModelMessage::Tool { .. } => Blk::User, + LanguageModelMessage::Assistant { .. } => Blk::Assistant, }; if blocks.last().map(|(lb, _)| *lb) == Some(b) { blocks.last_mut().unwrap().1.push(i); @@ -69,49 +76,84 @@ pub fn convert_prompt_to_bedrock(prompt: &LanguageModelPrompt) -> (Vec, V match blk { Blk::System => { for &i in idxs { - let msg = &prompt[i]; - for p in &msg.content { - if let ContentPart::Text { text, .. } = p { - system.push(json!({ "text": text })); - } - if let Some(cp) = part_cache_point(p) { - system.push(cp); - } + let LanguageModelMessage::System { + content, + provider_options, + } = &prompt[i] + else { + unreachable!() + }; + system.push(json!({ "text": content })); + if let Some(cp) = cache_point(provider_options) { + system.push(cp); } } } Blk::User => { let mut content: Vec = Vec::new(); for &i in idxs { - let msg = &prompt[i]; - for p in &msg.content { - push_user_part(p, &mut content, &mut document_counter); - if let Some(cp) = part_cache_point(p) { - content.push(cp); + match &prompt[i] { + LanguageModelMessage::User { content: parts, .. } => { + for part in parts { + let provider_options = match part { + UserPart::Text(part) => { + content.push(json!({ "text": part.text })); + &part.provider_options + } + UserPart::File(file) => { + if !push_file_part( + file, + &mut content, + &mut document_counter, + ) { + continue; + } + &file.provider_options + } + }; + if let Some(cp) = cache_point(provider_options) { + content.push(cp); + } + } + } + LanguageModelMessage::Tool { content: parts, .. } => { + for part in parts { + let ToolPart::ToolResult(part) = part else { + continue; + }; + let result_content = + resolve_tool_result_output(&part.output, &mut document_counter); + content.push(json!({ + "toolResult": { + "toolUseId": part.tool_call_id, + "content": result_content, + } + })); + } } + _ => unreachable!(), } } - messages.push(json!({ "role": "user", "content": content })); + append_user_message(&mut messages, content); } Blk::Assistant => { let mut content: Vec = Vec::new(); let num_msgs = idxs.len(); for (mj, &i) in idxs.iter().enumerate() { let is_last_message = mj == num_msgs - 1; - let msg = &prompt[i]; - let has_reasoning = msg - .content + let LanguageModelMessage::Assistant { content: parts, .. } = &prompt[i] else { + unreachable!() + }; + let has_reasoning = parts .iter() - .any(|p| matches!(p, ContentPart::Reasoning { .. })); - let num_parts = msg.content.len(); - for (kj, p) in msg.content.iter().enumerate() { + .any(|p| matches!(p, AssistantPart::Reasoning(_))); + let num_parts = parts.len(); + for (kj, p) in parts.iter().enumerate() { let is_last_content_part = kj == num_parts - 1; match p { - ContentPart::Text { text, .. } => { + AssistantPart::Text(TextPart { text, .. }) => { // Skip empty text unless the message has reasoning. - if text.trim().is_empty() && !has_reasoning { - // skipped - } else { + if !text.trim().is_empty() || has_reasoning { let t = if is_last_block && is_last_message && is_last_content_part { @@ -122,11 +164,30 @@ pub fn convert_prompt_to_bedrock(prompt: &LanguageModelPrompt) -> (Vec, V content.push(json!({ "text": t })); } } - ContentPart::Reasoning { + AssistantPart::Reasoning(ReasoningPart { text, - signature: Some(sig), - .. - } => { + provider_options, + }) => { + let options = provider_options.as_ref().and_then(|options| { + options + .get("amazonBedrock") + .or_else(|| options.get("bedrock")) + }); + let Some(sig) = options + .and_then(|options| options.get("signature")) + .and_then(Value::as_str) + else { + if let Some(redacted) = + options.and_then(|options| options.get("redactedContent")) + { + content.push(json!({ "reasoningContent": { "redactedContent": redacted } })); + } else if let Some(redacted) = + options.and_then(|options| options.get("redactedData")) + { + content.push(json!({ "reasoningContent": { "redactedReasoning": { "data": redacted } } })); + } + continue; + }; // Only signed reasoning is replayed; unsigned // reasoning is intentionally omitted. content.push(json!({ @@ -138,13 +199,25 @@ pub fn convert_prompt_to_bedrock(prompt: &LanguageModelPrompt) -> (Vec, V } })); } - ContentPart::Reasoning { .. } => {} - ContentPart::ToolCall { + AssistantPart::ToolResult(part) => { + if !content.is_empty() { + messages.push(json!({ "role": "assistant", "content": std::mem::take(&mut content) })); + } + let mut result = vec![ + json!({ "toolResult": { "toolUseId": part.tool_call_id, "content": resolve_tool_result_output(&part.output, &mut document_counter) } }), + ]; + if let Some(cp) = cache_point(&part.provider_options) { + result.push(cp); + } + append_user_message(&mut messages, result); + continue; + } + AssistantPart::ToolCall(ToolCallPart { tool_call_id, tool_name, input, .. - } => { + }) => { let input_val = if input.is_object() { input.clone() } else { @@ -160,79 +233,187 @@ pub fn convert_prompt_to_bedrock(prompt: &LanguageModelPrompt) -> (Vec, V } _ => {} } - if let Some(cp) = part_cache_point(p) { + let provider_options = match p { + AssistantPart::Text(part) => &part.provider_options, + AssistantPart::File(part) + if matches!( + part.data, + FileData::Data { .. } | FileData::Text { .. } + ) => + { + &part.provider_options + } + _ => continue, + }; + if let Some(cp) = cache_point(provider_options) { content.push(cp); } } } - messages.push(json!({ "role": "assistant", "content": content })); + if !content.is_empty() { + messages.push(json!({ "role": "assistant", "content": content })); + } } } } - (system, messages) + Ok((system, messages)) } -/// Push a single user/tool content part into a Bedrock content array. -fn push_user_part(part: &ContentPart, content: &mut Vec, doc_counter: &mut u32) { - match part { - ContentPart::Text { text, .. } => { - content.push(json!({ "text": text })); - } - ContentPart::Image { - image, media_type, .. - } => { - let b64 = base64::engine::general_purpose::STANDARD.encode(image); - push_file_block(&b64, media_type, None, &None, content, doc_counter); - } - ContentPart::File { - data, - media_type, - filename, - provider_options, - } => { - let b64 = base64::engine::general_purpose::STANDARD.encode(data); - push_file_block( - &b64, - media_type, - filename.as_deref(), - provider_options, - content, - doc_counter, - ); - } - ContentPart::FileBase64 { - data, - media_type, - filename, - provider_options, - } => { - // `data` is already a base64 string — use verbatim as `bytes`. - push_file_block( - data, - media_type, - filename.as_deref(), - provider_options, - content, - doc_counter, - ); +fn append_user_message(messages: &mut Vec, content: Vec) { + if let Some(last) = messages.last_mut() + && last.get("role").and_then(Value::as_str) == Some("user") + && let Some(parts) = last.get_mut("content").and_then(Value::as_array_mut) + { + parts.extend(content); + } else { + messages.push(json!({ "role": "user", "content": content })); + } +} + +pub(crate) fn validate_tool_result_content( + prompt: &LanguageModelPrompt, +) -> Result<(), aimux_core::error::AiMuxError> { + for message in prompt { + let results: Vec<&ToolResultOutput> = match message { + LanguageModelMessage::Tool { content, .. } => content + .iter() + .filter_map(|part| { + if let ToolPart::ToolResult(part) = part { + Some(&part.output) + } else { + None + } + }) + .collect(), + LanguageModelMessage::Assistant { content, .. } => content + .iter() + .filter_map(|part| { + if let AssistantPart::ToolResult(part) = part { + Some(&part.output) + } else { + None + } + }) + .collect(), + _ => continue, + }; + for output in results { + if let ToolResultOutput::Content { value } = output { + for part in value { + let unsupported = match part { + ToolResultContent::Text(_) => false, + ToolResultContent::File(file) => { + let allowed_data = matches!(file.data, FileData::Data { .. }) + || matches!(&file.data, FileData::Url { url, .. } if url.starts_with("s3:")); + if !allowed_data { + return Err( + aimux_core::error::AiMuxError::UnsupportedFunctionality( + "tool result content part".to_string(), + ), + ); + } + let media_type = crate::google::convert::tool_file_media_type(file)?; + let supported = matches!( + media_type.as_str(), + "image/jpeg" + | "image/png" + | "image/gif" + | "image/webp" + | "video/x-matroska" + | "video/quicktime" + | "video/mp4" + | "video/webm" + | "video/x-flv" + | "video/mpeg" + | "video/mpg" + | "video/wmv" + | "video/x-ms-wmv" + | "video/3gpp" + | "application/pdf" + | "text/csv" + | "application/msword" + | "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + | "application/vnd.ms-excel" + | "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" + | "text/html" + | "text/plain" + | "text/markdown" + ); + if !supported { + let kind = match media_type.split('/').next() { + Some("image") => "image", + Some("video") => "video", + _ => "file", + }; + return Err( + aimux_core::error::AiMuxError::UnsupportedFunctionality( + format!("{kind} mime type: {media_type}"), + ), + ); + } + matches!(file.data, FileData::Url { .. }) + && !(media_type.starts_with("image/") + || media_type.starts_with("video/")) + } + ToolResultContent::Custom { .. } => true, + }; + if unsupported { + return Err(aimux_core::error::AiMuxError::UnsupportedFunctionality( + "tool result content part".to_string(), + )); + } + } + } } - ContentPart::ToolResult { - tool_call_id, - result, - .. - } => { - let result_content = resolve_tool_result_output(result, doc_counter); + } + Ok(()) +} + +/// Push a supported file into a Bedrock content array. +fn push_file_part(file: &FilePart, content: &mut Vec, doc_counter: &mut u32) -> bool { + let b64 = match &file.data { + FileData::Data { + data: FileBytes::Binary(bytes), + } => base64::engine::general_purpose::STANDARD.encode(bytes), + FileData::Data { + data: FileBytes::Base64(data), + } => data.clone(), + FileData::Text { text } => base64::engine::general_purpose::STANDARD.encode(text), + FileData::Url { url, .. } + if url.starts_with("s3:") && file.media_type.starts_with("image/") => + { content.push(json!({ - "toolResult": { - "toolUseId": tool_call_id, - "content": result_content, + "image": { + "format": mime_to_image_format(&file.media_type), + "source": { "s3Location": { "uri": url } } } })); + return true; } - // Variants not yet modelled for Bedrock; no test exercises these paths. - _ => {} - } + FileData::Url { .. } | FileData::Reference { .. } => return false, + }; + let is_text = matches!(file.data, FileData::Text { .. }); + let media_type = if is_text + && !file + .media_type + .split_once('/') + .is_some_and(|(_, subtype)| !subtype.is_empty() && subtype != "*") + { + "text/plain" + } else { + &file.media_type + }; + push_file_block( + &b64, + media_type, + file.filename.as_deref(), + &file.provider_options, + is_text, + content, + doc_counter, + ); + true } /// Build a Bedrock `image` or `document` block from an already-base64 `bytes` @@ -245,11 +426,12 @@ fn push_file_block( media_type: &str, filename: Option<&str>, provider_options: &Option, + is_text: bool, content: &mut Vec, doc_counter: &mut u32, ) { let top_level = media_type.split('/').next().unwrap_or(""); - if top_level == "image" { + if top_level == "image" && !is_text { let format = mime_to_image_format(media_type); content.push(json!({ "image": { "format": format, "source": { "bytes": b64 } } @@ -276,22 +458,8 @@ fn push_file_block( /// Extract a `{ cachePoint: {...} }` block from a part's `providerOptions` /// (`bedrock.cachePoint` or `amazonBedrock.cachePoint`), if present. -fn part_cache_point(part: &ContentPart) -> Option { - let po = match part { - ContentPart::Text { - provider_options, .. - } - | ContentPart::Image { - provider_options, .. - } - | ContentPart::File { - provider_options, .. - } - | ContentPart::FileBase64 { - provider_options, .. - } => provider_options.as_ref()?, - _ => return None, - }; +fn cache_point(provider_options: &Option) -> Option { + let po = provider_options.as_ref()?; for key in ["bedrock", "amazonBedrock"] { if let Some(cp) = po.get(key).and_then(|v| v.get("cachePoint")) { return Some(json!({ "cachePoint": cp.clone() })); @@ -321,57 +489,50 @@ fn citations_enabled(provider_options: &Option) -> bool { /// Resolve a tool result `output` value into Bedrock's `toolResult.content` /// array, mirroring the TS SDK. -fn resolve_tool_result_output(output: &Value, doc_counter: &mut u32) -> Vec { - let (t, v) = match ( - output.get("type").and_then(|x| x.as_str()), - output.get("value"), - ) { - (Some(t), Some(v)) => (t, v), - _ => return vec![json!({ "text": output.to_string() })], - }; - match t { - "json" | "error-json" => vec![json!({ "text": v.to_string() })], - "text" | "error" | "error-text" => vec![json!({ "text": v })], - "execution-denied" => { - let reason = v.as_str().unwrap_or("Tool call execution denied."); - vec![json!({ "text": reason })] +fn resolve_tool_result_output(output: &ToolResultOutput, doc_counter: &mut u32) -> Vec { + match output { + ToolResultOutput::Json { value, .. } | ToolResultOutput::ErrorJson { value, .. } => { + vec![json!({ "text": value.to_string() })] } - "content" => { - if let Some(arr) = v.as_array() { - let mut out = Vec::new(); - for part in arr { - convert_tool_result_content_part(part, &mut out, doc_counter); - } - out - } else { - vec![json!({ "text": output.to_string() })] - } - } - _ => vec![json!({ "text": output.to_string() })], - } -} - -/// Convert a single content part inside a tool-result `content` array. -fn convert_tool_result_content_part(part: &Value, content: &mut Vec, doc_counter: &mut u32) { - let part_type = part.get("type").and_then(|t| t.as_str()).unwrap_or(""); - match part_type { - "text" => { - let text = part.get("text").and_then(|t| t.as_str()).unwrap_or(""); - content.push(json!({ "text": text })); + ToolResultOutput::Text { value, .. } | ToolResultOutput::ErrorText { value, .. } => { + vec![json!({ "text": value })] } - "file" => { - let media_type = part.get("mediaType").and_then(|m| m.as_str()).unwrap_or(""); - // Tool-result file parts carry an already-base64 `data` string. - let b64 = part - .get("data") - .and_then(|d| d.get("data")) - .and_then(|d| d.as_str()) - .unwrap_or(""); - let filename = part.get("filename").and_then(|f| f.as_str()); - push_file_block(b64, media_type, filename, &None, content, doc_counter); + ToolResultOutput::ExecutionDenied { reason, .. } => { + vec![json!({ "text": reason.as_deref().unwrap_or("Tool call execution denied.") })] } - _ => { - content.push(json!({ "text": part.to_string() })); + ToolResultOutput::Content { value } => { + let mut content = Vec::new(); + for part in value { + match part { + ToolResultContent::Text(part) => content.push(json!({ "text": part.text })), + ToolResultContent::File(part) => { + let mut resolved = part.clone(); + if let Ok(media_type) = crate::google::convert::tool_file_media_type(part) { + resolved.media_type = media_type; + } + let part = &resolved; + if let FileData::Url { url, .. } = &part.data { + let format = part.media_type.split('/').nth(1).unwrap_or(""); + let source = json!({ "s3Location": { "uri": url } }); + content.push(if part.media_type.starts_with("image/") { json!({ "image": { "format": mime_to_image_format(&part.media_type), "source": source } }) } else { json!({ "video": { "format": format, "source": source } }) }); + } else if part.media_type.starts_with("video/") { + if let FileData::Data { data } = &part.data { + let bytes = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + content.push(json!({ "video": { "format": part.media_type.split('/').nth(1).unwrap_or(""), "source": { "bytes": bytes } } })); + } + } else { + push_file_part(part, &mut content, doc_counter); + } + } + ToolResultContent::Custom { .. } => {} + } + } + content } } } @@ -521,9 +682,15 @@ pub fn prepare_tools( /// and `inferenceConfig.maxTokens` is bumped by `budgetTokens`. The /// `reasoningConfig` key never appears in the request body. User-supplied /// `additionalModelRequestFields` are merged with the derived `thinking` field. -#[must_use] -pub fn build_request_body(model_id: &str, options: &CallOptions) -> Value { - let (system, messages) = convert_prompt_to_bedrock(&options.prompt); +/// +/// # Errors +/// +/// Returns an unsupported-functionality error for unsupported tool result content. +pub fn build_request_body( + model_id: &str, + options: &CallOptions, +) -> Result { + let (system, messages) = convert_prompt_to_bedrock(&options.prompt)?; // Extract Bedrock-specific provider options. let bedrock_opts = options @@ -616,7 +783,7 @@ pub fn build_request_body(model_id: &str, options: &CallOptions) -> Value { body.insert("toolConfig".to_string(), tool_config); } - Value::Object(body) + Ok(Value::Object(body)) } /// Map a Bedrock `stopReason` to the unified `FinishReason`. diff --git a/aimux-providers/src/bedrock/model.rs b/aimux-providers/src/bedrock/model.rs index a5c77492f..e88314a8a 100644 --- a/aimux-providers/src/bedrock/model.rs +++ b/aimux-providers/src/bedrock/model.rs @@ -197,7 +197,7 @@ impl LanguageModel for BedrockModel { } async fn do_generate(&self, options: &CallOptions) -> Result { - let body = build_request_body(&self.model_id, options); + let body = build_request_body(&self.model_id, options)?; let body_str = serde_json::to_string(&body).unwrap_or_default(); let url = self.endpoint(false); let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; @@ -268,7 +268,7 @@ impl LanguageModel for BedrockModel { } async fn do_stream(&self, options: &CallOptions) -> Result { - let body = build_request_body(&self.model_id, options); + let body = build_request_body(&self.model_id, options)?; let body_str = serde_json::to_string(&body).unwrap_or_default(); let url = self.endpoint(true); let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; diff --git a/aimux-providers/src/codex.rs b/aimux-providers/src/codex.rs index 64bb10f7f..5cf6b16f7 100644 --- a/aimux-providers/src/codex.rs +++ b/aimux-providers/src/codex.rs @@ -258,7 +258,7 @@ impl CodexModel { options: &CallOptions, stream: bool, ) -> Result<(Value, Vec), AiMuxError> { - let mut result = build_responses_request_body(&self.model_id, options, stream); + let mut result = build_responses_request_body(&self.model_id, options, stream)?; if self.config.mode == CodexMode::Subscription { result.body["store"] = Value::Bool(false); } diff --git a/aimux-providers/src/cohere/convert.rs b/aimux-providers/src/cohere/convert.rs index c10ffb8b9..0c2d866a2 100644 --- a/aimux-providers/src/cohere/convert.rs +++ b/aimux-providers/src/cohere/convert.rs @@ -3,12 +3,13 @@ //! Mirrors the TS `convert-to-cohere-chat-prompt.ts`, //! `cohere-prepare-tools.ts`, and `map-cohere-finish-reason.ts`. -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; -use aimux_core::language_model_message::LanguageModelPrompt; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolCallPart, ToolPart, + ToolResultPart, UserPart, +}; use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; -use aimux_core::shared::SharedProviderOptions; +use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; use aimux_core::tool::Tool; use aimux_core::types::{FinishReason, FinishReasonUnified, ReasoningEffort, Warning}; use serde_json::{Value, json}; @@ -120,72 +121,90 @@ pub struct ConvertedPrompt { /// - Non-image file parts are extracted as `documents` (RAG), not in the message. /// - Assistant content is a plain string; with tool calls, `content` is omitted. /// - Tool messages are flat `{role:"tool", content, tool_call_id}`. -#[must_use] -pub fn convert_prompt_to_cohere(prompt: &LanguageModelPrompt) -> ConvertedPrompt { +/// +/// # Errors +/// Returns an error for unsupported text file data. +pub fn convert_prompt_to_cohere( + prompt: &LanguageModelPrompt, +) -> Result { let mut messages = Vec::new(); let mut documents = Vec::new(); for msg in prompt { - match msg.role { - Role::System => { - let text = join_text_parts(&msg.content); - messages.push(json!({ "role": "system", "content": text })); + match msg { + LanguageModelMessage::System { content, .. } => { + messages.push(json!({ "role": "system", "content": content })); } - Role::User => { + LanguageModelMessage::User { content, .. } => { let mut parts: Vec = Vec::new(); let mut has_image = false; - for part in &msg.content { + for part in content { match part { - ContentPart::Text { text, .. } => { + UserPart::Text(TextPart { text, .. }) => { if !text.is_empty() { parts.push(json!({ "type": "text", "text": text })); } } - ContentPart::Image { - image, media_type, .. - } => { - has_image = true; + UserPart::File(file) => { use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(image); - parts.push(json!({ - "type": "image_url", - "image_url": { - "url": format!("data:{};base64,{}", media_type, b64), - } - })); - } - ContentPart::File { - data, - media_type, - filename, - .. - } => { - let top_level = media_type.split('/').next().unwrap_or(""); - if top_level == "image" { - has_image = true; - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(data); - parts.push(json!({ - "type": "image_url", - "image_url": { - "url": format!("data:{};base64,{}", media_type, b64), + if file.media_type.split('/').next() == Some("image") { + let url = match &file.data { + FileData::Data { data } => { + let b64 = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD + .encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + let media_type = + crate::google::convert::tool_file_media_type(file)?; + format!("data:{media_type};base64,{b64}") } - })); + FileData::Url { url, .. } => url.clone(), + FileData::Reference { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "image file parts with provider references".to_string(), + )); + } + FileData::Text { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "image file parts with text data".to_string(), + )); + } + }; + has_image = true; + parts.push( + json!({ "type": "image_url", "image_url": { "url": url } }), + ); } else { - // Non-image files become RAG documents. - // Mirrors TS: `documents.push({ data: { text, title: - // part.filename } })` — `title` is only present when a - // filename is supplied. - let text = String::from_utf8_lossy(data).to_string(); + let text = match &file.data { + FileData::Data { + data: FileBytes::Binary(bytes), + } => String::from_utf8_lossy(bytes).into_owned(), + FileData::Data { + data: FileBytes::Base64(data), + } => data.clone(), + FileData::Text { text } => text.clone(), + FileData::Url { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "File URL data: URLs should be downloaded by the AI SDK and not reach this point. This indicates a configuration issue.".to_string(), + )); + } + FileData::Reference { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "file parts with provider references".to_string(), + )); + } + }; let mut doc_data = json!({ "text": text }); - if let Some(fname) = filename { - doc_data["title"] = json!(fname); + if let Some(filename) = &file.filename { + doc_data["title"] = json!(filename); } documents.push(json!({ "data": doc_data })); } } - _ => {} } } @@ -208,29 +227,23 @@ pub fn convert_prompt_to_cohere(prompt: &LanguageModelPrompt) -> ConvertedPrompt messages.push(json!({ "role": "user", "content": text })); } } - Role::Assistant => { - let text = join_text_parts(&msg.content); - let has_tool_calls = msg - .content + LanguageModelMessage::Assistant { content, .. } => { + let text = join_text_parts(content); + let has_tool_calls = content .iter() - .any(|p| matches!(p, ContentPart::ToolCall { .. })); + .any(|p| matches!(p, AssistantPart::ToolCall(_))); if has_tool_calls { - let tool_calls: Vec = msg - .content + let tool_calls: Vec = content .iter() .filter_map(|p| match p { - ContentPart::ToolCall { + AssistantPart::ToolCall(ToolCallPart { tool_call_id, tool_name, input, .. - } => { - let arguments = if input.is_null() { - "{}".to_string() - } else { - input.to_string() - }; + }) => { + let arguments = input.to_string(); Some(json!({ "id": tool_call_id, "type": "function", @@ -251,50 +264,44 @@ pub fn convert_prompt_to_cohere(prompt: &LanguageModelPrompt) -> ConvertedPrompt messages.push(json!({ "role": "assistant", "content": text })); } } - Role::Tool => { - for part in &msg.content { - if let ContentPart::ToolResult { + LanguageModelMessage::Tool { content, .. } => { + for part in content { + let ToolPart::ToolResult(ToolResultPart { tool_call_id, - result, + output, .. - } = part - { - let content = tool_result_to_content(result); - messages.push(json!({ - "role": "tool", - "content": content, - "tool_call_id": tool_call_id, - })); - } + }) = part + else { + continue; + }; + let content = crate::openai::convert::tool_result_to_content(output); + messages.push(json!({ + "role": "tool", + "content": content, + "tool_call_id": tool_call_id, + })); } } } } - ConvertedPrompt { + Ok(ConvertedPrompt { messages, documents, - } + }) } -fn join_text_parts(content: &[ContentPart]) -> String { +fn join_text_parts(content: &[AssistantPart]) -> String { content .iter() .filter_map(|p| match p { - ContentPart::Text { text, .. } => Some(text.as_str()), + AssistantPart::Text(TextPart { text, .. }) => Some(text.as_str()), _ => None, }) .collect::>() .join("") } -fn tool_result_to_content(output: &Value) -> Value { - match output { - Value::String(s) => Value::String(s.clone()), - other => Value::String(other.to_string()), - } -} - // ── Request body ──────────────────────────────────────────────────────────── /// Result of building a Cohere request body, including warnings collected from @@ -312,13 +319,13 @@ pub struct RequestBodyResult { /// and the `thinking` config resolved from `reasoning` / provider options. /// /// # Errors -/// Returns a parse error for invalid Cohere provider options. +/// Returns an error for unsupported text file data. pub fn build_request_body( model_id: &str, options: &CallOptions, stream: bool, ) -> Result { - let converted = convert_prompt_to_cohere(&options.prompt); + let converted = convert_prompt_to_cohere(&options.prompt)?; let mut warnings: Vec = Vec::new(); let mut body = json!({ diff --git a/aimux-providers/src/google/convert.rs b/aimux-providers/src/google/convert.rs index 9c871ebe0..43fbca977 100644 --- a/aimux-providers/src/google/convert.rs +++ b/aimux-providers/src/google/convert.rs @@ -18,13 +18,14 @@ //! fields we actually read back are `text`, function tool parts, and native //! provider-executed tool parts. -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; -use aimux_core::language_model_message::LanguageModelPrompt; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, +}; use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; use aimux_core::result::{GenerateContent, Source}; -use aimux_core::shared::{JsonObject, SharedProviderOptions}; +use aimux_core::shared::{FileBytes, FileData, JsonObject, SharedProviderOptions}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified, Warning}; use base64::Engine; @@ -102,24 +103,25 @@ fn read_provider_options( /// as a string (JSON-stringified for non-string outputs, matching the TS /// `output.type === 'json'` path). /// -/// Thought signatures from per-part provider options are echoed as a +/// Thought signatures from provider options are echoed as a /// `thoughtSignature` sibling of the `functionCall` part. #[must_use] pub fn convert_to_google_messages(prompt: &LanguageModelPrompt) -> GooglePrompt { - convert_to_google_messages_for_namespace(prompt, ProviderMetadataNamespace::Google) + convert_to_google_messages_for_namespace(prompt, ProviderMetadataNamespace::Google, true) } fn convert_to_google_messages_for_namespace( prompt: &LanguageModelPrompt, namespace: ProviderMetadataNamespace, + supports_function_response_parts: bool, ) -> GooglePrompt { let mut system_parts: Vec = Vec::new(); let mut contents: Vec = Vec::new(); let mut system_messages_allowed = true; for msg in prompt { - match msg.role { - Role::System => { + match msg { + LanguageModelMessage::System { content, .. } => { if !system_messages_allowed { // The TS SDK throws `UnsupportedFunctionalityError` here: // system messages are only valid at the start of the @@ -131,31 +133,54 @@ fn convert_to_google_messages_for_namespace( // TODO: change the signature to `Result` to match TS semantics. continue; } - for part in &msg.content { - if let ContentPart::Text { text, .. } = part { - system_parts.push(json!({ "text": text })); - } - } + system_parts.push(json!({ "text": content })); } - Role::User => { + LanguageModelMessage::User { content, .. } => { system_messages_allowed = false; - let parts = convert_user_parts(&msg.content); + let parts = convert_user_parts(content, namespace); contents.push(json!({ "role": "user", "parts": parts })); } - Role::Assistant => { + LanguageModelMessage::Assistant { content, .. } => { system_messages_allowed = false; - let parts = convert_assistant_parts(&msg.content, namespace); + let parts = convert_assistant_parts(content, namespace); if !parts.is_empty() { contents.push(json!({ "role": "model", "parts": parts })); } } - Role::Tool => { + LanguageModelMessage::Tool { content, .. } => { system_messages_allowed = false; // Gemini folds tool results into a `user`-role message as // `functionResponse` parts. We emit them as their own user // turn (matching how the TS SDK pushes a new `{role:'user'}` // entry per tool-role message). - let parts = convert_tool_parts(&msg.content); + let mut ordinary = Vec::new(); + for part in content { + if let ToolPart::ToolResult(result) = part + && let Some(options) = + read_provider_options(result.provider_options.as_ref(), namespace) + && let (Some(server_id), Some(server_type)) = ( + options.get("serverToolCallId"), + options.get("serverToolType"), + ) + && let Some(last) = contents.last_mut() + && last.get("role").and_then(Value::as_str) == Some("model") + && let Some(parts) = last.get_mut("parts").and_then(Value::as_array_mut) + { + let response = match &result.output { + ToolResultOutput::Json { value, .. } => value.clone(), + _ => json!({}), + }; + let mut value = json!({ "toolResponse": { "toolType": server_type, "response": response, "id": server_id } }); + if let Some(signature) = options.get("thoughtSignature") { + value["thoughtSignature"] = signature.clone(); + } + parts.push(value); + } else { + ordinary.push(part.clone()); + } + } + let parts = + convert_tool_parts(&ordinary, namespace, supports_function_response_parts); if !parts.is_empty() { contents.push(json!({ "role": "user", "parts": parts })); } @@ -176,32 +201,61 @@ fn convert_to_google_messages_for_namespace( } /// Convert user-role content parts into Google parts. -fn convert_user_parts(content: &[ContentPart]) -> Vec { +fn convert_user_parts(content: &[UserPart], namespace: ProviderMetadataNamespace) -> Vec { let mut parts = Vec::new(); for part in content { match part { - ContentPart::Text { text, .. } => { + UserPart::Text(TextPart { text, .. }) => { parts.push(json!({ "text": text })); } - ContentPart::Image { - image, media_type, .. - } => { - let b64 = base64::engine::general_purpose::STANDARD.encode(image); - parts.push(json!({ - "inlineData": { "mimeType": media_type, "data": b64 } - })); - } - ContentPart::File { + UserPart::File(FilePart { data, media_type, .. - } => { - let b64 = base64::engine::general_purpose::STANDARD.encode(data); + }) => { + let inline_media_type = if matches!(data, FileData::Text { .. }) + && !media_type + .split_once('/') + .is_some_and(|(_, subtype)| !subtype.is_empty() && subtype != "*") + { + "text/plain" + } else { + media_type.as_str() + }; + let data = match data { + FileData::Data { + data: FileBytes::Binary(bytes), + } => base64::engine::general_purpose::STANDARD.encode(bytes), + FileData::Text { text } => { + base64::engine::general_purpose::STANDARD.encode(text.as_bytes()) + } + FileData::Data { + data: FileBytes::Base64(data), + } => data.clone(), + FileData::Url { url, original_url } => { + let uri = if url.starts_with("gs:") { + original_url.as_deref().unwrap_or(url) + } else { + url + }; + parts.push(json!({ + "fileData": { "mimeType": media_type, "fileUri": uri } + })); + continue; + } + FileData::Reference { reference } => { + if namespace == ProviderMetadataNamespace::Google + && let Some(reference) = reference.get("google") + { + parts.push(json!({ + "fileData": { "mimeType": media_type, "fileUri": reference } + })); + } + continue; + } + }; parts.push(json!({ - "inlineData": { "mimeType": media_type, "data": b64 } + "inlineData": { "mimeType": inline_media_type, "data": data } })); } - // Tool calls / results inside a user message are unusual; skip - // them rather than trying to stringify a `ContentPart`. - _ => {} } } if parts.is_empty() { @@ -215,16 +269,16 @@ fn convert_user_parts(content: &[ContentPart]) -> Vec { /// - `Text` → `{ text }` (skipped when empty, matching the TS SDK). /// - `ToolCall` → `{ functionCall: { id?, name, args } }`. fn convert_assistant_parts( - content: &[ContentPart], + content: &[AssistantPart], namespace: ProviderMetadataNamespace, ) -> Vec { let mut parts = Vec::new(); for part in content { match part { - ContentPart::Text { + AssistantPart::Text(TextPart { text, provider_options, - } => { + }) => { if !text.is_empty() { let mut p = json!({ "text": text }); // Echo thoughtSignature from provider_options if present @@ -238,53 +292,28 @@ fn convert_assistant_parts( parts.push(p); } } - ContentPart::Reasoning { + AssistantPart::Reasoning(ReasoningPart { text, - signature, provider_options, - } => { + }) => { if !text.is_empty() { let mut p = json!({ "text": text, "thought": true }); - // Prefer explicit signature field, then fall back to provider_options. - if let Some(sig) = signature.as_ref() { - p["thoughtSignature"] = json!(sig); - } else if let Some(sig) = - read_provider_options(provider_options.as_ref(), namespace) - .and_then(|g| g.get("thoughtSignature")) - .and_then(|v| v.as_str()) + if let Some(sig) = read_provider_options(provider_options.as_ref(), namespace) + .and_then(|g| g.get("thoughtSignature")) + .and_then(Value::as_str) { p["thoughtSignature"] = json!(sig); } parts.push(p); } } - ContentPart::ReasoningFile { - data: aimux_core::shared::GeneratedFileData::Data { data }, - media_type, - provider_options, - } => { - use base64::Engine; - let data = match data { - aimux_core::shared::FileBytes::Binary(bytes) => { - base64::engine::general_purpose::STANDARD.encode(bytes) - } - aimux_core::shared::FileBytes::Base64(data) => data.clone(), - }; - let mut value = json!({ "inlineData": { "mimeType": media_type, "data": data }, "thought": true }); - if let Some(signature) = read_provider_options(provider_options.as_ref(), namespace) - .and_then(|options| options.get("thoughtSignature")) - { - value["thoughtSignature"] = signature.clone(); - } - parts.push(value); - } - ContentPart::ToolCall { + AssistantPart::ToolCall(ToolCallPart { tool_call_id, tool_name, input, provider_options, .. - } => { + }) => { let google_options = read_provider_options(provider_options.as_ref(), namespace); let server_tool_call_id = google_options .and_then(|options| options.get("serverToolCallId")) @@ -294,7 +323,7 @@ fn convert_assistant_parts( .and_then(|value| value.as_str()); let signature = google_options .and_then(|options| options.get("thoughtSignature")) - .and_then(|value| value.as_str()); + .and_then(Value::as_str); let mut part_value = if let (Some(server_id), Some(server_type)) = (server_tool_call_id, server_tool_type) @@ -318,7 +347,7 @@ fn convert_assistant_parts( } } else { let mut function_call = Map::new(); - if !tool_call_id.is_empty() { + if !tool_call_id.is_empty() && namespace != ProviderMetadataNamespace::Vertex { function_call.insert("id".to_string(), json!(tool_call_id)); } function_call.insert("name".to_string(), json!(tool_name)); @@ -330,15 +359,21 @@ fn convert_assistant_parts( } parts.push(part_value); } - ContentPart::ToolResult { - tool_call_id: _, - tool_name: _, - result, - is_error: _, - preliminary: _, - dynamic: _, + AssistantPart::ToolResult(ToolResultPart { + output, + tool_name, provider_options, - } => { + .. + }) => { + let result = match output { + ToolResultOutput::Json { value, .. } => value.clone(), + _ => json!({}), + }; + if tool_name == "code_execution" && matches!(output, ToolResultOutput::Json { .. }) + { + parts.push(json!({ "codeExecutionResult": result })); + continue; + } // Provider-executed tool result in an assistant message — // upstream convert-to-google-messages.ts:518-540. // If it carries serverToolCallId + serverToolType, emit as @@ -368,9 +403,60 @@ fn convert_assistant_parts( } } } - _ => { - // Files / images in assistant messages: skip. + AssistantPart::ReasoningFile(part) => { + if let aimux_core::shared::GeneratedFileData::Data { data } = &part.data { + let data = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + let mut value = json!({ "inlineData": { "mimeType": part.media_type, "data": data }, "thought": true }); + if let Some(signature) = + read_provider_options(part.provider_options.as_ref(), namespace) + .and_then(|options| options.get("thoughtSignature")) + { + value["thoughtSignature"] = signature.clone(); + } + parts.push(value); + } } + AssistantPart::File(file) => { + let mut value = match &file.data { + FileData::Data { data } => { + let data = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + json!({ "inlineData": { "mimeType": file.media_type, "data": data } }) + } + FileData::Text { text } => { + json!({ "inlineData": { "mimeType": if file.media_type.split_once('/').is_some_and(|(_, subtype)| !subtype.is_empty() && subtype != "*") { file.media_type.as_str() } else { "text/plain" }, "data": base64::engine::general_purpose::STANDARD.encode(text.as_bytes()) } }) + } + FileData::Reference { reference } => { + if let Some(uri) = reference.get("google") { + json!({ "fileData": { "mimeType": file.media_type, "fileUri": uri } }) + } else { + continue; + } + } + FileData::Url { .. } => continue, + }; + if let Some(options) = + read_provider_options(file.provider_options.as_ref(), namespace) + { + if options.get("thought").and_then(Value::as_bool) == Some(true) { + value["thought"] = json!(true); + } + if let Some(signature) = options.get("thoughtSignature") { + value["thoughtSignature"] = signature.clone(); + } + } + parts.push(value); + } + AssistantPart::Custom(_) => {} } } parts @@ -382,40 +468,180 @@ fn convert_assistant_parts( /// `{ functionResponse: { id?, name, response: { name, content } } }`. /// `content` is the tool's output serialized to a string (string outputs /// pass through; JSON outputs are stringified). -fn convert_tool_parts(content: &[ContentPart]) -> Vec { +pub(crate) fn tool_file_media_type( + file: &FilePart, +) -> Result { + if let FileData::Data { data } = &file.data { + let bytes = match data { + FileBytes::Binary(bytes) => bytes.clone(), + FileBytes::Base64(data) => base64::engine::general_purpose::STANDARD + .decode(data) + .unwrap_or_default(), + }; + crate::anthropic::convert::resolve_full_media_type(&file.media_type, &bytes) + } else { + crate::anthropic::convert::resolve_full_media_type(&file.media_type, &[]) + } +} + +pub(crate) fn validate_tool_result_files( + prompt: &LanguageModelPrompt, +) -> Result<(), aimux_core::error::AiMuxError> { + for message in prompt { + if let LanguageModelMessage::Tool { content, .. } = message { + for part in content { + if let ToolPart::ToolResult(part) = part + && let ToolResultOutput::Content { value } = &part.output + { + for item in value { + if let ToolResultContent::File(file) = item + && matches!(file.data, FileData::Data { .. }) + { + tool_file_media_type(file)?; + } + } + } + } + } + } + Ok(()) +} + +fn contains_schema_reference(value: &Value) -> bool { + match value { + Value::Object(object) => { + object.contains_key("$ref") || object.values().any(contains_schema_reference) + } + Value::Array(array) => array.iter().any(contains_schema_reference), + _ => false, + } +} + +fn convert_tool_parts( + content: &[ToolPart], + namespace: ProviderMetadataNamespace, + supports_function_response_parts: bool, +) -> Vec { let mut parts = Vec::new(); for part in content { - if let ContentPart::ToolResult { + let ToolPart::ToolResult(ToolResultPart { tool_call_id, tool_name, - result, + output, .. - } = part - { - let content_str = match result { - Value::String(s) => s.clone(), - other => other.to_string(), - }; - let mut function_response = Map::new(); - if !tool_call_id.is_empty() { - function_response.insert("id".to_string(), json!(tool_call_id)); + }) = part + else { + continue; + }; + let response = |content: Value| { + let mut value = json!({ "functionResponse": { "name": tool_name, "response": { "name": tool_name, "content": content } } }); + if namespace != ProviderMetadataNamespace::Vertex && !tool_call_id.is_empty() { + value["functionResponse"]["id"] = json!(tool_call_id); } - // `name` is required by the API and must be the name of the tool - // that was called (Gemini pairs functionResponse with the prior - // functionCall by name; an empty or mismatched name is rejected - // with HTTP 400 on multi-turn tool calls). Prefer the framework - // provided `tool_name`, falling back to the call id for callers - // that never set one. - let name = tool_name - .clone() - .filter(|s| !s.is_empty()) - .unwrap_or_else(|| tool_call_id.clone()); - function_response.insert("name".to_string(), json!(name)); - function_response.insert( - "response".to_string(), - json!({ "name": name, "content": content_str }), - ); - parts.push(json!({ "functionResponse": function_response })); + value + }; + match output { + ToolResultOutput::Content { value } => { + if supports_function_response_parts { + let mut texts = Vec::new(); + let mut files = Vec::new(); + for part in value { + match part { + ToolResultContent::Text(part) => texts.push(part.text.clone()), + ToolResultContent::File(file) => match &file.data { + FileData::Data { data } => { + let data = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + files.push(json!({ "inlineData": { "mimeType": tool_file_media_type(file).unwrap_or_else(|_| file.media_type.clone()), "data": data } })); + } + FileData::Url { url, .. } if url.starts_with("data:") => { + if let Some((media_type, data)) = url + .strip_prefix("data:") + .and_then(|url| url.split_once(";base64,")) + { + files.push(json!({ "inlineData": { "mimeType": media_type, "data": data } })); + } else { + texts.push( + crate::openai::convert::tool_result_content_value( + std::slice::from_ref(part), + )[0] + .to_string(), + ); + } + } + FileData::Url { url, original_url } + if namespace == ProviderMetadataNamespace::Vertex + && url.starts_with("gs:") + && original_url + .as_deref() + .unwrap_or(url) + .starts_with("gs://") + && matches!( + file.media_type.as_str(), + "image/png" + | "image/jpeg" + | "image/webp" + | "application/pdf" + | "text/plain" + ) => + { + files.push(json!({ "fileData": { "mimeType": file.media_type, "fileUri": original_url.as_deref().unwrap_or(url) } })); + } + _ => texts.push( + crate::openai::convert::tool_result_content_value( + std::slice::from_ref(part), + )[0] + .to_string(), + ), + }, + _ => texts.push( + crate::openai::convert::tool_result_content_value( + std::slice::from_ref(part), + )[0] + .to_string(), + ), + } + } + let mut result = response(json!(if texts.is_empty() { + "Tool executed successfully.".to_string() + } else { + texts.join("\n") + })); + if !files.is_empty() { + result["functionResponse"]["parts"] = json!(files); + } + parts.push(result); + continue; + } + for part in value { + match part { + ToolResultContent::Text(part) => parts.push(response(json!(part.text))), + ToolResultContent::File(file) => if let FileData::Data { data } = &file.data { + let data = match data { FileBytes::Binary(bytes) => base64::engine::general_purpose::STANDARD.encode(bytes), FileBytes::Base64(data) => data.clone() }; + parts.push(json!({ "inlineData": { "mimeType": tool_file_media_type(file).unwrap_or_else(|_| file.media_type.clone()), "data": data } })); + parts.push(json!({ "text": format!("Tool executed successfully and returned this {} as a response", if file.media_type.starts_with("image/") { "image" } else { "file" }) })); + } else { parts.push(json!({ "text": crate::openai::convert::tool_result_content_value(std::slice::from_ref(part))[0].to_string() })); }, + _ => parts.push(json!({ "text": crate::openai::convert::tool_result_content_value(std::slice::from_ref(part))[0].to_string() })), + } + } + } + ToolResultOutput::Text { value, .. } | ToolResultOutput::ErrorText { value, .. } => { + parts.push(response(json!(value))) + } + ToolResultOutput::Json { value, .. } | ToolResultOutput::ErrorJson { value, .. } => { + parts.push(response(if contains_schema_reference(value) { + json!(value.to_string()) + } else { + value.clone() + })) + } + ToolResultOutput::ExecutionDenied { reason, .. } => parts.push(response(json!( + reason.as_deref().unwrap_or("Tool call execution denied.") + ))), } } parts @@ -1122,7 +1348,11 @@ fn build_request_body_with_warnings_for_namespace( let GooglePrompt { system_instruction, contents, - } = convert_to_google_messages_for_namespace(&options.prompt, namespace); + } = convert_to_google_messages_for_namespace( + &options.prompt, + namespace, + model_id.contains("gemini-3"), + ); let mut generation_config = Map::new(); diff --git a/aimux-providers/src/google/model.rs b/aimux-providers/src/google/model.rs index 1d976fcf2..712c51b38 100644 --- a/aimux-providers/src/google/model.rs +++ b/aimux-providers/src/google/model.rs @@ -126,15 +126,28 @@ impl LanguageModel for GoogleModel { if options .prompt .iter() - .filter(|message| message.role == aimux_core::message::Role::Assistant) - .flat_map(|message| &message.content) + .filter_map(|message| match message { + aimux_core::language_model_message::LanguageModelMessage::Assistant { + content, + .. + } => Some(content), + _ => None, + }) + .flatten() .any(|part| { matches!( part, - aimux_core::content::ContentPart::ReasoningFile { - data: aimux_core::shared::GeneratedFileData::Url { .. }, - .. - } + aimux_core::language_model_message::AssistantPart::ReasoningFile( + aimux_core::language_model_message::ReasoningFilePart { + data: aimux_core::shared::GeneratedFileData::Url { .. }, + .. + } + ) | aimux_core::language_model_message::AssistantPart::File( + aimux_core::language_model_message::FilePart { + data: aimux_core::shared::FileData::Url { .. }, + .. + } + ) ) }) { @@ -142,6 +155,7 @@ impl LanguageModel for GoogleModel { "File data URLs in assistant messages are not supported".to_string(), )); } + crate::google::convert::validate_tool_result_files(&options.prompt)?; let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); let (body, tool_warnings) = build_request_body_with_warnings(&self.model_id, options)?; let headers = self.build_headers(options.headers.as_ref()); @@ -224,15 +238,28 @@ impl LanguageModel for GoogleModel { if options .prompt .iter() - .filter(|message| message.role == aimux_core::message::Role::Assistant) - .flat_map(|message| &message.content) + .filter_map(|message| match message { + aimux_core::language_model_message::LanguageModelMessage::Assistant { + content, + .. + } => Some(content), + _ => None, + }) + .flatten() .any(|part| { matches!( part, - aimux_core::content::ContentPart::ReasoningFile { - data: aimux_core::shared::GeneratedFileData::Url { .. }, - .. - } + aimux_core::language_model_message::AssistantPart::ReasoningFile( + aimux_core::language_model_message::ReasoningFilePart { + data: aimux_core::shared::GeneratedFileData::Url { .. }, + .. + } + ) | aimux_core::language_model_message::AssistantPart::File( + aimux_core::language_model_message::FilePart { + data: aimux_core::shared::FileData::Url { .. }, + .. + } + ) ) }) { @@ -240,6 +267,7 @@ impl LanguageModel for GoogleModel { "File data URLs in assistant messages are not supported".to_string(), )); } + crate::google::convert::validate_tool_result_files(&options.prompt)?; let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); let (body, tool_warnings) = build_request_body_with_warnings(&self.model_id, options)?; let headers = self.build_headers(options.headers.as_ref()); diff --git a/aimux-providers/src/huggingface/responses.rs b/aimux-providers/src/huggingface/responses.rs index b75da8726..adb27505e 100644 --- a/aimux-providers/src/huggingface/responses.rs +++ b/aimux-providers/src/huggingface/responses.rs @@ -23,15 +23,16 @@ use base64::Engine; use futures::StreamExt; use serde_json::{Map, Value, json}; -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; use aimux_core::error::ApiCallError; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::LanguageModelPrompt; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, UserPart, +}; use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; use aimux_core::shared::provider_namespace; +use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; use aimux_core::tool::Tool; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage, Warning}; @@ -732,74 +733,59 @@ pub fn convert_to_huggingface_responses_messages( let mut warnings: Vec = Vec::new(); for msg in prompt { - match msg.role { - Role::System => { - // System content is a plain string (concatenation of text parts). - let content: String = msg - .content - .iter() - .filter_map(|p| match p { - ContentPart::Text { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join(""); + match msg { + LanguageModelMessage::System { content, .. } => { messages.push(json!({ "role": "system", "content": content })); } - - Role::User => { + LanguageModelMessage::User { content, .. } => { let mut parts: Vec = Vec::new(); - for part in &msg.content { + for part in content { match part { - ContentPart::Text { text, .. } => { + UserPart::Text(TextPart { text, .. }) => { parts.push(json!({ "type": "input_text", "text": text })); } - ContentPart::FileBase64 { - data, media_type, .. - } => { - parts.push(convert_file_part_base64(media_type, data)?); - } - ContentPart::FileUrl { - url, media_type, .. - } => { - parts.push(convert_file_part_url(media_type, url)?); - } - ContentPart::FileReference { .. } => { - return Err(AiMuxError::UnsupportedFunctionality( - "file parts with provider references".into(), - )); - } - ContentPart::Image { - image, media_type, .. - } => { - let encoded = base64::engine::general_purpose::STANDARD.encode(image); - parts.push(convert_file_part_base64(media_type, &encoded)?); - } - ContentPart::File { - data, media_type, .. - } => { - let encoded = base64::engine::general_purpose::STANDARD.encode(data); - parts.push(convert_file_part_base64(media_type, &encoded)?); - } - _ => { - // Other part types (tool-call, tool-result, reasoning) - // are not expected in user messages — skip. + UserPart::File(file) => { + parts.push(match &file.data { + FileData::Data { + data: FileBytes::Binary(bytes), + } => { + let encoded = + base64::engine::general_purpose::STANDARD.encode(bytes); + convert_file_part_base64(&file.media_type, &encoded)? + } + FileData::Data { + data: FileBytes::Base64(data), + } => convert_file_part_base64(&file.media_type, data)?, + FileData::Url { url, .. } => { + convert_file_part_url(&file.media_type, url)? + } + FileData::Reference { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "file parts with provider references".into(), + )); + } + FileData::Text { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "text file parts".into(), + )); + } + }); } } } messages.push(json!({ "role": "user", "content": parts })); } - Role::Assistant => { - for part in &msg.content { + LanguageModelMessage::Assistant { content, .. } => { + for part in content { match part { - ContentPart::Text { text, .. } => { + AssistantPart::Text(TextPart { text, .. }) => { messages.push(json!({ "role": "assistant", "content": [{ "type": "output_text", "text": text }] })); } - ContentPart::Reasoning { text, .. } => { + AssistantPart::Reasoning(ReasoningPart { text, .. }) => { messages.push(json!({ "role": "assistant", "content": [{ "type": "output_text", "text": text }] @@ -807,13 +793,13 @@ pub fn convert_to_huggingface_responses_messages( } // Tool calls and tool results are handled by the // Responses API — skip (no warning, matching TS). - ContentPart::ToolCall { .. } | ContentPart::ToolResult { .. } => {} + AssistantPart::ToolCall(_) | AssistantPart::ToolResult(_) => {} _ => {} } } } - Role::Tool => { + LanguageModelMessage::Tool { .. } => { warnings.push(Warning::Unsupported { feature: "tool messages".into(), details: None, diff --git a/aimux-providers/src/mistral/convert.rs b/aimux-providers/src/mistral/convert.rs index 0f6c3a5eb..683392aa2 100644 --- a/aimux-providers/src/mistral/convert.rs +++ b/aimux-providers/src/mistral/convert.rs @@ -3,11 +3,13 @@ //! Mirrors the TS `convert-to-mistral-chat-messages.ts`, //! `mistral-prepare-tools.ts`, and `map-mistral-finish-reason.ts`. -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolCallPart, ToolPart, + ToolResultOutput, ToolResultPart, UserPart, +}; use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; +use aimux_core::shared::{FileBytes, FileData}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified}; use serde_json::{Value, json}; @@ -102,54 +104,68 @@ pub fn prepare_tools( /// message if it is an assistant message (continuation mode). /// - Tool messages include `tool_call_id` (no `name` — the Rust data model /// does not carry the tool name on `ToolResult` parts). -#[must_use] -pub fn convert_prompt_to_mistral_messages(prompt: &LanguageModelPrompt) -> Vec { +/// +/// # Errors +/// Returns an error for unsupported text file data. +pub fn convert_prompt_to_mistral_messages( + prompt: &LanguageModelPrompt, +) -> Result, AiMuxError> { let mut result = Vec::new(); let last_idx = prompt.len().saturating_sub(1); for (i, msg) in prompt.iter().enumerate() { let is_last = i == last_idx; - for value in convert_message_to_mistral(msg, is_last) { + for value in convert_message_to_mistral(msg, is_last)? { result.push(value); } } - result + Ok(result) } -fn convert_message_to_mistral(msg: &LanguageModelPromptMessage, is_last: bool) -> Vec { - match msg.role { - Role::System => { - let text = join_text_parts(&msg.content); - vec![json!({ "role": "system", "content": text })] +fn convert_message_to_mistral( + msg: &LanguageModelMessage, + is_last: bool, +) -> Result, AiMuxError> { + Ok(match msg { + LanguageModelMessage::System { content, .. } => { + vec![json!({ "role": "system", "content": content })] } - Role::User => { - let parts: Vec = msg.content.iter().map(convert_part_to_mistral).collect(); + LanguageModelMessage::User { content, .. } => { + let parts: Vec = content + .iter() + .map(convert_part_to_mistral) + .collect::>()?; vec![json!({ "role": "user", "content": parts })] } - Role::Assistant => { - let text = join_text_parts(&msg.content); - let has_tool_calls = msg - .content + LanguageModelMessage::Assistant { content, .. } => { + let text = join_text_parts(content); + let has_reasoning = content.iter().any(|part| matches!(part, AssistantPart::Reasoning(_))); + let mut content_parts = Vec::new(); + for part in content { + match part { + AssistantPart::Text(part) => content_parts.push(json!({ "type": "text", "text": part.text })), + AssistantPart::Reasoning(part) => content_parts.push(json!({ "type": "thinking", "thinking": [{ "type": "text", "text": part.text }], "closed": true })), + AssistantPart::ToolCall(_) => {}, + _ => return Err(AiMuxError::UnsupportedFunctionality("assistant content part".to_string())), + } + } + let has_tool_calls = content .iter() - .any(|p| matches!(p, ContentPart::ToolCall { .. })); + .any(|p| matches!(p, AssistantPart::ToolCall(_))); let mut msg_json = json!({ "role": "assistant", "content": text }); + if has_reasoning { msg_json["content"] = json!(content_parts); } if has_tool_calls { - let tool_calls: Vec = msg - .content + let tool_calls: Vec = content .iter() .filter_map(|p| match p { - ContentPart::ToolCall { + AssistantPart::ToolCall(ToolCallPart { tool_call_id, tool_name, input, .. - } => { - let arguments = if input.is_null() { - "{}".to_string() - } else { - input.to_string() - }; + }) => { + let arguments = input.to_string(); Some(json!({ "id": tool_call_id, "type": "function", @@ -170,115 +186,92 @@ fn convert_message_to_mistral(msg: &LanguageModelPromptMessage, is_last: bool) - } vec![msg_json] } - Role::Tool => msg - .content + LanguageModelMessage::Tool { content, .. } => content .iter() - .filter_map(|part| match part { - ContentPart::ToolResult { - tool_call_id, - result, - .. - } => { - let content = tool_result_to_content(result); - Some(json!({ - "role": "tool", - "tool_call_id": tool_call_id, - "content": content, - })) - } - _ => None, + .filter_map(|part| { + let ToolPart::ToolResult(ToolResultPart { tool_call_id, tool_name, output, .. }) = part else { return None; }; + Some(json!({ "role": "tool", "name": tool_name, "tool_call_id": tool_call_id, "content": tool_result_to_content(output) })) }) .collect(), - } + }) } -fn join_text_parts(content: &[ContentPart]) -> String { +fn join_text_parts(content: &[AssistantPart]) -> String { content .iter() .filter_map(|p| match p { - ContentPart::Text { text, .. } => Some(text.as_str()), + AssistantPart::Text(TextPart { text, .. }) => Some(text.as_str()), _ => None, }) .collect::>() .join("") } -fn tool_result_to_content(output: &Value) -> Value { - match output { - Value::String(s) => Value::String(s.clone()), - other => Value::String(other.to_string()), - } +fn tool_result_to_content(output: &ToolResultOutput) -> Value { + Value::String(match output { + ToolResultOutput::Text { value, .. } | ToolResultOutput::ErrorText { value, .. } => { + value.clone() + } + ToolResultOutput::ExecutionDenied { reason, .. } => reason + .clone() + .unwrap_or_else(|| "Tool call execution denied.".to_string()), + ToolResultOutput::Json { value, .. } | ToolResultOutput::ErrorJson { value, .. } => { + value.to_string() + } + ToolResultOutput::Content { value } => { + crate::openai::convert::tool_result_content_value(value).to_string() + } + }) } -fn convert_part_to_mistral(part: &ContentPart) -> Value { - match part { - ContentPart::Text { text, .. } => json!({ "type": "text", "text": text }), - ContentPart::Image { - image, media_type, .. - } => { +fn convert_part_to_mistral(part: &UserPart) -> Result { + Ok(match part { + UserPart::Text(TextPart { text, .. }) => json!({ "type": "text", "text": text }), + UserPart::File(file) => { use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(image); - json!({ - "type": "image_url", - "image_url": format!("data:{};base64,{}", media_type, b64), - }) - } - ContentPart::File { - data, media_type, .. - } => { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(data); - let top_level = media_type.split('/').next().unwrap_or(""); - if top_level == "image" { - json!({ - "type": "image_url", - "image_url": format!("data:{};base64,{}", media_type, b64), - }) + let is_image = file.media_type.split('/').next() == Some("image"); + let url = match &file.data { + FileData::Data { data } => { + let b64 = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + let media_type = crate::google::convert::tool_file_media_type(file)?; + if !is_image && media_type != "application/pdf" { + return Err(AiMuxError::UnsupportedFunctionality( + "Only images and PDF file parts are supported".to_string(), + )); + } + format!("data:{media_type};base64,{b64}") + } + FileData::Url { url, .. } => { + if !is_image && file.media_type != "application/pdf" { + return Err(AiMuxError::UnsupportedFunctionality( + "Only images and PDF file parts are supported".to_string(), + )); + } + url.clone() + } + FileData::Reference { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "file parts with provider references".to_string(), + )); + } + FileData::Text { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "text file parts".to_string(), + )); + } + }; + if is_image { + json!({ "type": "image_url", "image_url": url }) } else { - json!({ - "type": "document_url", - "document_url": format!("data:{};base64,{}", media_type, b64), - }) + json!({ "type": "document_url", "document_url": url }) } } - // These variants are handled by `convert_message_to_mistral` for - // assistant/tool roles. - ContentPart::ToolCall { - tool_call_id, - tool_name, - input, - .. - } => { - json!({ - "type": "tool_call", - "id": tool_call_id, - "function": { - "name": tool_name, - "arguments": input.to_string(), - } - }) - } - ContentPart::ToolResult { - tool_call_id, - result, - .. - } => { - json!({ - "type": "tool_result", - "tool_call_id": tool_call_id, - "content": result, - }) - } - - // Variants not yet modelled for Mistral; no test exercises these paths. - ContentPart::FileBase64 { .. } - | ContentPart::FileUrl { .. } - | ContentPart::FileReference { .. } - | ContentPart::Reasoning { .. } - | ContentPart::Custom { .. } - | ContentPart::ReasoningFile { .. } - | ContentPart::ToolApprovalRequest { .. } => Value::Null, - } + }) } // ── Request body ──────────────────────────────────────────────────────────── @@ -286,7 +279,7 @@ fn convert_part_to_mistral(part: &ContentPart) -> Value { /// Convert `CallOptions` to a Mistral request body. /// /// # Errors -/// Returns a parse error for invalid Mistral provider options. +/// Returns an error for unsupported text file data. pub fn build_request_body( model_id: &str, options: &CallOptions, @@ -313,7 +306,7 @@ pub fn build_request_body( }) .transpose()? .unwrap_or_default(); - let messages = convert_prompt_to_mistral_messages(&options.prompt); + let messages = convert_prompt_to_mistral_messages(&options.prompt)?; let mut body = json!({ "model": model_id, diff --git a/aimux-providers/src/open_responses.rs b/aimux-providers/src/open_responses.rs index 304ab3770..c684cc6eb 100644 --- a/aimux-providers/src/open_responses.rs +++ b/aimux-providers/src/open_responses.rs @@ -14,15 +14,17 @@ use async_trait::async_trait; use futures::StreamExt; use serde_json::{Map, Value, json}; -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; use aimux_core::error::ApiCallError; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::LanguageModelPrompt; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolCallPart, + ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, +}; use aimux_core::options::{CallOptions, ToolChoice}; use aimux_core::provider::Provider; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; +use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; use aimux_core::tool::Tool; use aimux_core::types::{ @@ -312,6 +314,7 @@ impl LanguageModel for OpenResponsesModel { async fn do_generate(&self, options: &CallOptions) -> Result { let headers = self.build_headers(options.headers.as_ref()); + validate_tool_output_media(&options.prompt)?; let (body, warnings) = build_request_body(&self.model_id, options, &self.config.provider_options_name); @@ -455,6 +458,7 @@ impl LanguageModel for OpenResponsesModel { async fn do_stream(&self, options: &CallOptions) -> Result { let headers = self.build_headers(options.headers.as_ref()); + validate_tool_output_media(&options.prompt)?; let (body, warnings) = build_request_body(&self.model_id, options, &self.config.provider_options_name); @@ -903,7 +907,8 @@ fn build_request_body( } // Convert prompt to input + instructions. - let (input, instructions, input_warnings) = convert_to_open_responses_input(&options.prompt); + let (input, instructions, input_warnings) = + convert_to_open_responses_input_with_namespace(&options.prompt, provider_options_name); warnings.extend(input_warnings); // Convert function tools. @@ -1063,87 +1068,152 @@ fn build_request_body( #[must_use] pub fn convert_to_open_responses_input( prompt: &LanguageModelPrompt, +) -> (Value, Option, Vec) { + convert_to_open_responses_input_with_namespace(prompt, "open-responses") +} + +fn convert_to_open_responses_input_with_namespace( + prompt: &LanguageModelPrompt, + provider_options_name: &str, ) -> (Value, Option, Vec) { let mut input: Vec = Vec::new(); let mut warnings = Vec::new(); let mut system_messages: Vec = Vec::new(); for msg in prompt { - match msg.role { - Role::System => { - for part in &msg.content { - if let ContentPart::Text { text, .. } = part { - system_messages.push(text.clone()); - } - } + match msg { + LanguageModelMessage::System { content, .. } => { + system_messages.push(content.clone()); } - Role::User => { - let user_content = convert_user_content(&msg.content, &mut warnings); + LanguageModelMessage::User { content, .. } => { + let user_content = convert_user_content(content, &mut warnings); input.push(json!({ "type": "message", "role": "user", "content": user_content, })); } - Role::Assistant => { + LanguageModelMessage::Assistant { content, .. } => { let mut assistant_content: Vec = Vec::new(); - let mut tool_calls: Vec = Vec::new(); - for part in &msg.content { + for part in content { match part { - ContentPart::Text { text, .. } => { + AssistantPart::Text(TextPart { text, .. }) => { assistant_content.push(json!({ "type": "output_text", "text": text, })); } - ContentPart::ToolCall { + AssistantPart::ToolCall(ToolCallPart { tool_call_id, tool_name, input: tool_input, .. - } => { + }) => { let arguments = match tool_input { Value::String(s) => s.clone(), other => other.to_string(), }; - tool_calls.push(json!({ + flush_assistant_content(&mut assistant_content, &mut input); + input.push(json!({ "type": "function_call", "call_id": tool_call_id, "name": tool_name, "arguments": arguments, })); } + AssistantPart::Reasoning(part) => { + flush_assistant_content(&mut assistant_content, &mut input); + let metadata = part + .provider_options + .as_ref() + .and_then(|options| options.get(provider_options_name)); + let parse_parts = |key: &str, kind: &str| { + metadata + .and_then(|data| data.get(key)) + .and_then(Value::as_array) + .filter(|parts| { + parts.iter().all(|part| { + part["type"] == kind && part["text"].is_string() + }) + }) + .map(|parts| { + json!( + parts + .iter() + .map( + |part| json!({"type":kind, "text":part["text"]}) + ) + .collect::>() + ) + }) + }; + let mut reasoning = json!({"type":"reasoning", "summary":parse_parts("reasoningSummary", "summary_text").unwrap_or_else(|| json!([]))}); + if let Some(id) = metadata + .and_then(|data| data.get("itemId")) + .and_then(Value::as_str) + { + reasoning["id"] = json!(id); + } + if let Some(content) = parse_parts("reasoningContent", "reasoning_text") + { + reasoning["content"] = content; + } else if !metadata + .is_some_and(|data| data.contains_key("reasoningContent")) + && !part.text.is_empty() + { + reasoning["content"] = + json!([{"type":"reasoning_text", "text":part.text}]); + } + if let Some(encrypted) = metadata + .and_then(|data| data.get("reasoningEncryptedContent")) + .and_then(Value::as_str) + { + reasoning["encrypted_content"] = json!(encrypted); + } + if let Some(previous) = input.last_mut().filter(|previous| { + reasoning.get("id").is_some() + && previous["type"] == "reasoning" + && previous.get("id") == reasoning.get("id") + }) { + if let Some(content) = + reasoning.get("content").and_then(Value::as_array) + { + if previous.get("content").is_none() { + previous["content"] = json!([]); + } + if let Some(previous_content) = + previous["content"].as_array_mut() + { + previous_content.extend(content.iter().cloned()); + } + } + } else { + input.push(reasoning); + } + } _ => {} } } - - if !assistant_content.is_empty() { - input.push(json!({ - "type": "message", - "role": "assistant", - "content": assistant_content, - })); - } - for tc in tool_calls { - input.push(tc); - } + flush_assistant_content(&mut assistant_content, &mut input); } - Role::Tool => { - for part in &msg.content { - if let ContentPart::ToolResult { + LanguageModelMessage::Tool { content, .. } => { + for part in content { + let ToolPart::ToolResult(ToolResultPart { tool_call_id, - result, + output, .. - } = part - { - let content_value = resolve_tool_result_output(result, &mut warnings); - input.push(json!({ - "type": "function_call_output", - "call_id": tool_call_id, - "output": content_value, - })); - } + }) = part + else { + continue; + }; + let content_value = + resolve_tool_result_output(output, provider_options_name, &mut warnings); + input.push(json!({ + "type": "function_call_output", + "call_id": tool_call_id, + "output": content_value, + })); } } } @@ -1158,183 +1228,207 @@ pub fn convert_to_open_responses_input( (json!(input), instructions, warnings) } +fn flush_assistant_content(content: &mut Vec, input: &mut Vec) { + if !content.is_empty() { + input + .push(json!({"type":"message", "role":"assistant", "content":std::mem::take(content)})); + } +} + /// Convert user message content parts to the Open Responses format. -fn convert_user_content(content: &[ContentPart], warnings: &mut Vec) -> Value { +fn convert_user_content(content: &[UserPart], warnings: &mut Vec) -> Value { + use base64::Engine; + let mut parts: Vec = Vec::new(); for part in content { match part { - ContentPart::Text { text, .. } => { + UserPart::Text(TextPart { text, .. }) => { parts.push(json!({ "type": "input_text", "text": text })); } - ContentPart::Image { - image, media_type, .. - } => { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(image); - parts.push(json!({ - "type": "input_image", - "image_url": format!("data:{};base64,{}", media_type, b64), - })); - } - ContentPart::File { + UserPart::File(FilePart { data, media_type, filename, .. - } => { - let top_level = top_level_media_type(media_type); - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(data); - if top_level == "image" { - parts.push(json!({ - "type": "input_image", - "image_url": format!("data:{};base64,{}", media_type, b64), - })); - } else { - parts.push(json!({ - "type": "input_file", - "filename": filename.as_deref().unwrap_or("data"), - "file_data": format!("data:{};base64,{}", media_type, b64), - })); + }) => { + let image = top_level_media_type(media_type) == "image"; + match data { + FileData::Data { data } => { + let b64 = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + let data_url = format!("data:{media_type};base64,{b64}"); + parts.push(if image { + json!({ "type": "input_image", "image_url": data_url }) + } else { + json!({ + "type": "input_file", + "filename": filename.as_deref().unwrap_or("data"), + "file_data": data_url, + }) + }); + } + FileData::Url { url, .. } => { + parts.push(if image { + json!({ "type": "input_image", "image_url": url }) + } else { + json!({ "type": "input_file", "file_url": url }) + }); + } + FileData::Reference { .. } => { + warnings.push(Warning::Other { + message: "unsupported file part with provider reference".to_string(), + }); + } + FileData::Text { .. } => { + warnings.push(Warning::Other { + message: "unsupported text file part".to_string(), + }); + } } } - ContentPart::FileBase64 { - data, - media_type, - filename, + } + } + json!(parts) +} + +fn validate_tool_output_media(prompt: &LanguageModelPrompt) -> Result<(), AiMuxError> { + use base64::Engine; + for message in prompt { + let LanguageModelMessage::Tool { content, .. } = message else { + continue; + }; + for part in content { + let ToolPart::ToolResult(ToolResultPart { + output: ToolResultOutput::Content { value }, .. - } => { - let top_level = top_level_media_type(media_type); - if top_level == "image" { - parts.push(json!({ - "type": "input_image", - "image_url": format!("data:{};base64,{}", media_type, data), - })); - } else { - parts.push(json!({ - "type": "input_file", - "filename": filename.as_deref().unwrap_or("data"), - "file_data": format!("data:{};base64,{}", media_type, data), - })); - } - } - ContentPart::FileUrl { - url, media_type, .. - } => { - let top_level = top_level_media_type(media_type); - if top_level == "image" { - parts.push(json!({ - "type": "input_image", - "image_url": url, - })); - } else { - parts.push(json!({ - "type": "input_file", - "file_url": url, - })); - } - } - ContentPart::FileReference { .. } => { - warnings.push(Warning::Other { - message: "unsupported file part with provider reference".to_string(), - }); - } - _ => { - warnings.push(Warning::Other { - message: format!("unsupported content part type: {}", part_variant_name(part)), - }); + }) = part + else { + continue; + }; + for item in value { + let ToolResultContent::File(FilePart { + data: FileData::Data { data }, + media_type, + .. + }) = item + else { + continue; + }; + let bytes = match data { + FileBytes::Binary(bytes) => std::borrow::Cow::Borrowed(bytes.as_slice()), + FileBytes::Base64(data) => std::borrow::Cow::Owned( + base64::engine::general_purpose::STANDARD + .decode(data) + .map_err(|error| { + AiMuxError::InvalidArgument(format!( + "invalid base64 file data: {error}" + )) + })?, + ), + }; + crate::anthropic::convert::resolve_full_media_type(media_type, &bytes)?; } } } - json!(parts) + Ok(()) } /// Resolve a tool-result `output` value into the Open Responses `output` /// field, mirroring the TS convert logic. -fn resolve_tool_result_output(output: &Value, warnings: &mut Vec) -> Value { - let output_type = output.get("type").and_then(|x| x.as_str()); - - match output_type { - Some("text") | Some("error-text") => { - output.get("value").cloned().unwrap_or_else(|| Value::Null) - } - Some("execution-denied") => output - .get("reason") - .and_then(|r| r.as_str()) - .map(|s| Value::String(s.to_string())) - .unwrap_or_else(|| Value::String("Tool call execution denied.".to_string())), - Some("json") | Some("error-json") => { - let v = output.get("value").unwrap_or(&Value::Null); - Value::String(v.to_string()) - } - Some("content") => { - let v = output.get("value"); - if let Some(arr) = v.and_then(|v| v.as_array()) { - let mut parts: Vec = Vec::new(); - for item in arr { - match item.get("type").and_then(|t| t.as_str()) { - Some("text") => { - if let Some(text) = item.get("text").and_then(|t| t.as_str()) { - parts.push(json!({ "type": "input_text", "text": text })); +fn resolve_tool_result_output( + output: &ToolResultOutput, + provider_options_name: &str, + warnings: &mut Vec, +) -> Value { + let ToolResultOutput::Content { value } = output else { + return crate::openai::convert::tool_result_to_content(output); + }; + let mut parts = Vec::new(); + for item in value { + match item { + ToolResultContent::Text(part) => { + parts.push(json!({"type":"input_text", "text":part.text})) + } + ToolResultContent::File(part) => { + let image = top_level_media_type(&part.media_type) == "image"; + let mut converted = match &part.data { + FileData::Data { data } => { + use base64::Engine; + let bytes = match data { + FileBytes::Binary(bytes) => bytes.clone(), + FileBytes::Base64(data) => base64::engine::general_purpose::STANDARD + .decode(data) + .unwrap_or_default(), + }; + let media_type = match crate::anthropic::convert::resolve_full_media_type( + &part.media_type, + &bytes, + ) { + Ok(media_type) => media_type, + Err(error) => { + warnings.push(Warning::Other { + message: error.to_string(), + }); + continue; } - } - Some("file") => { - if let Some(media_type) = item.get("mediaType").and_then(|m| m.as_str()) - { - let top_level = top_level_media_type(media_type); - if let Some(data) = item.get("data") { - if let Some(data_obj) = - data.get("data").and_then(|d| d.as_str()) - { - if top_level == "image" { - parts.push(json!({ - "type": "input_image", - "image_url": format!("data:{};base64,{}", media_type, data_obj), - })); - } else { - parts.push(json!({ - "type": "input_file", - "filename": item.get("filename").and_then(|f| f.as_str()).unwrap_or("data"), - "file_data": format!("data:{};base64,{}", media_type, data_obj), - })); - } - } else if let Some(url) = - data.get("url").and_then(|u| u.as_str()) - { - if top_level == "image" { - parts.push(json!({ - "type": "input_image", - "image_url": url, - })); - } else { - parts.push(json!({ - "type": "input_file", - "file_url": url, - })); - } - } - } + }; + let b64 = match data { + FileBytes::Binary(_) => { + base64::engine::general_purpose::STANDARD.encode(&bytes) } + FileBytes::Base64(data) => data.clone(), + }; + let data_url = format!("data:{media_type};base64,{b64}"); + if image { + json!({"type":"input_image", "image_url":data_url}) + } else { + json!({"type":"input_file", "filename":part.filename.as_deref().unwrap_or("data"), "file_data":data_url}) } - _ => { - warnings.push(Warning::Other { - message: format!( - "unsupported tool content part type: {}", - item.get("type") - .and_then(|t| t.as_str()) - .unwrap_or("unknown") - ), - }); + } + FileData::Url { url, .. } => { + if image { + json!({"type":"input_image", "image_url":url}) + } else { + json!({"type":"input_file", "file_url":url}) } } + data => { + warnings.push(Warning::Other { + message: format!( + "unsupported tool content part type: file with data type: {}", + if matches!(data, FileData::Reference { .. }) { + "reference" + } else { + "text" + } + ), + }); + continue; + } + }; + if image { + let detail = part + .provider_options + .as_ref() + .and_then(|options| options.get(provider_options_name)) + .and_then(|options| options.get("imageDetail")) + .and_then(Value::as_str) + .filter(|detail| matches!(*detail, "low" | "high" | "auto")) + .unwrap_or("auto"); + converted["detail"] = json!(detail); } - Value::Array(parts) - } else { - output.clone() + parts.push(converted); } + ToolResultContent::Custom { .. } => warnings.push(Warning::Other { + message: "unsupported tool content part type: custom".into(), + }), } - _ => output.clone(), } + json!(parts) } /// Extract the top-level media type (e.g. "image" from "image/png"). @@ -1342,24 +1436,6 @@ fn top_level_media_type(media_type: &str) -> &str { media_type.split('/').next().unwrap_or("") } -/// Get a human-readable name for a ContentPart variant (for warnings). -fn part_variant_name(part: &ContentPart) -> &'static str { - match part { - ContentPart::Text { .. } => "text", - ContentPart::Image { .. } => "image", - ContentPart::File { .. } => "file", - ContentPart::FileBase64 { .. } => "file-base64", - ContentPart::FileUrl { .. } => "file-url", - ContentPart::FileReference { .. } => "file-reference", - ContentPart::Reasoning { .. } => "reasoning", - ContentPart::Custom { .. } => "custom", - ContentPart::ReasoningFile { .. } => "reasoning-file", - ContentPart::ToolApprovalRequest { .. } => "tool-approval-request", - ContentPart::ToolCall { .. } => "tool-call", - ContentPart::ToolResult { .. } => "tool-result", - } -} - // == Usage extraction == /// Extract `Usage` from a response body `Value`. diff --git a/aimux-providers/src/openai/convert.rs b/aimux-providers/src/openai/convert.rs index 363a93f10..f79d608f0 100644 --- a/aimux-providers/src/openai/convert.rs +++ b/aimux-providers/src/openai/convert.rs @@ -1,11 +1,12 @@ //! Conversion between `LanguageModelPrompt` and OpenAI API format. -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolPart, + ToolResultContent, ToolResultOutput, UserPart, +}; use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; -use aimux_core::shared::{JsonObject, SharedProviderOptions}; +use aimux_core::shared::{FileBytes, FileData, JsonObject, SharedProviderOptions}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified, ReasoningEffort, Warning}; use serde::Serialize; @@ -269,6 +270,15 @@ pub fn convert_prompt_to_openai_messages_with_provider( prompt: &LanguageModelPrompt, system_message_mode: SystemMessageMode, provider: &str, +) -> Result, AiMuxError> { + convert_prompt_with_warnings(prompt, system_message_mode, provider, &mut Vec::new()) +} + +fn convert_prompt_with_warnings( + prompt: &LanguageModelPrompt, + system_message_mode: SystemMessageMode, + provider: &str, + warnings: &mut Vec, ) -> Result, AiMuxError> { let mut result = Vec::new(); for msg in prompt { @@ -276,6 +286,7 @@ pub fn convert_prompt_to_openai_messages_with_provider( msg, system_message_mode, provider, + warnings, )?); } Ok(result) @@ -332,48 +343,63 @@ fn resolve_full_media_type(media_type: &str, b64_data: &str) -> String { } /// Resolve a provider reference, throwing if the provider is not found. -fn resolve_provider_reference(reference: &Value, provider: &str) -> Result { - if let Some(val) = reference.get(provider) { - if let Some(s) = val.as_str() { - return Ok(s.to_string()); - } - return Ok(val.to_string()); - } - let available: Vec = reference - .as_object() - .map(|m| m.keys().cloned().collect()) - .unwrap_or_default(); - Err(format!( - "No provider reference found for provider '{}'. Available providers: {}", - provider, - available.join(", ") - )) +fn resolve_provider_reference( + reference: &std::collections::HashMap, + provider: &str, +) -> Result { + reference.get(provider).cloned().ok_or_else(|| { + let mut available: Vec<&str> = reference.keys().map(String::as_str).collect(); + available.sort_unstable(); + format!( + "No provider reference found for provider '{}'. Available providers: {}", + provider, + available.join(", ") + ) + }) } /// Convert a file part to the OpenAI format, handling images, audio, and PDF. fn convert_file_part_to_openai( - media_type: &str, - data_b64: Option<&str>, - url: Option<&str>, - reference: Option<&Value>, - filename: Option<&str>, - provider_options: &Option, + file: &FilePart, part_index: usize, + provider: &str, ) -> Result { - let prompt_cache_breakpoint = get_prompt_cache_breakpoint(provider_options); - - // Reference type - if let Some(ref_val) = reference { - let file_id = resolve_provider_reference(ref_val, "openai")?; - let mut part = json!({ - "type": "file", - "file": { "file_id": file_id } - }); - if let Some(bpt) = prompt_cache_breakpoint { - part["prompt_cache_breakpoint"] = bpt; + use base64::Engine; + + let FilePart { + data, + media_type, + filename, + provider_options, + } = file; + let prompt_cache_breakpoint = (provider == "openai") + .then(|| get_prompt_cache_breakpoint(provider_options)) + .flatten(); + let (data_b64, url) = match data { + FileData::Reference { reference } => { + let file_id = resolve_provider_reference(reference, provider)?; + let mut part = if provider == "deepseek" { + json!({ "type": "file", "file_id": file_id }) + } else { + json!({ "type": "file", "file": { "file_id": file_id } }) + }; + if let Some(bpt) = prompt_cache_breakpoint { + part["prompt_cache_breakpoint"] = bpt; + } + return Ok(part); } - return Ok(part); - } + FileData::Text { .. } => return Err("text file parts".to_string()), + FileData::Url { url, .. } => (None, Some(url.as_str())), + FileData::Data { data } => { + let b64 = match data { + FileBytes::Binary(bytes) => base64::engine::general_purpose::STANDARD.encode(bytes), + FileBytes::Base64(data) => data.clone(), + }; + (Some(b64), None) + } + }; + let data_b64 = data_b64.as_deref(); + let filename = filename.as_deref(); let top_level = get_top_level_media_type(media_type); @@ -459,438 +485,414 @@ fn convert_file_part_to_openai( /// Convert a single provider-facing message into one or more OpenAI messages. fn convert_message_to_openai( - msg: &LanguageModelPromptMessage, + msg: &LanguageModelMessage, system_message_mode: SystemMessageMode, provider: &str, + warnings: &mut Vec, ) -> Result, AiMuxError> { - let role = match msg.role { - Role::System => "system", - Role::User => "user", - Role::Assistant => "assistant", - Role::Tool => "tool", - }; - - // System messages: respect systemMessageMode. - if msg.role == Role::System { - match system_message_mode { - SystemMessageMode::Remove => return Ok(vec![]), - SystemMessageMode::Developer => { - return Ok(convert_system_message(msg, "developer")); - } - SystemMessageMode::System => { - return Ok(convert_system_message(msg, "system")); - } + if provider != "openai" && provider != "deepseek" { + let has_reference = match msg { + LanguageModelMessage::User { content, .. } => content.iter().any(|part| { + matches!( + part, + UserPart::File(FilePart { + data: FileData::Reference { .. }, + .. + }) + ) + }), + LanguageModelMessage::Assistant { content, .. } => content.iter().any(|part| { + matches!( + part, + AssistantPart::File(FilePart { + data: FileData::Reference { .. }, + .. + }) + ) + }), + _ => false, + }; + if has_reference { + return Err(AiMuxError::UnsupportedFunctionality( + "file parts with provider references".into(), + )); } } - - // Tool-role messages: each ToolResult part becomes its own OpenAI message. - if msg.role == Role::Tool { - let messages = msg - .content - .iter() - .filter_map(|part| match part { - ContentPart::ToolResult { - tool_call_id, - result, - .. - } => { - let content = tool_result_to_content(result); - Some(json!({ - "role": "tool", - "content": content, - "tool_call_id": tool_call_id, - })) - } - _ => None, - }) - .collect(); - return Ok(messages); - } - - // Assistant messages with tool calls - let has_tool_calls = msg - .content - .iter() - .any(|p| matches!(p, ContentPart::ToolCall { .. })); - - // Groq: assistant messages always collect text, reasoning, and tool_calls - // together (mirrors TS `convertToGroqChatMessages`). The `reasoning` field - // is added when non-empty, and `content` is "" (not null) when there is no - // text. - if msg.role == Role::Assistant && provider == "groq" { - let text: String = msg - .content - .iter() - .filter_map(|p| match p { - ContentPart::Text { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join(""); - - let reasoning: String = msg - .content - .iter() - .filter_map(|p| match p { - ContentPart::Reasoning { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join(""); - - let tool_calls_json: Vec = msg - .content - .iter() - .filter_map(|p| match p { - ContentPart::ToolCall { - tool_call_id, - tool_name, - input, - .. - } => { - let arguments = if input.is_null() { - "{}".to_string() - } else { - input.to_string() - }; - Some(json!({ - "type": "function", - "id": tool_call_id, - "function": { - "name": tool_name, - "arguments": arguments, - } - })) + let message = match msg { + LanguageModelMessage::System { + content, + provider_options, + } => { + let role = match system_message_mode { + SystemMessageMode::Remove => return Ok(vec![]), + SystemMessageMode::Developer => "developer", + SystemMessageMode::System => "system", + }; + let content = match (provider == "openai") + .then(|| get_prompt_cache_breakpoint(provider_options)) + .flatten() + { + None => json!(content), + Some(bpt) => json!([{ + "type": "text", "text": content, "prompt_cache_breakpoint": bpt, + }]), + }; + json!({ "role": role, "content": content }) + } + LanguageModelMessage::Tool { content, .. } => { + let mut messages = Vec::new(); + for part in content { + let ToolPart::ToolResult(result) = part else { + continue; + }; + let mut content = if provider == "deepseek" { + deepseek_tool_result_content(&result.output, warnings)? + } else { + tool_result_to_content(&result.output) + }; + if provider == "openai" + && let Some(breakpoint) = tool_result_cache_breakpoint(&result.output) + .or_else(|| get_prompt_cache_breakpoint(&result.provider_options)) + { + content = json!([{"type":"text", "text":content, "prompt_cache_breakpoint":breakpoint}]); } - _ => None, - }) - .collect(); - - let mut msg_obj = json!({ - "role": "assistant", - "content": text, - }); - if !reasoning.is_empty() { - msg_obj["reasoning"] = json!(reasoning); + messages.push( + json!({"role":"tool", "content":content, "tool_call_id":result.tool_call_id}), + ); + } + return Ok(messages); } - if !tool_calls_json.is_empty() { - msg_obj["tool_calls"] = json!(tool_calls_json); + LanguageModelMessage::User { content, .. } => { + let all_plain_text = content.len() == 1 && content.iter().all(|part| { + matches!(part, UserPart::Text(text) + if provider != "openai" || get_prompt_cache_breakpoint(&text.provider_options).is_none()) + }); + let content = if all_plain_text { + json!( + content + .iter() + .filter_map(|part| match part { + UserPart::Text(text) => Some(text.text.as_str()), + UserPart::File(_) => None, + }) + .collect::() + ) + } else { + json!( + content + .iter() + .enumerate() + .map(|(index, part)| match part { + UserPart::Text(text) => Ok(convert_text_part_to_openai(text, provider)), + UserPart::File(file) => + convert_file_part_to_openai(file, index, provider) + .map_err(AiMuxError::InvalidArgument), + }) + .collect::, _>>()? + ) + }; + json!({ "role": "user", "content": content }) } - return Ok(vec![msg_obj]); - } - - if msg.role == Role::Assistant && has_tool_calls { - let text: String = msg - .content - .iter() - .filter_map(|p| match p { - ContentPart::Text { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join(""); - - // Reasoning / thinking content. DeepSeek V4 thinking mode (and other - // OpenAI-compatible reasoning models such as xAI) require prior - // assistant `reasoning_content` to be replayed on subsequent turns, - // including tool-call turns. Mirrors the Vercel AI SDK - // `openai-compatible` assistant conversion, which emits - // `reasoning_content` whenever a reasoning part is present. Groq uses - // the `reasoning` field name and is handled in its own branch above. - let reasoning: String = collect_reasoning(&msg.content); - - let tool_calls_json: Vec = msg - .content - .iter() - .filter_map(|p| match p { - ContentPart::ToolCall { - tool_call_id, - tool_name, - input, - .. - } => { - let arguments = if input.is_null() { - "{}".to_string() - } else { - input.to_string() - }; - Some(json!({ - "type": "function", - "id": tool_call_id, - "function": { - "name": tool_name, - "arguments": arguments, - } - })) + LanguageModelMessage::Assistant { content, .. } => { + let mut text = String::new(); + let mut reasoning = String::new(); + let mut tool_calls = Vec::new(); + for part in content { + match part { + AssistantPart::Text(part) => text.push_str(&part.text), + AssistantPart::Reasoning(part) => reasoning.push_str(&part.text), + AssistantPart::ToolCall(part) => { + let arguments = if provider == "openai" && !part.input.is_object() { + "{}".to_string() + } else { + part.input.to_string() + }; + tool_calls.push(json!({ + "type": "function", "id": part.tool_call_id, + "function": { "name": part.tool_name, "arguments": arguments }, + })); + } + AssistantPart::File(_) + | AssistantPart::ToolResult(_) + | AssistantPart::ReasoningFile(_) + | AssistantPart::Custom(_) => {} } - _ => None, - }) - .collect(); - - let content = if text.is_empty() { - Value::Null - } else { - Value::String(text) - }; - - let mut msg_obj = json!({ - "role": "assistant", - "content": content, - "tool_calls": tool_calls_json, - }); - if !reasoning.is_empty() { - msg_obj["reasoning_content"] = json!(reasoning); + } + let mut message = if provider == "groq" || !tool_calls.is_empty() { + let value = if provider != "groq" && text.is_empty() { + Value::Null + } else { + json!(text) + }; + let mut message = json!({ "role": "assistant", "content": value }); + if !tool_calls.is_empty() { + message["tool_calls"] = json!(tool_calls); + } + message + } else { + let all_plain_text = content.iter().all(|part| { + matches!( + part, + AssistantPart::Text(TextPart { + provider_options: None, + .. + }) | AssistantPart::Reasoning(_) + | AssistantPart::ToolResult(_) + | AssistantPart::ReasoningFile(_) + | AssistantPart::Custom(_) + ) + }); + let value = if all_plain_text { + json!(text) + } else { + let mut parts = Vec::new(); + for part in content { + match part { + AssistantPart::Text(text) => { + parts.push(convert_text_part_to_openai(text, provider)) + } + AssistantPart::File(file) => parts.push( + convert_file_part_to_openai(file, parts.len(), provider) + .map_err(AiMuxError::InvalidArgument)?, + ), + AssistantPart::Reasoning(_) + | AssistantPart::ToolResult(_) + | AssistantPart::ReasoningFile(_) + | AssistantPart::Custom(_) + | AssistantPart::ToolCall(_) => {} + } + } + json!(parts) + }; + json!({ "role": "assistant", "content": value }) + }; + if !reasoning.is_empty() { + let field = if provider == "groq" { + "reasoning" + } else { + "reasoning_content" + }; + message[field] = json!(reasoning); + } + message } - return Ok(vec![msg_obj]); - } + }; + Ok(vec![message]) +} - // Default path (non-Groq, no tool calls). Assistant reasoning / thinking - // parts are lifted to a top-level `reasoning_content` string (DeepSeek V4 - // thinking mode and other OpenAI-compatible reasoning models require it to - // be replayed on later turns); they are never valid OpenAI content parts, - // so they are excluded from the content shape below. Non-assistant roles do - // not carry reasoning, but the filter is harmless. - let reasoning = if msg.role == Role::Assistant { - collect_reasoning(&msg.content) - } else { - String::new() +/// Serialize a tool-result value into the OpenAI tool message content string. +fn deepseek_tool_result_content( + output: &ToolResultOutput, + warnings: &mut Vec, +) -> Result { + let ToolResultOutput::Content { value } = output else { + return Ok(tool_result_to_content(output)); + }; + let is_image = |part: &FilePart| { + get_top_level_media_type(&part.media_type) == "image" + && !matches!(part.data, FileData::Text { .. }) }; - let has_reasoning = !reasoning.is_empty(); - - // Consider only non-reasoning parts for the content shape. When they are - // all plain text (without providerOptions), collapse to a string — matching - // the Vercel AI SDK `openai-compatible` assistant conversion - // (`content: toolCalls.length > 0 ? text || null : text`). - // - // Provider-executed tool results are dropped from assistant messages: the - // OpenAI chat wire format has no content part for them, and upstream's - // assistant branch handles only `text` and `tool-call` - // (convert-to-openai-chat-messages.ts:246-300). Emitting one produces a - // part type the API rejects. - let content_parts: Vec<&ContentPart> = msg - .content + if !value .iter() - .filter(|p| match p { - ContentPart::Reasoning { .. } => false, - ContentPart::ToolResult { .. } => msg.role != Role::Assistant, - _ => true, - }) - .collect(); - let all_plain_text = content_parts.iter().all(|p| { - matches!( - p, - ContentPart::Text { - provider_options: None, - .. + .any(|part| matches!(part, ToolResultContent::File(file) if is_image(file))) + { + return Ok(tool_result_to_content(output)); + } + let mut parts = Vec::new(); + for part in value { + match part { + ToolResultContent::Text(part) => parts.push(json!({"type":"text", "text":part.text})), + ToolResultContent::File(part) if is_image(part) => { + if let FileData::Reference { reference } = &part.data { + let file_id = resolve_provider_reference(reference, "deepseek") + .map_err(AiMuxError::InvalidArgument)?; + parts.push(json!({"type":"file", "file_id":file_id})); + continue; + } + let url = match &part.data { + FileData::Url { url, .. } => { + let media_type = resolve_full_media_type(&part.media_type, ""); + if !matches!( + media_type.as_str(), + "image/jpeg" | "image/jpg" | "image/png" | "image/gif" | "image/webp" + ) { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "DeepSeek image media type {media_type}" + ))); + } + if url.len() > 8192 { + return Err(AiMuxError::InvalidPrompt( + "DeepSeek image URLs must not exceed 8192 characters.".into(), + )); + } + url.clone() + } + FileData::Data { data } => { + use base64::Engine; + let b64 = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + let media_type = resolve_full_media_type(&part.media_type, &b64); + let media_type = if media_type == "image/jpg" { + "image/jpeg" + } else { + media_type.as_str() + }; + if !matches!( + media_type, + "image/jpeg" | "image/png" | "image/gif" | "image/webp" + ) { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "DeepSeek image media type {media_type}" + ))); + } + format!("data:{media_type};base64,{b64}") + } + _ => continue, + }; + let mut image_url = json!({"url":url}); + if let Some(detail) = part + .provider_options + .as_ref() + .and_then(|options| options.get("deepseek")) + .and_then(|options| options.get("imageDetail")) + { + image_url["detail"] = detail.clone(); + } + parts.push(json!({"type":"image_url", "image_url":image_url})); } - ) - }); - - let mut msg_obj = if all_plain_text { - let text: String = content_parts - .iter() - .filter_map(|p| match p { - ContentPart::Text { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join(""); - json!({ "role": role, "content": text }) - } else { - let parts: Vec = content_parts - .iter() - .enumerate() - .map(|(i, part)| convert_part_to_openai(part, i)) - .collect::, _>>()? - .into_iter() - .filter(|p| !p.is_null()) - .collect(); - json!({ "role": role, "content": parts }) - }; - if has_reasoning { - msg_obj["reasoning_content"] = json!(reasoning); + part => warnings.push(Warning::Unsupported { + feature: format!( + "tool result content part type: {}", + if matches!(part, ToolResultContent::File(_)) { + "file" + } else { + "custom" + } + ), + details: None, + }), + } } - Ok(vec![msg_obj]) + Ok(json!(parts)) } -/// Collect and concatenate the `text` of all `ContentPart::Reasoning` parts in -/// `content`, mirroring the Vercel AI SDK assistant-message conversion. Used to -/// build the OpenAI-compatible top-level `reasoning_content` / `reasoning` -/// field that thinking models (DeepSeek V4, xAI, Groq) require to be replayed -/// across turns. -fn collect_reasoning(content: &[ContentPart]) -> String { - content - .iter() - .filter_map(|p| match p { - ContentPart::Reasoning { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join("") +pub(crate) fn tool_result_to_content(output: &ToolResultOutput) -> Value { + Value::String(match output { + ToolResultOutput::Text { value, .. } | ToolResultOutput::ErrorText { value, .. } => { + value.clone() + } + ToolResultOutput::ExecutionDenied { reason, .. } => reason + .clone() + .unwrap_or_else(|| "Tool call execution denied.".to_string()), + ToolResultOutput::Json { value, .. } | ToolResultOutput::ErrorJson { value, .. } => { + value.to_string() + } + ToolResultOutput::Content { value } => tool_result_content_value(value).to_string(), + }) } -/// Convert a system message, respecting promptCacheBreakpoint. -fn convert_system_message(msg: &LanguageModelPromptMessage, role: &str) -> Vec { - // System messages in the Rust model are always a single text part. - let text = msg - .content - .iter() - .filter_map(|p| match p { - ContentPart::Text { - text, - provider_options, - } => { - let bpt = get_prompt_cache_breakpoint(provider_options); - Some((text.clone(), bpt)) - } - _ => None, - }) - .next(); - - match text { - Some((t, None)) => vec![json!({ "role": role, "content": t })], - Some((t, Some(bpt))) => vec![json!({ - "role": role, - "content": [{ - "type": "text", - "text": t, - "prompt_cache_breakpoint": bpt, - }] - })], - None => vec![json!({ "role": role, "content": "" })], - } +pub(crate) fn tool_result_content_value(content: &[ToolResultContent]) -> Value { + json!( + content + .iter() + .map(|part| { + let (mut item, provider_options) = match part { + ToolResultContent::Text(part) => ( + json!({"type":"text", "text":part.text}), + &part.provider_options, + ), + ToolResultContent::File(part) => { + let data = match &part.data { + FileData::Data { data } => { + let data = match data { + FileBytes::Base64(data) => json!(data), + FileBytes::Binary(bytes) => Value::Object( + bytes + .iter() + .enumerate() + .map(|(i, byte)| (i.to_string(), json!(byte))) + .collect(), + ), + }; + json!({"type":"data", "data":data}) + } + FileData::Url { url, original_url } => { + let mut data = json!({"type":"url", "url":url}); + if let Some(original_url) = original_url { + data["originalUrl"] = json!(original_url); + } + data + } + FileData::Reference { reference } => { + json!({"type":"reference", "reference":reference}) + } + FileData::Text { text } => json!({"type":"text", "text":text}), + }; + let mut item = + json!({"type":"file", "data":data, "mediaType":part.media_type}); + if let Some(filename) = &part.filename { + item["filename"] = json!(filename); + } + (item, &part.provider_options) + } + ToolResultContent::Custom { provider_options } => { + (json!({"type":"custom"}), provider_options) + } + }; + if let Some(provider_options) = provider_options { + item["providerOptions"] = json!(provider_options); + } + item + }) + .collect::>() + ) } -/// Serialize a tool-result `output` value into the OpenAI tool message -/// `content` string. -fn tool_result_to_content(output: &Value) -> Value { +pub(crate) fn tool_result_cache_breakpoint(output: &ToolResultOutput) -> Option { match output { - Value::String(s) => Value::String(s.clone()), - other => Value::String(other.to_string()), - } -} - -/// Convert a content part to the OpenAI format. -fn convert_part_to_openai(part: &ContentPart, index: usize) -> Result { - match part { - ContentPart::Text { - text, - provider_options, - } => { - let bpt = get_prompt_cache_breakpoint(provider_options); - let mut p = json!({ "type": "text", "text": text }); - if let Some(b) = bpt { - p["prompt_cache_breakpoint"] = b; - } - Ok(p) + ToolResultOutput::Text { + provider_options, .. } - ContentPart::Image { - image, - media_type, - provider_options, - } => { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(image); - convert_file_part_to_openai( - media_type, - Some(&b64), - None, - None, - None, - provider_options, - index, - ) - .map_err(AiMuxError::InvalidArgument) + | ToolResultOutput::Json { + provider_options, .. } - ContentPart::File { - data, - media_type, - filename, - provider_options, - } => { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(data); - convert_file_part_to_openai( - media_type, - Some(&b64), - None, - None, - filename.as_deref(), - provider_options, - index, - ) - .map_err(AiMuxError::InvalidArgument) + | ToolResultOutput::ErrorText { + provider_options, .. } - ContentPart::FileBase64 { - data, - media_type, - filename, - provider_options, - } => convert_file_part_to_openai( - media_type, - Some(data), - None, - None, - filename.as_deref(), - provider_options, - index, - ) - .map_err(AiMuxError::InvalidArgument), - ContentPart::FileUrl { - url, - media_type, - provider_options, - } => convert_file_part_to_openai( - media_type, - None, - Some(url), - None, - None, - provider_options, - index, - ) - .map_err(AiMuxError::InvalidArgument), - ContentPart::FileReference { - media_type, - reference, - filename, - provider_options, - } => convert_file_part_to_openai( - media_type, - None, - None, - Some(reference), - filename.as_deref(), - provider_options, - index, - ) - .map_err(AiMuxError::InvalidArgument), - ContentPart::Reasoning { .. } => Ok(Value::Null), - // These variants are handled by `convert_message_to_openai` for - // assistant/tool roles; kept here as a defensive fallback. - ContentPart::ToolCall { - tool_call_id, - tool_name, - input, - .. - } => Ok(json!({ - "type": "tool_call", - "id": tool_call_id, - "function": { - "name": tool_name, - "arguments": input.to_string(), + | ToolResultOutput::ErrorJson { + provider_options, .. + } + | ToolResultOutput::ExecutionDenied { + provider_options, .. + } => get_prompt_cache_breakpoint(provider_options), + ToolResultOutput::Content { value } => value.iter().find_map(|part| match part { + aimux_core::language_model_message::ToolResultContent::Text(part) => { + get_prompt_cache_breakpoint(&part.provider_options) + } + aimux_core::language_model_message::ToolResultContent::File(part) => { + get_prompt_cache_breakpoint(&part.provider_options) + } + aimux_core::language_model_message::ToolResultContent::Custom { provider_options } => { + get_prompt_cache_breakpoint(provider_options) } - })), - // Tool results reach the wire as a separate `role: "tool"` message - // (built by `convert_message_to_openai`), never as a content part — - // there is no such part type in the OpenAI chat format. - ContentPart::ToolResult { .. } - | ContentPart::Custom { .. } - | ContentPart::ReasoningFile { .. } - | ContentPart::ToolApprovalRequest { .. } => Ok(Value::Null), + }), } } +fn convert_text_part_to_openai(part: &TextPart, provider: &str) -> Value { + let mut value = json!({ "type": "text", "text": part.text }); + if provider == "openai" + && let Some(bpt) = get_prompt_cache_breakpoint(&part.provider_options) + { + value["prompt_cache_breakpoint"] = bpt; + } + value +} + // ── Request body ──────────────────────────────────────────────────────────── /// Result of building a request body, including warnings. @@ -1642,10 +1644,11 @@ pub fn build_request_body_with_warnings( }); } - let messages = convert_prompt_to_openai_messages_with_provider( + let messages = convert_prompt_with_warnings( &options.prompt, system_message_mode, provider, + &mut warnings, )?; let mut body = json!({ diff --git a/aimux-providers/src/openai/responses/convert.rs b/aimux-providers/src/openai/responses/convert.rs index bc5768aba..2ed3a40dc 100644 --- a/aimux-providers/src/openai/responses/convert.rs +++ b/aimux-providers/src/openai/responses/convert.rs @@ -12,11 +12,13 @@ use serde_json::{Value, json}; -use aimux_core::content::ContentPart; -use aimux_core::language_model_message::LanguageModelPrompt; -use aimux_core::message::Role; +use aimux_core::error::AiMuxError; +use aimux_core::language_model_message::{ + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, +}; use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; -use aimux_core::shared::SharedProviderOptions; +use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified, ReasoningEffort, Usage, Warning}; @@ -69,29 +71,33 @@ pub struct ResponsesInputResult { /// `has_previous_response_id` is true, assistant reasoning/function-call items /// that already carry an `itemId` are skipped (they live in the previous /// response chain). +/// +/// # Errors +/// Returns an unsupported-functionality error for text file parts. pub fn convert_to_responses_input( prompt: &LanguageModelPrompt, system_message_mode: SystemMessageMode, store: bool, has_previous_response_id: bool, has_conversation: bool, -) -> ResponsesInputResult { +) -> Result { let mut input: Vec = Vec::new(); let mut warnings: Vec = Vec::new(); + let mut processed_approval_ids = std::collections::HashSet::new(); for msg in prompt { - match msg.role { - Role::System => match system_message_mode { + match msg { + LanguageModelMessage::System { content, .. } => match system_message_mode { SystemMessageMode::System => { input.push(json!({ "role": "system", - "content": join_text_parts(&msg.content), + "content": content, })); } SystemMessageMode::Developer => { input.push(json!({ "role": "developer", - "content": join_text_parts(&msg.content), + "content": content, })); } SystemMessageMode::Remove => { @@ -100,36 +106,22 @@ pub fn convert_to_responses_input( }); } }, - Role::User => { - let content: Vec = msg.content.iter().map(convert_user_part).collect(); + LanguageModelMessage::User { content, .. } => { + let content: Vec = content + .iter() + .map(convert_user_part) + .collect::>()?; input.push(json!({ "role": "user", "content": content })); } - Role::Assistant => { - for part in &msg.content { + LanguageModelMessage::Assistant { content, .. } => { + for part in content { match part { - ContentPart::Custom { - kind, - provider_options, - } if kind == "openai.compaction" => { - if let Some(id) = item_id(provider_options) { - if has_conversation { - continue; - } - if store { - input.push(json!({ "type": "item_reference", "id": id })); - } else { - input.push(json!({ "type": "compaction", "id": id, - "encrypted_content": openai_sub_option(provider_options, "encryptedContent") })); - } - } - } - - ContentPart::Text { + AssistantPart::Text(TextPart { text, provider_options, - } => { + }) => { let id = item_id(provider_options); - if has_previous_response_id && id.is_some() { + if (has_previous_response_id || has_conversation) && id.is_some() { continue; } if store && let Some(ref id) = id { @@ -149,15 +141,15 @@ pub fn convert_to_responses_input( } input.push(item); } - ContentPart::ToolCall { + AssistantPart::ToolCall(ToolCallPart { tool_call_id, tool_name, input: tool_input, provider_options, .. - } => { + }) => { let id = item_id(provider_options); - if has_previous_response_id && id.is_some() { + if (has_previous_response_id || has_conversation) && id.is_some() { continue; } let namespace = namespace_from_provider_options(provider_options); @@ -172,13 +164,15 @@ pub fn convert_to_responses_input( } input.push(item); } - ContentPart::Reasoning { + AssistantPart::Reasoning(ReasoningPart { text, provider_options, .. - } => { + }) => { let reasoning_id = openai_sub_option(provider_options, "itemId"); - if has_previous_response_id && reasoning_id.is_some() { + if (has_previous_response_id || has_conversation) + && reasoning_id.is_some() + { continue; } if let Some(ref rid) = reasoning_id { @@ -227,28 +221,84 @@ pub fn convert_to_responses_input( } } } + AssistantPart::ToolResult(part) => { + if has_conversation + || matches!(part.output, ToolResultOutput::ExecutionDenied { .. }) + || matches!(&part.output, ToolResultOutput::Json { value, .. } if value.get("type").and_then(Value::as_str) == Some("execution-denied")) + { + continue; + } + if store { + let id = item_id(&part.provider_options) + .unwrap_or_else(|| part.tool_call_id.clone()); + input.push(json!({"type":"item_reference", "id":id})); + } else { + warnings.push(Warning::Other { + message:format!("Results for OpenAI tool {} are not sent to the API when store is false", part.tool_name), + }); + } + } + AssistantPart::Custom(part) if part.kind == "openai.compaction" => { + let id = item_id(&part.provider_options); + if has_conversation && id.is_some() { + continue; + } + if let Some(id) = id { + if store { + input.push(json!({"type":"item_reference", "id":id})); + } else { + input.push(json!({"type":"compaction", "id":id, "encrypted_content":openai_sub_option(&part.provider_options, "encryptedContent")})); + } + } + } _ => {} } } } - Role::Tool => { - for part in &msg.content { - if let ContentPart::ToolResult { + LanguageModelMessage::Tool { content, .. } => { + for part in content { + if let ToolPart::ToolApprovalResponse(approval) = part { + if !processed_approval_ids.insert(&approval.approval_id) { + continue; + } + if store && !has_conversation && !has_previous_response_id { + input.push(json!({"type":"item_reference", "id":approval.approval_id})); + } + input.push(json!({"type":"mcp_approval_response", "approval_request_id":approval.approval_id, "approve":approval.approved})); + continue; + } + let ToolPart::ToolResult(ToolResultPart { tool_call_id, - result, + output, + provider_options, .. - } = part + }) = part + else { + continue; + }; + if let ToolResultOutput::ExecutionDenied { + provider_options, .. + } = output + && openai_sub_option(provider_options, "approvalId") + .and_then(|value| value.as_str().map(|id| !id.is_empty())) + == Some(true) + { + continue; + } + let mut content_value = convert_tool_result_output(output, &mut warnings)?; + if !matches!(output, ToolResultOutput::Content { .. }) + && let Some(breakpoint) = + crate::openai::convert::tool_result_cache_breakpoint(output).or_else( + || openai_sub_option(provider_options, "promptCacheBreakpoint"), + ) { - let content_value = match result { - Value::String(s) => Value::String(s.clone()), - other => Value::String(other.to_string()), - }; - input.push(json!({ - "type": "function_call_output", - "call_id": tool_call_id, - "output": content_value, - })); + content_value = json!([{"type":"input_text", "text":content_value, "prompt_cache_breakpoint":breakpoint}]); } + input.push(json!({ + "type": "function_call_output", + "call_id": tool_call_id, + "output": content_value, + })); } } } @@ -272,94 +322,122 @@ pub fn convert_to_responses_input( }); } - ResponsesInputResult { input, warnings } -} - -/// Join all text parts of a message into a single string (for system messages). -fn join_text_parts(parts: &[ContentPart]) -> String { - parts - .iter() - .filter_map(|p| match p { - ContentPart::Text { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join("") + Ok(ResponsesInputResult { input, warnings }) } /// Convert a single user-message content part into the Responses input shape. -fn convert_user_part(part: &ContentPart) -> Value { - match part { - ContentPart::Text { text, .. } => json!({ "type": "input_text", "text": text }), - ContentPart::Image { - image, - media_type, - provider_options, - } => { - let b64 = { - use base64::Engine; - base64::engine::general_purpose::STANDARD.encode(image) - }; - let detail = openai_sub_option(provider_options, "imageDetail"); - let mut img = json!({ - "type": "input_image", - "image_url": format!("data:{};base64,{}", media_type, b64), - }); - if let Some(d) = detail { - img["detail"] = d; - } - img - } - ContentPart::File { +fn convert_user_part(part: &UserPart) -> Result { + Ok(match part { + UserPart::Text(TextPart { text, .. }) => json!({ "type": "input_text", "text": text }), + UserPart::File(FilePart { data, media_type, filename, - .. - } => { - let b64 = { - use base64::Engine; - base64::engine::general_purpose::STANDARD.encode(data) + provider_options, + }) => { + let is_image = media_type.split('/').next() == Some("image"); + let mut item = match data { + FileData::Data { data } => { + let b64 = match data { + FileBytes::Binary(bytes) => { + use base64::Engine; + base64::engine::general_purpose::STANDARD.encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + let data_url = format!("data:{media_type};base64,{b64}"); + if is_image { + json!({ "type": "input_image", "image_url": data_url }) + } else { + json!({ + "type": "input_file", + "filename": filename.as_deref().unwrap_or("part.pdf"), + "file_data": data_url, + }) + } + } + FileData::Url { url, .. } => { + if is_image { + json!({ "type": "input_image", "image_url": url }) + } else { + json!({ "type": "input_file", "file_url": url }) + } + } + FileData::Reference { reference } => { + let file_id = reference.get("openai").ok_or_else(|| { + AiMuxError::InvalidArgument( + "missing file reference for provider openai".into(), + ) + })?; + json!({ "type": if is_image { "input_image" } else { "input_file" }, "file_id": file_id }) + } + FileData::Text { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "text file parts".into(), + )); + } }; - let fname = filename.clone().unwrap_or_else(|| "part.pdf".to_string()); - json!({ - "type": "input_file", - "filename": fname, - "file_data": format!("data:{};base64,{}", media_type, b64), - }) - } - ContentPart::FileBase64 { - data, - media_type, - filename, - .. - } => { - let fname = filename.clone().unwrap_or_else(|| "part.pdf".to_string()); - json!({ - "type": "input_file", - "filename": fname, - "file_data": format!("data:{};base64,{}", media_type, data), - }) + if item["type"] == "input_image" + && let Some(detail) = openai_sub_option(provider_options, "imageDetail") + { + item["detail"] = detail; + } + item } - ContentPart::FileUrl { - url, media_type, .. - } => { - if media_type.starts_with("image") { - json!({ "type": "input_image", "image_url": url }) - } else { - json!({ "type": "input_file", "file_url": url }) + }) +} + +fn convert_tool_result_output( + output: &ToolResultOutput, + warnings: &mut Vec, +) -> Result { + let ToolResultOutput::Content { value } = output else { + return Ok(crate::openai::convert::tool_result_to_content(output)); + }; + let mut parts = Vec::new(); + for item in value { + let (mut converted, provider_options) = match item { + ToolResultContent::Text(part) => ( + json!({"type":"input_text", "text":part.text}), + &part.provider_options, + ), + ToolResultContent::File(part) => { + if matches!(part.data, FileData::Text { .. }) { + warnings.push(Warning::Other { + message: "unsupported tool content part type: file with data type: text" + .into(), + }); + continue; + } + let mut converted = convert_user_part(&UserPart::File(part.clone()))?; + if converted["type"] == "input_file" && converted.get("file_data").is_some() { + converted["filename"] = json!(part.filename.as_deref().unwrap_or("data")); + } + if converted["type"] == "input_image" + && let Some(detail) = openai_sub_option(&part.provider_options, "imageDetail") + { + converted["detail"] = detail; + } + (converted, &part.provider_options) + } + ToolResultContent::Custom { .. } => { + warnings.push(Warning::Other { + message: "unsupported tool content part type: custom".into(), + }); + continue; } + }; + if let Some(breakpoint) = openai_sub_option(provider_options, "promptCacheBreakpoint") { + converted["prompt_cache_breakpoint"] = breakpoint; } - _ => json!({ "type": "input_text", "text": "" }), + parts.push(converted); } + Ok(json!(parts)) } -/// Serialize tool-call arguments: null/empty -> `"{}"`, objects -> JSON string. +/// Serialize tool-call arguments as JSON, mirroring JSON.stringify. fn serialize_arguments(input: &Value) -> String { - match input { - Value::Null => "{}".to_string(), - Value::String(s) if s.is_empty() => "{}".to_string(), - other => other.to_string(), - } + input.to_string() } /// Read the `itemId` from a content part's `providerOptions.openai.itemId`. @@ -856,12 +934,14 @@ fn apply_responses_reasoning_block( /// /// Splits the original ~380-line function into focused helpers (issue M11); /// behavior is unchanged. -#[must_use] +/// +/// # Errors +/// Returns an unsupported-functionality error for text file parts. pub fn build_responses_request_body( model_id: &str, options: &CallOptions, stream: bool, -) -> ResponsesRequestBodyResult { +) -> Result { let mut warnings: Vec = Vec::new(); let caps = get_model_capabilities(model_id); let provider_opts = &options.provider_options; @@ -891,7 +971,7 @@ pub fn build_responses_request_body( store_bool, has_previous_response_id, openai_option(provider_opts, "conversation").is_some(), - ); + )?; warnings.extend(input_result.warnings); // -- Base body -- @@ -960,7 +1040,7 @@ pub fn build_responses_request_body( warnings.push(tw); } - ResponsesRequestBodyResult { body, warnings } + Ok(ResponsesRequestBodyResult { body, warnings }) } // -- Usage conversion -------------------------------------------------------- diff --git a/aimux-providers/src/openai/responses/mod.rs b/aimux-providers/src/openai/responses/mod.rs index b5b05668c..38ed6f196 100644 --- a/aimux-providers/src/openai/responses/mod.rs +++ b/aimux-providers/src/openai/responses/mod.rs @@ -121,7 +121,7 @@ impl LanguageModel for OpenAIResponsesModel { async fn do_generate(&self, options: &CallOptions) -> Result { let headers = self.build_headers(options.headers.as_ref()); - let request_result = build_responses_request_body(&self.model_id, options, false); + let request_result = build_responses_request_body(&self.model_id, options, false)?; let body = request_result.body; let provider_key = self.provider_options_name().to_string(); @@ -155,7 +155,7 @@ impl LanguageModel for OpenAIResponsesModel { async fn do_stream(&self, options: &CallOptions) -> Result { let headers = self.build_headers(options.headers.as_ref()); - let request_result = build_responses_request_body(&self.model_id, options, true); + let request_result = build_responses_request_body(&self.model_id, options, true)?; let body = request_result.body; let warnings = request_result.warnings; let provider_key = self.provider_options_name().to_string(); diff --git a/aimux-providers/src/replay.rs b/aimux-providers/src/replay.rs index 35376e726..ce41c6d1f 100644 --- a/aimux-providers/src/replay.rs +++ b/aimux-providers/src/replay.rs @@ -455,14 +455,7 @@ mod tests { let ring = aimux_core::recording::RingRecorder::with_capacity(8); let options = aimux_core::generate::GenerateTextOptions::default().into_call_options(vec![ - aimux_core::language_model_message::LanguageModelPromptMessage { - role: aimux_core::message::Role::User, - content: vec![aimux_core::content::ContentPart::Text { - text: "ping".into(), - provider_options: None, - }], - provider_options: None, - }, + aimux_core::language_model_message::LanguageModelMessage::user_text("ping"), ]); ring.record_input("c1", &options, "openai", "gpt-4o"); ring.record_provider("c1", &snap); diff --git a/aimux-providers/src/vertex/anthropic_model.rs b/aimux-providers/src/vertex/anthropic_model.rs index 83e904e7f..68d84b36c 100644 --- a/aimux-providers/src/vertex/anthropic_model.rs +++ b/aimux-providers/src/vertex/anthropic_model.rs @@ -205,6 +205,7 @@ impl LanguageModel for VertexAnthropicModel { let content = crate::anthropic::stream::parse_anthropic_content( &data.content, &ToolNameMapping::new(options.tools.as_deref()), + &crate::anthropic::stream::generate_id, ); let finish_reason = data @@ -421,6 +422,7 @@ impl LanguageModel for VertexAnthropicModel { &tool_names, &mcp_tool_calls, &server_tool_calls, + &crate::anthropic::stream::generate_id, ) { yield Ok(part); } diff --git a/aimux-providers/src/vertex/model.rs b/aimux-providers/src/vertex/model.rs index c187b8d91..20a7163ed 100644 --- a/aimux-providers/src/vertex/model.rs +++ b/aimux-providers/src/vertex/model.rs @@ -164,15 +164,28 @@ impl LanguageModel for VertexModel { if options .prompt .iter() - .filter(|message| message.role == aimux_core::message::Role::Assistant) - .flat_map(|message| &message.content) + .filter_map(|message| match message { + aimux_core::language_model_message::LanguageModelMessage::Assistant { + content, + .. + } => Some(content), + _ => None, + }) + .flatten() .any(|part| { matches!( part, - aimux_core::content::ContentPart::ReasoningFile { - data: aimux_core::shared::GeneratedFileData::Url { .. }, - .. - } + aimux_core::language_model_message::AssistantPart::ReasoningFile( + aimux_core::language_model_message::ReasoningFilePart { + data: aimux_core::shared::GeneratedFileData::Url { .. }, + .. + } + ) | aimux_core::language_model_message::AssistantPart::File( + aimux_core::language_model_message::FilePart { + data: aimux_core::shared::FileData::Url { .. }, + .. + } + ) ) }) { @@ -180,6 +193,40 @@ impl LanguageModel for VertexModel { "File data URLs in assistant messages are not supported".to_string(), )); } + if options.prompt.iter().any(|message| match message { + aimux_core::language_model_message::LanguageModelMessage::Assistant { + content, .. + } => content.iter().any(|part| { + matches!( + part, + aimux_core::language_model_message::AssistantPart::File( + aimux_core::language_model_message::FilePart { + data: aimux_core::shared::FileData::Reference { .. }, + .. + } + ) + ) + }), + aimux_core::language_model_message::LanguageModelMessage::User { content, .. } => { + content.iter().any(|part| { + matches!( + part, + aimux_core::language_model_message::UserPart::File( + aimux_core::language_model_message::FilePart { + data: aimux_core::shared::FileData::Reference { .. }, + .. + } + ) + ) + }) + } + _ => false, + }) { + return Err(AiMuxError::UnsupportedFunctionality( + "file parts with provider references".to_string(), + )); + } + crate::google::convert::validate_tool_result_files(&options.prompt)?; let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); let body = build_vertex_request_body(&self.model_id, options)?; let headers = self.build_headers(options.headers.as_ref()); @@ -260,15 +307,28 @@ impl LanguageModel for VertexModel { if options .prompt .iter() - .filter(|message| message.role == aimux_core::message::Role::Assistant) - .flat_map(|message| &message.content) + .filter_map(|message| match message { + aimux_core::language_model_message::LanguageModelMessage::Assistant { + content, + .. + } => Some(content), + _ => None, + }) + .flatten() .any(|part| { matches!( part, - aimux_core::content::ContentPart::ReasoningFile { - data: aimux_core::shared::GeneratedFileData::Url { .. }, - .. - } + aimux_core::language_model_message::AssistantPart::ReasoningFile( + aimux_core::language_model_message::ReasoningFilePart { + data: aimux_core::shared::GeneratedFileData::Url { .. }, + .. + } + ) | aimux_core::language_model_message::AssistantPart::File( + aimux_core::language_model_message::FilePart { + data: aimux_core::shared::FileData::Url { .. }, + .. + } + ) ) }) { @@ -276,6 +336,40 @@ impl LanguageModel for VertexModel { "File data URLs in assistant messages are not supported".to_string(), )); } + if options.prompt.iter().any(|message| match message { + aimux_core::language_model_message::LanguageModelMessage::Assistant { + content, .. + } => content.iter().any(|part| { + matches!( + part, + aimux_core::language_model_message::AssistantPart::File( + aimux_core::language_model_message::FilePart { + data: aimux_core::shared::FileData::Reference { .. }, + .. + } + ) + ) + }), + aimux_core::language_model_message::LanguageModelMessage::User { content, .. } => { + content.iter().any(|part| { + matches!( + part, + aimux_core::language_model_message::UserPart::File( + aimux_core::language_model_message::FilePart { + data: aimux_core::shared::FileData::Reference { .. }, + .. + } + ) + ) + }) + } + _ => false, + }) { + return Err(AiMuxError::UnsupportedFunctionality( + "file parts with provider references".to_string(), + )); + } + crate::google::convert::validate_tool_result_files(&options.prompt)?; let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); let body = build_vertex_request_body(&self.model_id, options)?; let headers = self.build_headers(options.headers.as_ref()); diff --git a/aimux-providers/src/xai/convert.rs b/aimux-providers/src/xai/convert.rs index e43528dc9..3434df4bc 100644 --- a/aimux-providers/src/xai/convert.rs +++ b/aimux-providers/src/xai/convert.rs @@ -5,12 +5,13 @@ //! `supports-reasoning-effort.ts`, `map-xai-finish-reason.ts`, and //! `remove-additional-properties.ts`. -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; -use aimux_core::language_model_message::LanguageModelPrompt; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolCallPart, ToolPart, + ToolResultPart, UserPart, +}; use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; -use aimux_core::shared::SharedProviderOptions; +use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; use aimux_core::tool::Tool; use aimux_core::types::{FinishReason, FinishReasonUnified, ReasoningEffort, Warning}; @@ -258,17 +259,14 @@ pub fn resolve_full_media_type(media_type: &str, b64_data: &str) -> String { /// /// Returns a `String` listing the available providers when the requested key /// is absent. -pub fn resolve_provider_reference(reference: &Value, provider: &str) -> Result { - if let Some(val) = reference.get(provider) { - if let Some(s) = val.as_str() { - return Ok(s.to_string()); - } - return Ok(val.to_string()); - } - let available: Vec = reference - .as_object() - .map(|m| m.keys().cloned().collect()) - .unwrap_or_default(); +pub fn resolve_provider_reference( + reference: &std::collections::HashMap, + provider: &str, +) -> Result { + if let Some(value) = reference.get(provider) { + return Ok(value.clone()); + } + let available: Vec = reference.keys().cloned().collect(); Err(format!( "No provider reference found for provider '{}'. Available providers: {}", provider, @@ -291,49 +289,36 @@ pub fn convert_to_xai_messages( let warnings: Vec = Vec::new(); for msg in prompt { - match msg.role { - Role::System => { - let content: String = msg - .content - .iter() - .filter_map(|p| match p { - ContentPart::Text { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join(""); + match msg { + LanguageModelMessage::System { content, .. } => { messages.push(json!({ "role": "system", "content": content })); } - Role::User => { - if msg.content.len() == 1 - && let ContentPart::Text { text, .. } = &msg.content[0] + LanguageModelMessage::User { content, .. } => { + if content.len() == 1 + && let UserPart::Text(TextPart { text, .. }) = &content[0] { messages.push(json!({ "role": "user", "content": text })); continue; } let mut user_content = Vec::new(); - for (i, part) in msg.content.iter().enumerate() { + for (i, part) in content.iter().enumerate() { user_content.push(convert_user_part(part, i)?); } messages.push(json!({ "role": "user", "content": user_content })); } - Role::Assistant => { + LanguageModelMessage::Assistant { content, .. } => { let mut text = String::new(); let mut tool_calls: Vec = Vec::new(); - for part in &msg.content { + for part in content { match part { - ContentPart::Text { text: t, .. } => text.push_str(t), - ContentPart::ToolCall { + AssistantPart::Text(TextPart { text: t, .. }) => text.push_str(t), + AssistantPart::ToolCall(ToolCallPart { tool_call_id, tool_name, input, .. - } => { - let arguments = if input.is_null() { - "{}".to_string() - } else { - input.to_string() - }; + }) => { + let arguments = input.to_string(); tool_calls.push(json!({ "id": tool_call_id, "type": "function", "function": { "name": tool_name, "arguments": arguments } })); } _ => {} @@ -345,20 +330,20 @@ pub fn convert_to_xai_messages( } messages.push(msg_obj); } - Role::Tool => { - for part in &msg.content { - if let ContentPart::ToolResult { + LanguageModelMessage::Tool { content, .. } => { + for part in content { + let ToolPart::ToolResult(ToolResultPart { tool_call_id, - result, + output, .. - } = part - { - let content = match result { - Value::String(s) => s.clone(), - other => other.to_string(), - }; - messages.push(json!({ "role": "tool", "tool_call_id": tool_call_id, "content": content })); - } + }) = part + else { + continue; + }; + let content = crate::openai::convert::tool_result_to_content(output); + messages.push( + json!({ "role": "tool", "tool_call_id": tool_call_id, "content": content }), + ); } } } @@ -366,52 +351,32 @@ pub fn convert_to_xai_messages( Ok((messages, warnings)) } -fn convert_user_part(part: &ContentPart, _index: usize) -> Result { +fn convert_user_part(part: &UserPart, _index: usize) -> Result { match part { - ContentPart::Text { text, .. } => Ok(json!({ "type": "text", "text": text })), - ContentPart::Image { - image, - media_type, - provider_options, - } => { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(image); - convert_image_part(media_type, Some(&b64), None, provider_options) - } - ContentPart::File { - data, - media_type, - provider_options, - .. - } => { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(data); - convert_image_part(media_type, Some(&b64), None, provider_options) - } - ContentPart::FileBase64 { - data, - media_type, - provider_options, - .. - } => convert_image_part(media_type, Some(data), None, provider_options), - ContentPart::FileUrl { - url, - media_type, - provider_options, - } => convert_image_part(media_type, None, Some(url), provider_options), - ContentPart::FileReference { - media_type, - reference, - provider_options, - .. - } => { - let _ = media_type; - let _ = provider_options; - let file_id = resolve_provider_reference(reference, "xai") - .map_err(AiMuxError::InvalidArgument)?; - Ok(json!({ "type": "file", "file": { "file_id": file_id } })) - } - _ => Ok(Value::Null), + UserPart::Text(TextPart { text, .. }) => Ok(json!({ "type": "text", "text": text })), + UserPart::File(file) => match &file.data { + FileData::Data { + data: FileBytes::Binary(bytes), + } => { + use base64::Engine; + let b64 = base64::engine::general_purpose::STANDARD.encode(bytes); + convert_image_part(&file.media_type, Some(&b64), None, &file.provider_options) + } + FileData::Data { + data: FileBytes::Base64(data), + } => convert_image_part(&file.media_type, Some(data), None, &file.provider_options), + FileData::Url { url, .. } => { + convert_image_part(&file.media_type, None, Some(url), &file.provider_options) + } + FileData::Reference { reference } => { + let file_id = resolve_provider_reference(reference, "xai") + .map_err(AiMuxError::InvalidArgument)?; + Ok(json!({ "type": "file", "file": { "file_id": file_id } })) + } + FileData::Text { .. } => Err(AiMuxError::UnsupportedFunctionality( + "text file parts".into(), + )), + }, } } diff --git a/aimux-providers/src/xai/responses/convert.rs b/aimux-providers/src/xai/responses/convert.rs index 6efaf41a2..0855f9ce2 100644 --- a/aimux-providers/src/xai/responses/convert.rs +++ b/aimux-providers/src/xai/responses/convert.rs @@ -6,12 +6,13 @@ //! - `map-xai-responses-finish-reason.ts` //! - `xai-responses-prepare-tools.ts` -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; -use aimux_core::language_model_message::LanguageModelPrompt; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, +}; use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; -use aimux_core::shared::SharedProviderOptions; +use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; use aimux_core::tool::Tool; use aimux_core::types::{FinishReasonUnified, ReasoningEffort, Warning}; @@ -295,114 +296,85 @@ pub fn convert_to_xai_responses_input( let mut warnings: Vec = Vec::new(); for msg in prompt { - match msg.role { - Role::System => { - let content: String = msg - .content - .iter() - .filter_map(|p| match p { - ContentPart::Text { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join(""); + match msg { + LanguageModelMessage::System { content, .. } => { input.push(json!({ "role": "system", "content": content })); } - Role::User => { + LanguageModelMessage::User { content, .. } => { let mut content_parts = Vec::new(); - for part in &msg.content { + for part in content { match part { - ContentPart::Text { text, .. } => { + UserPart::Text(TextPart { text, .. }) => { content_parts.push(json!({ "type": "input_text", "text": text })); } - ContentPart::Image { - image, - media_type, - provider_options, - } => { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(image); - content_parts.push(convert_image_part( - media_type, - Some(&b64), - None, - provider_options, - )); - } - ContentPart::File { - data, - media_type, - provider_options, - .. - } => { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(data); - content_parts.push(convert_image_part( - media_type, - Some(&b64), - None, - provider_options, - )); - } - ContentPart::FileBase64 { - data, - media_type, - provider_options, - .. - } => { - content_parts.push(convert_image_part( - media_type, - Some(data), - None, - provider_options, - )); - } - ContentPart::FileUrl { - url, - media_type, - provider_options, - } => { - let top_level = media_type.split('/').next().unwrap_or(""); - if top_level == "image" { - content_parts.push(convert_image_part( - media_type, - None, - Some(url), - provider_options, - )); - } else { - content_parts - .push(json!({ "type": "input_file", "file_url": url })); + UserPart::File(file) => { + if matches!(file.data, FileData::Data { .. }) + && file.media_type.split('/').next() != Some("image") + { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "file part media type {} as inline data (xAI Responses requires a URL or a Files API reference for non-image files)", + file.media_type + ))); } - } - ContentPart::FileReference { - media_type, - reference, - provider_options, - .. - } => { - let _ = (media_type, provider_options); - let file_id = resolve_provider_reference(reference, "xai") - .map_err(AiMuxError::InvalidArgument)?; - content_parts.push(json!({ "type": "input_file", "file_id": file_id })); - } - _ => { - warnings.push(Warning::Other { - message: "xAI Responses API does not support this content type in user messages".to_string(), + content_parts.push(match &file.data { + FileData::Data { + data: FileBytes::Binary(bytes), + } => { + use base64::Engine; + let b64 = + base64::engine::general_purpose::STANDARD.encode(bytes); + convert_image_part( + &file.media_type, + Some(&b64), + None, + &file.provider_options, + ) + } + FileData::Data { + data: FileBytes::Base64(data), + } => convert_image_part( + &file.media_type, + Some(data), + None, + &file.provider_options, + ), + FileData::Url { url, .. } + if file.media_type.split('/').next() == Some("image") => + { + convert_image_part( + &file.media_type, + None, + Some(url), + &file.provider_options, + ) + } + FileData::Url { url, .. } => { + json!({ "type": "input_file", "file_url": url }) + } + FileData::Reference { reference } => { + let file_id = resolve_provider_reference(reference, "xai") + .map_err(AiMuxError::InvalidArgument)?; + json!({ "type": "input_file", "file_id": file_id }) + } + FileData::Text { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "text file parts".into(), + )); + } }); } } } input.push(json!({ "role": "user", "content": content_parts })); } - Role::Assistant => { - for part in &msg.content { + LanguageModelMessage::Assistant { content, .. } => { + for part in content { match part { - ContentPart::Text { + AssistantPart::Text(TextPart { text, provider_options, .. - } => { + }) => { let id = provider_options .as_ref() .and_then(|po| po.get("xai")) @@ -415,14 +387,14 @@ pub fn convert_to_xai_responses_input( } input.push(msg_obj); } - ContentPart::ToolCall { + AssistantPart::ToolCall(ToolCallPart { tool_call_id, tool_name, input: tool_input, provider_executed, provider_options, .. - } => { + }) => { // Skip provider-executed tool calls. // The standardized field is authoritative. Fall back // to the legacy provider option only when it is absent. @@ -444,11 +416,7 @@ pub fn convert_to_xai_responses_input( .and_then(|v| v.as_str()) .map(std::string::ToString::to_string); let item_id = id.unwrap_or_else(|| tool_call_id.clone()); - let arguments = if tool_input.is_null() { - "{}".to_string() - } else { - tool_input.to_string() - }; + let arguments = tool_input.to_string(); input.push(json!({ "type": "function_call", "id": item_id, @@ -458,12 +426,12 @@ pub fn convert_to_xai_responses_input( "status": "completed" })); } - ContentPart::ToolResult { .. } => {} - ContentPart::Reasoning { + AssistantPart::ToolResult(_) => {} + AssistantPart::Reasoning(ReasoningPart { text, provider_options, .. - } => { + }) => { let item_id = provider_options .as_ref() .and_then(|po| po.get("xai")) @@ -506,30 +474,72 @@ pub fn convert_to_xai_responses_input( } } } - Role::Tool => { - for part in &msg.content { - if let ContentPart::ToolResult { + LanguageModelMessage::Tool { content, .. } => { + for part in content { + let ToolPart::ToolResult(ToolResultPart { tool_call_id, - result, + output, .. - } = part + }) = part + else { + continue; + }; + let output_value = convert_tool_result_output(output); + input.push(json!({ + "type": "function_call_output", + "call_id": tool_call_id, + "output": output_value + })); + } + } + } + } + + Ok((input, warnings)) +} + +fn convert_tool_result_output(output: &ToolResultOutput) -> Value { + match output { + ToolResultOutput::ExecutionDenied { reason, .. } => { + json!(reason.as_deref().unwrap_or("tool execution denied")) + } + ToolResultOutput::Content { value } => { + let mut parts = Vec::new(); + for item in value { + match item { + ToolResultContent::Text(part) => { + parts.push(json!({"type":"input_text", "text":part.text})) + } + ToolResultContent::File(part) + if part.media_type.split('/').next() == Some("image") => { - let output_value = match result { - Value::String(s) => s.clone(), - other => other.to_string(), + let url = match &part.data { + FileData::Url { url, .. } => url.clone(), + FileData::Data { data } => { + use base64::Engine; + let b64 = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + format!( + "data:{};base64,{}", + resolve_full_media_type(&part.media_type, &b64), + b64 + ) + } + _ => continue, }; - input.push(json!({ - "type": "function_call_output", - "call_id": tool_call_id, - "output": output_value - })); + parts.push(json!({"type":"input_image", "image_url":url})); } + _ => {} } } + json!(parts) } + _ => crate::openai::convert::tool_result_to_content(output), } - - Ok((input, warnings)) } fn convert_image_part( diff --git a/aimux-providers/tests/alibaba_test.rs b/aimux-providers/tests/alibaba_test.rs index 30d1ec812..9ff1525c6 100644 --- a/aimux-providers/tests/alibaba_test.rs +++ b/aimux-providers/tests/alibaba_test.rs @@ -31,11 +31,9 @@ use serial_test::serial; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; @@ -47,11 +45,7 @@ use aimux_providers::{ProviderOptions, provider, provider_from_env}; /// TS `TEST_PROMPT`: a single user text message "Hello". fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn default_options(prompt: LanguageModelPrompt) -> CallOptions { @@ -285,61 +279,6 @@ async fn sends_correct_request_body() { assert_eq!(body["messages"][0]["content"], "Hello"); } -/// TS: a tool-call assistant message round-trips as `tool_calls` and a tool -/// result as a `tool` role message (shared OpenAI conversion). -#[tokio::test] -async fn converts_tool_call_and_tool_result_messages() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/chat/completions")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_completion_body())) - .mount(&server) - .await; - - let model = make_provider(&server); - - let prompt: LanguageModelPrompt = vec![ - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::ToolCall { - tool_call_id: "call-1".to_string(), - tool_name: "get_weather".to_string(), - input: json!({"location": "SF"}), - provider_executed: None, - provider_options: None, - }], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::ToolResult { - tool_call_id: "call-1".to_string(), - result: json!({"temp": 72}), - tool_name: None, - is_error: None, - preliminary: None, - dynamic: None, - provider_options: None, - }], - ..Default::default() - }, - ]; - let _ = model.do_generate(&default_options(prompt)).await.unwrap(); - - let requests = server.received_requests().await.expect("requests recorded"); - let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - // Assistant tool-call message. - assert_eq!(body["messages"][0]["role"], "assistant"); - assert_eq!(body["messages"][0]["tool_calls"][0]["id"], "call-1"); - assert_eq!( - body["messages"][0]["tool_calls"][0]["function"]["name"], - "get_weather" - ); - // Tool result message. - assert_eq!(body["messages"][1]["role"], "tool"); - assert_eq!(body["messages"][1]["tool_call_id"], "call-1"); -} - // ════════════════════════════════════════════════════════════════════════════ // Usage conversion (convert-alibaba-usage.test.ts + chat model usage tests) // ════════════════════════════════════════════════════════════════════════════ diff --git a/aimux-providers/tests/anthropic_aws_model_test.rs b/aimux-providers/tests/anthropic_aws_model_test.rs deleted file mode 100644 index 1c0dc0bfe..000000000 --- a/aimux-providers/tests/anthropic_aws_model_test.rs +++ /dev/null @@ -1,469 +0,0 @@ -//! Wiremock tests for the Anthropic-AWS (Claude Platform on AWS) provider. -//! -//! Tests cover: -//! - Non-streaming text generation (Anthropic Messages API JSON response) -//! - Non-streaming tool call extraction -//! - Streaming via SSE (Anthropic message events) -//! - API key authentication -//! - Error handling - -use aimux_core::tool::RawToolCall; -use futures::StreamExt; -use serde_json::{Value, json}; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::content::ContentPart; -use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; -use aimux_core::options::CallOptions; -use aimux_core::result::{GenerateContent, StreamResult}; -use aimux_core::stream_part::StreamPart; -use aimux_core::types::FinishReasonUnified; - -use aimux_providers::anthropic_aws::{AnthropicAwsAuth, AnthropicAwsConfig, AnthropicAwsModel}; - -// ── Helpers ────────────────────────────────────────────────────────────────── - -fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] -} - -fn default_options(prompt: LanguageModelPrompt) -> CallOptions { - CallOptions::new(prompt) -} - -fn make_model(server: &MockServer) -> AnthropicAwsModel { - AnthropicAwsModel::new( - "claude-sonnet-4-20250514".to_string(), - AnthropicAwsConfig { - base_url: server.uri(), - auth: AnthropicAwsAuth::ApiKey("test-api-key".to_string()), - api_version: "2023-06-01".to_string(), - workspace_id: None, - api_key_source: None, - retry_config: aimux_provider_utils::RetryConfig::default(), - }, - ) -} - -async fn mock_messages_json(server: &MockServer, status: u16, body: Value) { - Mock::given(method("POST")) - .and(path("/messages")) - .respond_with(ResponseTemplate::new(status).set_body_json(body)) - .mount(server) - .await; -} - -async fn mock_messages_sse(server: &MockServer, sse_body: &str) { - Mock::given(method("POST")) - .and(path("/messages")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body.to_string()), - ) - .mount(server) - .await; -} - -fn sse(data: &Value) -> String { - format!("data: {data}\n\n") -} - -fn sse_stream(events: &[Value]) -> String { - events.iter().map(sse).collect() -} - -fn as_text(item: &GenerateContent) -> &str { - match item { - GenerateContent::Text { text, .. } => text, - _ => panic!("expected Text content, got {item:?}"), - } -} - -fn as_tool_call(item: &GenerateContent) -> (&str, &str, &str) { - match item { - GenerateContent::ToolCall(RawToolCall { - tool_call_id, - tool_name, - input, - .. - }) => (tool_call_id, tool_name, input), - _ => panic!("expected ToolCall content, got {item:?}"), - } -} - -async fn collect_stream(result: StreamResult) -> Vec { - let mut parts = Vec::new(); - let mut stream = result.stream; - while let Some(part) = stream.next().await { - match part { - Ok(p) => parts.push(p), - Err(e) => panic!("stream error: {e:?}"), - } - } - parts -} - -fn text_response(text: &str) -> Value { - json!({ - "id": "msg_017TfcQ4AgGxKyBduUpqYPZn", - "type": "message", - "role": "assistant", - "content": [{ "type": "text", "text": text }], - "model": "claude-sonnet-4-20250514", - "stop_reason": "end_turn", - "stop_sequence": null, - "usage": { "input_tokens": 4, "output_tokens": 30 } - }) -} - -// ═════════════════════════════════════════════════════════════════════════════ -// doGenerate tests -// ═════════════════════════════════════════════════════════════════════════════ - -/// Test: non-streaming tool call extraction. -#[tokio::test] -async fn anthropic_aws_generate_tool_call() { - let server = MockServer::start().await; - mock_messages_json( - &server, - 200, - json!({ - "id": "msg_001", - "type": "message", - "role": "assistant", - "content": [ - { "type": "text", "text": "I'll check that for you." }, - { - "type": "tool_use", - "id": "toolu_01A", - "name": "getWeather", - "input": { "location": "Paris" } - } - ], - "model": "claude-sonnet-4-20250514", - "stop_reason": "tool_use", - "usage": { "input_tokens": 10, "output_tokens": 20 } - }), - ) - .await; - - let model = make_model(&server); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - assert_eq!(result.content.len(), 2); - assert_eq!(as_text(&result.content[0]), "I'll check that for you."); - - let (id, name, input) = as_tool_call(&result.content[1]); - assert_eq!(id, "toolu_01A"); - assert_eq!(name, "getWeather"); - assert_eq!(input, r#"{"location":"Paris"}"#); - assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); -} - -/// Test: HTTP error status is propagated. -#[tokio::test] -async fn anthropic_aws_generate_error() { - let server = MockServer::start().await; - mock_messages_json( - &server, - 429, - json!({ - "type": "error", - "error": { "type": "rate_limit_error", "message": "Too many requests" } - }), - ) - .await; - - let model = make_model(&server); - let result = model.do_generate(&default_options(test_prompt())).await; - - assert!(result.is_err(), "should return error for 429 status"); -} - -// ═════════════════════════════════════════════════════════════════════════════ -// doStream tests -// ═════════════════════════════════════════════════════════════════════════════ - -/// Test: streaming text via Anthropic SSE events. -#[tokio::test] -async fn anthropic_aws_stream_text() { - let server = MockServer::start().await; - - let sse_body = sse_stream(&[ - json!({ - "type": "message_start", - "message": { - "id": "msg_001", - "model": "claude-sonnet-4-20250514", - "usage": { "input_tokens": 3 } - } - }), - json!({ - "type": "content_block_start", - "index": 0, - "content_block": { "type": "text", "text": "" } - }), - json!({ - "type": "content_block_delta", - "index": 0, - "delta": { "type": "text_delta", "text": "Hello" } - }), - json!({ - "type": "content_block_delta", - "index": 0, - "delta": { "type": "text_delta", "text": " streamed!" } - }), - json!({ - "type": "content_block_stop", - "index": 0 - }), - json!({ - "type": "message_delta", - "delta": { "stop_reason": "end_turn" }, - "usage": { "output_tokens": 5 } - }), - json!({ "type": "message_stop" }), - ]); - - mock_messages_sse(&server, &sse_body).await; - - let model = make_model(&server); - let result = model - .do_stream(&default_options(test_prompt())) - .await - .expect("do_stream should succeed"); - - let parts = collect_stream(result).await; - - let text_deltas: Vec = parts - .iter() - .filter_map(|p| match p { - StreamPart::TextDelta { delta, .. } => Some(delta.clone()), - _ => None, - }) - .collect(); - assert_eq!(text_deltas, vec!["Hello", " streamed!"]); - - let finish = parts.iter().find_map(|p| match p { - StreamPart::Finish { finish_reason, .. } => Some(finish_reason.clone()), - _ => None, - }); - assert!(finish.is_some()); - assert_eq!(finish.unwrap().unified, FinishReasonUnified::Stop); -} - -/// Test: streaming tool calls via Anthropic SSE events. -#[tokio::test] -async fn anthropic_aws_stream_tool_call() { - let server = MockServer::start().await; - - let sse_body = sse_stream(&[ - json!({ - "type": "message_start", - "message": { - "id": "msg_002", - "model": "claude-sonnet-4-20250514", - "usage": { "input_tokens": 5 } - } - }), - json!({ - "type": "content_block_start", - "index": 0, - "content_block": { - "type": "tool_use", - "id": "toolu_01B", - "name": "getWeather", - "input": {} - } - }), - json!({ - "type": "content_block_delta", - "index": 0, - "delta": { "type": "input_json_delta", "partial_json": "{\"location\":" } - }), - json!({ - "type": "content_block_delta", - "index": 0, - "delta": { "type": "input_json_delta", "partial_json": "\"Berlin\"}" } - }), - json!({ - "type": "content_block_stop", - "index": 0 - }), - json!({ - "type": "message_delta", - "delta": { "stop_reason": "tool_use" }, - "usage": { "output_tokens": 10 } - }), - json!({ "type": "message_stop" }), - ]); - - mock_messages_sse(&server, &sse_body).await; - - let model = make_model(&server); - let result = model - .do_stream(&default_options(test_prompt())) - .await - .expect("do_stream should succeed"); - - let parts = collect_stream(result).await; - - let tool_calls: Vec<_> = parts - .iter() - .filter_map(|p| match p { - StreamPart::ToolCall(RawToolCall { - tool_call_id, - tool_name, - input, - .. - }) => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), - _ => None, - }) - .collect(); - assert_eq!(tool_calls.len(), 1); - assert_eq!(tool_calls[0].0, "toolu_01B"); - assert_eq!(tool_calls[0].1, "getWeather"); - assert_eq!(tool_calls[0].2, r#"{"location":"Berlin"}"#); -} - -/// Test: SigV4 authentication adds Authorization header. -#[tokio::test] -async fn anthropic_aws_sigv4_auth() { - let server = MockServer::start().await; - mock_messages_json(&server, 200, text_response("Signed!")).await; - - let model = AnthropicAwsModel::new( - "claude-sonnet-4-20250514".to_string(), - AnthropicAwsConfig { - base_url: server.uri(), - auth: AnthropicAwsAuth::SigV4(aimux_providers::bedrock::AwsCredentials { - access_key_id: "AKIAIOSFODNN7EXAMPLE".to_string(), - secret_access_key: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".to_string(), - session_token: None, - region: "us-east-1".to_string(), - }), - api_version: "2023-06-01".to_string(), - workspace_id: None, - api_key_source: None, - retry_config: aimux_provider_utils::RetryConfig::default(), - }, - ); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed with SigV4"); - - assert_eq!(as_text(&result.content[0]), "Signed!"); -} - -// ═════════════════════════════════════════════════════════════════════════════ -// Additional cases — finish reasons, settings, usage, headers, errors. -// ═════════════════════════════════════════════════════════════════════════════ - -/// TS: max_tokens stop reason → Length -#[tokio::test] -async fn anthropic_aws_generate_finish_reason_max_tokens() { - let server = MockServer::start().await; - mock_messages_json( - &server, - 200, - json!({ - "id": "msg_max", - "type": "message", - "role": "assistant", - "content": [{ "type": "text", "text": "truncated" }], - "model": "claude-sonnet-4-20250514", - "stop_reason": "max_tokens", - "usage": { "input_tokens": 4, "output_tokens": 30 } - }), - ) - .await; - - let model = make_model(&server); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - assert_eq!(result.finish_reason.unified, FinishReasonUnified::Length); - assert_eq!(result.finish_reason.raw.as_deref(), Some("max_tokens")); -} - -/// TS: an unknown stop reason → Other -#[tokio::test] -async fn anthropic_aws_generate_finish_reason_unknown() { - let server = MockServer::start().await; - mock_messages_json( - &server, - 200, - json!({ - "id": "msg_unk", - "type": "message", - "role": "assistant", - "content": [{ "type": "text", "text": "hi" }], - "model": "claude-sonnet-4-20250514", - "stop_reason": "pause_turn", - "usage": { "input_tokens": 4, "output_tokens": 30 } - }), - ) - .await; - - let model = make_model(&server); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - assert_eq!(result.finish_reason.unified, FinishReasonUnified::Other); - assert_eq!(result.finish_reason.raw.as_deref(), Some("pause_turn")); -} - -/// TS: a 401 response maps to `AiMuxError::ApiCall` (401 in `status_code`). -#[tokio::test] -async fn anthropic_aws_generate_auth_error() { - let server = MockServer::start().await; - mock_messages_json( - &server, - 401, - json!({ "type": "error", "error": { "type": "authentication_error", "message": "invalid x-api-key" } }), - ) - .await; - - let model = make_model(&server); - let result = model.do_generate(&default_options(test_prompt())).await; - assert!( - matches!(result, Err(ref e) if e.status_code() == Some(401)), - "expected Auth, got {result:?}" - ); -} - -/// TS: a 429 response maps to `AiMuxError::ApiCall` (429 in `status_code`). -#[tokio::test] -async fn anthropic_aws_generate_rate_limit_error() { - let server = MockServer::start().await; - mock_messages_json( - &server, - 429, - json!({ "type": "error", "error": { "type": "rate_limit_error", "message": "Too many requests" } }), - ) - .await; - - let model = make_model(&server); - let result = model.do_generate(&default_options(test_prompt())).await; - assert!( - matches!(result, Err(ref e) if e.status_code() == Some(429)), - "expected RateLimited, got {result:?}" - ); -} diff --git a/aimux-providers/tests/anthropic_convert_full_test.rs b/aimux-providers/tests/anthropic_convert_full_test.rs deleted file mode 100644 index a02fb47b0..000000000 --- a/aimux-providers/tests/anthropic_convert_full_test.rs +++ /dev/null @@ -1,712 +0,0 @@ -// Panic convert wrappers are #[deprecated]; these tests still use them. -#![allow(deprecated)] -//! Rust port of the remaining `convert-to-anthropic-prompt.test.ts` cases that -//! exercise features now supported by the Rust data model: mid-conversation -//! system messages, trailing-whitespace trimming, reasoning/thinking parts, URL -//! & base64 file parts, PDF beta, top-level media-type detection, and -//! provider-referenced files. -//! -//! Cases that remain out of scope (they need richer per-part provider options or -//! server-tool content parts not present in the Rust model) are documented in -//! the task summary: cache_control, citations, mid-conversation tool-changes, -//! server tools (web_search/web_fetch/code_execution/mcp/advisor), and -//! tool-result `content` outputs whose inner parts need conversion. - -use aimux_core::content::ContentPart; -use aimux_core::language_model_message::LanguageModelPromptMessage; -use aimux_core::message::Role; -use aimux_core::types::Warning; -use aimux_providers::anthropic::convert::convert_prompt_to_anthropic_full; -use serde_json::{Value, json}; -use std::collections::BTreeSet; - -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -fn msg(role: Role, parts: Vec) -> LanguageModelPromptMessage { - LanguageModelPromptMessage { - role, - content: parts, - provider_options: None, - } -} - -fn text_part(t: &str) -> ContentPart { - ContentPart::text(t) -} - -fn prompt(msgs: Vec) -> Vec { - msgs -} - -fn convert_full( - msgs: Vec, - send_reasoning: bool, -) -> aimux_providers::anthropic::convert::AnthropicPromptConversion { - convert_prompt_to_anthropic_full(&msgs, send_reasoning) -} - -fn assert_other_warning(warnings: &[Warning], message: &str) { - let found = warnings.iter().any(|w| match w { - Warning::Other { message: m } => m == message, - _ => false, - }); - assert!( - found, - "expected a Warning::Other with message {message:?}, got {warnings:?}" - ); -} - -fn reasoning_part(text: &str, signature: Option<&str>) -> ContentPart { - ContentPart::Reasoning { - text: text.to_string(), - signature: signature.map(std::string::ToString::to_string), - provider_options: None, - } -} - -fn file_url(url: &str, media_type: &str) -> ContentPart { - ContentPart::file_url(url.to_string(), media_type.to_string()) -} - -fn file_base64(data: &str, media_type: &str) -> ContentPart { - ContentPart::file_base64(data.to_string(), media_type.to_string()) -} - -fn file_bytes(data: &[u8], media_type: &str) -> ContentPart { - ContentPart::file(data.to_vec(), media_type.to_string()) -} - -fn file_bytes_named(data: &[u8], media_type: &str, filename: &str) -> ContentPart { - ContentPart::File { - data: data.to_vec(), - media_type: media_type.to_string(), - filename: Some(filename.to_string()), - provider_options: None, - } -} - -fn file_ref(media_type: &str, reference: Value) -> ContentPart { - ContentPart::file_reference(media_type.to_string(), reference) -} - -// PNG magic bytes and their base64 encoding (matches the TS `pngBase64`). -const PNG_BYTES: [u8; 8] = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]; -const PNG_BASE64: &str = "iVBORw0KGgo="; -// PDF magic bytes and their base64 encoding (matches the TS `pdfBase64`). -const PDF_BYTES: [u8; 5] = [0x25, 0x50, 0x44, 0x46, 0x2d]; -const PDF_BASE64: &str = "JVBERi0="; - -// --------------------------------------------------------------------------- -// mid-conversation system -// --------------------------------------------------------------------------- - -#[test] -fn should_emit_a_mid_conversation_system_message_inline_and_add_the_beta() { - let p = prompt(vec![ - msg(Role::System, vec![text_part("initial")]), - msg(Role::User, vec![text_part("hi")]), - msg(Role::Assistant, vec![text_part("hello")]), - msg(Role::System, vec![text_part("switch tone")]), - msg(Role::User, vec![text_part("go")]), - ]); - let result = convert_full(p, true); - assert_eq!( - result.system, - Some(vec![json!({ "type": "text", "text": "initial" })]) - ); - assert!( - result.messages.iter().any(|m| m - == &json!({ - "role": "system", - "content": [{ "type": "text", "text": "switch tone" }] - })), - "expected an inline system message, got {:?}", - result.messages - ); - assert!(result.betas.contains("mid-conversation-system-2026-04-07")); -} - -// --------------------------------------------------------------------------- -// trailing whitespace -// --------------------------------------------------------------------------- - -#[test] -fn should_remove_trailing_whitespace_from_last_assistant_message_when_no_further_user_message() { - let p = prompt(vec![ - msg(Role::User, vec![text_part("user content")]), - msg(Role::Assistant, vec![text_part("assistant content ")]), - ]); - let result = convert_full(p, true); - assert_eq!( - result.messages, - vec![ - json!({ "role": "user", "content": [{ "type": "text", "text": "user content" }] }), - json!({ "role": "assistant", "content": [{ "type": "text", "text": "assistant content" }] }), - ] - ); -} - -#[test] -fn should_remove_trailing_whitespace_from_last_assistant_message_with_multi_part_content() { - let p = prompt(vec![ - msg(Role::User, vec![text_part("user content")]), - msg( - Role::Assistant, - vec![text_part("assistant "), text_part("content ")], - ), - ]); - let result = convert_full(p, true); - assert_eq!( - result.messages, - vec![ - json!({ "role": "user", "content": [{ "type": "text", "text": "user content" }] }), - json!({ - "role": "assistant", - "content": [ - { "type": "text", "text": "assistant " }, - { "type": "text", "text": "content" }, - ] - }), - ] - ); -} - -#[test] -fn should_keep_trailing_whitespace_from_assistant_message_when_there_is_a_further_user_message() { - let p = prompt(vec![ - msg(Role::User, vec![text_part("user content")]), - msg(Role::Assistant, vec![text_part("assistant content ")]), - msg(Role::User, vec![text_part("user content 2")]), - ]); - let result = convert_full(p, true); - assert_eq!( - result.messages, - vec![ - json!({ "role": "user", "content": [{ "type": "text", "text": "user content" }] }), - json!({ "role": "assistant", "content": [{ "type": "text", "text": "assistant content " }] }), - json!({ "role": "user", "content": [{ "type": "text", "text": "user content 2" }] }), - ] - ); -} - -// --------------------------------------------------------------------------- -// reasoning / thinking -// --------------------------------------------------------------------------- - -#[test] -fn should_convert_reasoning_parts_with_signature_into_thinking_parts_when_send_reasoning_is_true() { - let p = prompt(vec![msg( - Role::Assistant, - vec![ - reasoning_part( - "I need to count the number of \"r\"s in the word \"strawberry\".", - Some("test-signature"), - ), - text_part("The word \"strawberry\" has 2 \"r\"s."), - ], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages, - vec![json!({ - "role": "assistant", - "content": [ - { - "type": "thinking", - "thinking": "I need to count the number of \"r\"s in the word \"strawberry\".", - "signature": "test-signature", - }, - { "type": "text", "text": "The word \"strawberry\" has 2 \"r\"s." }, - ] - })] - ); - assert!(result.warnings.is_empty()); -} - -#[test] -fn should_ignore_reasoning_parts_without_signature_when_send_reasoning_is_true() { - let p = prompt(vec![msg( - Role::Assistant, - vec![ - reasoning_part( - "I need to count the number of \"r\"s in the word \"strawberry\".", - None, - ), - text_part("The word \"strawberry\" has 2 \"r\"s."), - ], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages, - vec![json!({ - "role": "assistant", - "content": [ - { "type": "text", "text": "The word \"strawberry\" has 2 \"r\"s." }, - ] - })] - ); - assert_eq!(result.warnings.len(), 1); - assert_other_warning(&result.warnings, "unsupported reasoning metadata"); -} - -#[test] -fn should_omit_reasoning_parts_with_signature_when_send_reasoning_is_false() { - let p = prompt(vec![msg( - Role::Assistant, - vec![ - reasoning_part( - "I need to count the number of \"r\"s in the word \"strawberry\".", - Some("test-signature"), - ), - text_part("The word \"strawberry\" has 2 \"r\"s."), - ], - )]); - let result = convert_full(p, false); - assert_eq!( - result.messages, - vec![json!({ - "role": "assistant", - "content": [ - { "type": "text", "text": "The word \"strawberry\" has 2 \"r\"s." }, - ] - })] - ); - assert_eq!(result.warnings.len(), 1); - assert_other_warning( - &result.warnings, - "sending reasoning content is disabled for this model", - ); -} - -#[test] -fn should_omit_reasoning_parts_without_signature_when_send_reasoning_is_false() { - let p = prompt(vec![msg( - Role::Assistant, - vec![ - reasoning_part( - "I need to count the number of \"r\"s in the word \"strawberry\".", - None, - ), - text_part("The word \"strawberry\" has 2 \"r\"s."), - ], - )]); - let result = convert_full(p, false); - assert_eq!( - result.messages, - vec![json!({ - "role": "assistant", - "content": [ - { "type": "text", "text": "The word \"strawberry\" has 2 \"r\"s." }, - ] - })] - ); - assert_eq!(result.warnings.len(), 1); - assert_other_warning( - &result.warnings, - "sending reasoning content is disabled for this model", - ); -} - -// --------------------------------------------------------------------------- -// Extended-thinking input echo (build_request_body_with_warnings wiring) -// --------------------------------------------------------------------------- - -// --------------------------------------------------------------------------- -// URL images -// --------------------------------------------------------------------------- - -#[test] -fn should_add_image_parts_for_url_images() { - let p = prompt(vec![msg( - Role::User, - vec![file_url("https://example.com/image.png", "image/*")], - )]); - let result = convert_full(p, true); - assert!(result.system.is_none()); - assert_eq!( - result.messages, - vec![json!({ - "role": "user", - "content": [{ - "type": "image", - "source": { "type": "url", "url": "https://example.com/image.png" }, - }], - })] - ); - assert!(result.betas.is_empty()); -} - -#[test] -fn should_treat_url_strings_in_image_file_data_as_urls_not_base64() { - let p = prompt(vec![msg( - Role::User, - vec![file_url("https://example.com/image.png", "image/png")], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages, - vec![json!({ - "role": "user", - "content": [{ - "type": "image", - "source": { "type": "url", "url": "https://example.com/image.png" }, - }], - })] - ); -} - -#[test] -fn passes_through_url_for_top_level_only_image() { - let p = prompt(vec![msg( - Role::User, - vec![file_url("https://example.com/x.png", "image")], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0], - json!({ - "type": "image", - "source": { "type": "url", "url": "https://example.com/x.png" }, - }) - ); -} - -// --------------------------------------------------------------------------- -// URL PDFs -// --------------------------------------------------------------------------- - -#[test] -fn should_treat_url_strings_in_pdf_file_data_as_urls_not_base64() { - let p = prompt(vec![msg( - Role::User, - vec![file_url( - "https://example.com/document.pdf", - "application/pdf", - )], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages, - vec![json!({ - "role": "user", - "content": [{ - "type": "document", - "source": { "type": "url", "url": "https://example.com/document.pdf" }, - }], - })] - ); - let mut expected_betas = BTreeSet::new(); - expected_betas.insert("pdfs-2024-09-25".to_string()); - assert_eq!(result.betas, expected_betas); -} - -#[test] -fn should_add_pdf_file_parts_for_url_pdfs() { - let p = prompt(vec![msg( - Role::User, - vec![file_url( - "https://example.com/document.pdf", - "application/pdf", - )], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0], - json!({ - "type": "document", - "source": { "type": "url", "url": "https://example.com/document.pdf" }, - }) - ); - assert!(result.betas.contains("pdfs-2024-09-25")); -} - -// --------------------------------------------------------------------------- -// base64 PDFs -// --------------------------------------------------------------------------- - -#[test] -fn should_add_pdf_file_parts_for_base64_pdfs() { - // The TS test uses the opaque placeholder string 'base64PDFdata'; the Rust - // `FileBase64` variant holds the base64 string verbatim (it is NOT - // decoded/re-encoded), so the same placeholder round-trips exactly. - let p = prompt(vec![msg( - Role::User, - vec![file_base64("base64PDFdata", "application/pdf")], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages, - vec![json!({ - "role": "user", - "content": [{ - "type": "document", - "source": { - "type": "base64", - "media_type": "application/pdf", - "data": "base64PDFdata", - }, - }], - })] - ); - assert!(result.betas.contains("pdfs-2024-09-25")); -} - -// --------------------------------------------------------------------------- -// text/plain documents (filename -> title) -// --------------------------------------------------------------------------- - -#[test] -fn should_add_text_file_parts_for_text_plain_documents() { - let p = prompt(vec![msg( - Role::User, - vec![file_bytes_named( - b"sample text content", - "text/plain", - "sample.txt", - )], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0], - json!({ - "type": "document", - "source": { "type": "text", "media_type": "text/plain", "data": "sample text content" }, - "title": "sample.txt", - }) - ); - assert!(result.betas.is_empty()); -} - -#[test] -fn should_map_inline_text_file_parts_to_inline_text_document_source() { - let p = prompt(vec![msg( - Role::User, - vec![file_bytes_named( - b"inline text content", - "text/plain", - "inline.txt", - )], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0], - json!({ - "type": "document", - "source": { "type": "text", "media_type": "text/plain", "data": "inline text content" }, - "title": "inline.txt", - }) - ); -} - -#[test] -fn still_routes_text_plain_inline_text_through_document_source() { - let p = prompt(vec![msg( - Role::User, - vec![file_bytes(b"hello", "text/plain")], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0], - json!({ - "type": "document", - "source": { "type": "text", "media_type": "text/plain", "data": "hello" }, - }) - ); -} - -// --------------------------------------------------------------------------- -// unsupported file types -// --------------------------------------------------------------------------- - -#[test] -#[should_panic(expected = "media type: video/mp4")] -fn should_throw_error_for_unsupported_file_types() { - let p = prompt(vec![msg( - Role::User, - vec![file_base64("base64data", "video/mp4")], - )]); - let _ = convert_full(p, true); -} - -// --------------------------------------------------------------------------- -// top-level-only media type resolution (byte sniffing) -// --------------------------------------------------------------------------- - -#[test] -fn detects_image_subtype_from_inline_bytes_for_top_level_image() { - let p = prompt(vec![msg( - Role::User, - vec![file_base64(PNG_BASE64, "image")], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0], - json!({ - "type": "image", - "source": { "type": "base64", "media_type": "image/png", "data": PNG_BASE64 }, - }) - ); -} - -#[test] -fn normalizes_image_star_via_detection() { - let p = prompt(vec![msg( - Role::User, - vec![file_base64(PNG_BASE64, "image/*")], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0]["source"]["media_type"], - json!("image/png") - ); -} - -#[test] -fn detects_pdf_subtype_from_inline_bytes_for_top_level_application() { - let p = prompt(vec![msg( - Role::User, - vec![file_base64(PDF_BASE64, "application")], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0]["source"]["media_type"], - json!("application/pdf") - ); - assert!(result.betas.contains("pdfs-2024-09-25")); -} - -#[test] -fn preserves_full_image_png_pass_through() { - let p = prompt(vec![msg( - Role::User, - vec![file_base64(PNG_BASE64, "image/png")], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0], - json!({ - "type": "image", - "source": { "type": "base64", "media_type": "image/png", "data": PNG_BASE64 }, - }) - ); -} - -#[test] -fn detects_image_subtype_from_inline_bytes_for_top_level_image_via_file_bytes() { - // Same as above but through the raw-bytes `File` variant, which re-encodes - // the bytes to base64 (rather than passing a string through verbatim). - let p = prompt(vec![msg(Role::User, vec![file_bytes(&PNG_BYTES, "image")])]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0], - json!({ - "type": "image", - "source": { "type": "base64", "media_type": "image/png", "data": PNG_BASE64 }, - }) - ); -} - -#[test] -fn detects_pdf_subtype_from_inline_bytes_via_file_bytes() { - let p = prompt(vec![msg( - Role::User, - vec![file_bytes(&PDF_BYTES, "application")], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0]["source"]["media_type"], - json!("application/pdf") - ); - assert!(result.betas.contains("pdfs-2024-09-25")); -} - -// --------------------------------------------------------------------------- -// provider-referenced files -// --------------------------------------------------------------------------- - -#[test] -fn should_convert_messages_with_image_file_parts_using_provider_reference() { - let p = prompt(vec![msg( - Role::User, - vec![file_ref( - "image/png", - json!({ "anthropic": "file-img-12345" }), - )], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages, - vec![json!({ - "role": "user", - "content": [{ - "type": "image", - "source": { "type": "file", "file_id": "file-img-12345" }, - }], - })] - ); - assert!(result.betas.contains("files-api-2025-04-14")); -} - -#[test] -fn should_convert_messages_with_pdf_file_parts_using_provider_reference() { - let p = prompt(vec![msg( - Role::User, - vec![file_ref( - "application/pdf", - json!({ "anthropic": "file-pdf-12345" }), - )], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages[0]["content"][0], - json!({ - "type": "document", - "source": { "type": "file", "file_id": "file-pdf-12345" }, - }) - ); - assert!(result.betas.contains("files-api-2025-04-14")); -} - -#[test] -fn should_convert_provider_referenced_file_parts_to_container_uploads_when_requested() { - let part = ContentPart::FileReference { - media_type: "text/csv".to_string(), - reference: json!({ "anthropic": "file-csv-12345" }), - filename: None, - provider_options: Some( - serde_json::from_value(json!({ "anthropic": { "containerUpload": true } })).unwrap(), - ), - }; - let p = prompt(vec![msg( - Role::User, - vec![text_part("Analyze this data."), part], - )]); - let result = convert_full(p, true); - assert_eq!( - result.messages, - vec![json!({ - "role": "user", - "content": [ - { "type": "text", "text": "Analyze this data." }, - { "type": "container_upload", "file_id": "file-csv-12345" }, - ], - })] - ); - assert!(result.betas.contains("files-api-2025-04-14")); -} - -#[test] -#[should_panic(expected = "No provider reference found for provider 'anthropic'")] -fn should_throw_when_provider_reference_does_not_contain_anthropic_key() { - let p = prompt(vec![msg( - Role::User, - vec![file_ref("application/pdf", json!({ "openai": "file-xyz" }))], - )]); - let _ = convert_full(p, true); -} - -// Keep the `Value` import used even if some tests only use `json!`. -#[allow(dead_code)] -fn _ensure_value_import_used(_v: Value) {} diff --git a/aimux-providers/tests/anthropic_remaining_test.rs b/aimux-providers/tests/anthropic_remaining_test.rs deleted file mode 100644 index ed6c7545b..000000000 --- a/aimux-providers/tests/anthropic_remaining_test.rs +++ /dev/null @@ -1,675 +0,0 @@ -//! Remaining Anthropic provider tests — ported from the TS SDK suite. -//! -//! Covers the TS tests not already exercised by the existing six Anthropic -//! test files (`anthropic_model_test.rs`, `anthropic_convert_test.rs`, -//! `anthropic_convert_full_test.rs`, `anthropic_prepare_tools_test.rs`, -//! `anthropic_provider_tools_test.rs`, `anthropic_cache_control_test.rs`, -//! `anthropic_aws_model_test.rs`): -//! -//! - `anthropic-provider.test.ts` — provider configuration (baseURL, auth, -//! custom provider name, `supportedUrls`). -//! - `anthropic-unknown-model-max-output-tokens.test.ts` — unknown-model -//! `max_tokens` defaulting + compatibility warning. -//! - `sanitize-json-schema.test.ts` — JSON Schema sanitization (the -//! `sanitize_json_schema.rs` implementation had no dedicated tests). -//! - `anthropic-language-model.test.ts` → `mid-conversation tool changes` -//! describe block. -//! -//! HTTP is mocked with `wiremock` (a real loopback HTTP server). Env-var -//! tests are `#[serial]` and bracketed with `unsafe` `set_var`/`remove_var` -//! per the workspace convention. -//! -//! Tests for behaviour the Rust implementation does not yet expose -//! (`supportedUrls`, `toolChanges`, empty-`baseURL` rejection) are written -//! against the intended contract and marked `#[ignore]` with a -//! `// TODO: implementation gap` note, so they document the gap without -//! breaking the suite. - -use serde_json::{Value, json}; -use serial_test::serial; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::content::ContentPart; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; -use aimux_core::options::CallOptions; -use aimux_core::types::Warning; - -use aimux_providers::anthropic::AnthropicConfig; -use aimux_providers::anthropic::model::AnthropicModel; -use aimux_providers::anthropic::sanitize_json_schema::sanitize_json_schema; - -// ───────────────────────────────────────────────────────────────────────────── -// Shared helpers -// ───────────────────────────────────────────────────────────────────────────── - -/// The TS `TEST_PROMPT`: a single user text message "Hello". -fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] -} - -/// Build `CallOptions` with everything unset except `prompt`. -fn default_options(prompt: LanguageModelPrompt) -> CallOptions { - CallOptions::new(prompt) -} - -/// Build an `AnthropicModel` backed by `config` (model id `claude-3-haiku-20240307`). -fn make_model_with_config(config: AnthropicConfig) -> AnthropicModel { - AnthropicModel::new("claude-3-haiku-20240307".to_string(), config) -} - -/// A minimal Anthropic text response body (mirrors the TS `anthropic-text` fixture). -fn text_response(text: &str) -> Value { - json!({ - "id": "msg_123", - "type": "message", - "role": "assistant", - "content": [{ "type": "text", "text": text }], - "model": "claude-3-haiku-20240307", - "stop_reason": "end_turn", - "stop_sequence": null, - "usage": { "input_tokens": 1, "output_tokens": 1 }, - }) -} - -/// Mount a JSON 200 response on `POST /v1/messages`. -async fn mock_messages_ok(server: &MockServer, body: Value) { - Mock::given(method("POST")) - .and(path("/v1/messages")) - .respond_with(ResponseTemplate::new(200).set_body_json(body)) - .mount(server) - .await; -} - -/// Remove the Anthropic env vars consulted by `AnthropicConfig::from_env` -/// (test isolation for `#[serial]` env tests). -fn clear_anthropic_env() { - unsafe { - std::env::remove_var("ANTHROPIC_BASE_URL"); - std::env::remove_var("ANTHROPIC_API_KEY"); - } -} - -/// Set the API key (always required by `from_env`) plus an optional base URL. -fn set_anthropic_env(api_key: &str, base_url: Option<&str>) { - unsafe { - std::env::set_var("ANTHROPIC_API_KEY", api_key); - match base_url { - Some(url) => std::env::set_var("ANTHROPIC_BASE_URL", url), - None => std::env::remove_var("ANTHROPIC_BASE_URL"), - } - } -} - -// ═════════════════════════════════════════════════════════════════════════════ -// baseURL configuration (TS: anthropic-provider.test.ts → 'baseURL configuration') -// ═════════════════════════════════════════════════════════════════════════════ - -/// TS: "uses the default Anthropic base URL when not provided". -/// -/// The TS test mocks `fetch` and asserts the request URL is -/// `https://api.anthropic.com/v1/messages`. The Rust provider builds the -/// endpoint as `{base_url}/v1/messages`, so the equivalent observable -/// behaviour is: the default config's `base_url` is the bare Anthropic URL, -/// and a request against a stubbed server lands on `/v1/messages`. -#[tokio::test] -async fn default_base_url_when_not_provided() { - // Config-level: default base URL is the bare Anthropic API URL. - let config = AnthropicConfig::new("test-api-key"); - assert_eq!(config.base_url, "https://api.anthropic.com"); - - // HTTP-level: the endpoint path is `/v1/messages`. - let server = MockServer::start().await; - mock_messages_ok(&server, text_response("Hi")).await; - let model = - make_model_with_config(AnthropicConfig::new("test-api-key").with_base_url(server.uri())); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - assert_eq!(requests[0].url.path(), "/v1/messages"); -} - -/// TS: "uses ANTHROPIC_BASE_URL when set". -/// -/// The TS test mocks `fetch` and asserts the request URL is -/// `https://proxy.anthropic.example/v1/messages`. The Rust provider builds the -/// endpoint as `{base_url}/v1/messages`, so the observable behaviour is that -/// `from_env` picks up `ANTHROPIC_BASE_URL` and normalizes it into `base_url`. -/// (The HTTP `/v1/messages` path round-trip is covered by -/// `default_base_url_when_not_provided`.) -#[tokio::test] -#[serial] -async fn uses_anthropic_base_url_env_when_set() { - clear_anthropic_env(); - set_anthropic_env("test-api-key", Some("https://proxy.anthropic.example/v1/")); - - let config = AnthropicConfig::from_env().expect("from_env"); - // Trailing slash + `/v1` segment normalized away; the endpoint becomes - // `https://proxy.anthropic.example/v1/messages`. - assert_eq!(config.base_url, "https://proxy.anthropic.example"); - - clear_anthropic_env(); -} - -/// TS: "normalizes a bare Anthropic API URL from ANTHROPIC_BASE_URL". -#[tokio::test] -#[serial] -async fn normalizes_bare_anthropic_url_from_env() { - clear_anthropic_env(); - set_anthropic_env("test-api-key", Some("https://api.anthropic.com/")); - let config = AnthropicConfig::from_env().expect("from_env"); - // Trailing slash stripped; endpoint becomes `https://api.anthropic.com/v1/messages`. - assert_eq!(config.base_url, "https://api.anthropic.com"); - clear_anthropic_env(); -} - -/// TS: "normalizes a bare Anthropic API URL from the baseURL option". -#[tokio::test] -async fn normalizes_bare_anthropic_url_from_base_url_option() { - let config = AnthropicConfig::new("test-api-key").with_base_url("https://api.anthropic.com/"); - assert_eq!(config.base_url, "https://api.anthropic.com"); -} - -/// TS: "prefers the baseURL option over ANTHROPIC_BASE_URL". -#[tokio::test] -#[serial] -async fn prefers_base_url_option_over_env() { - clear_anthropic_env(); - set_anthropic_env("test-api-key", Some("https://env.anthropic.example/v1")); - let config = AnthropicConfig::from_env() - .expect("from_env") - .with_base_url("https://option.anthropic.example/v1/"); - // The explicit option wins; its trailing slash and `/v1` are normalized away. - assert_eq!(config.base_url, "https://option.anthropic.example"); - clear_anthropic_env(); -} - -/// TS: "rejects an empty baseURL option during provider creation". -// TODO: implementation gap — the Rust `normalize_base_url` returns the default -// Anthropic URL for an empty input instead of rejecting it with -// `InvalidArgumentError` (`argument: "baseURL"`, message: "baseURL must be a -// non-empty string."). Re-enable once `with_base_url` / the builder validate -// emptiness. -#[tokio::test] -#[ignore] -async fn rejects_empty_base_url_option() { - let result = AnthropicConfig::builder() - .api_key("test-api-key") - .base_url("") - .build(); - match result { - Err(AiMuxError::InvalidArgument(msg)) => { - assert_eq!(msg, "baseURL must be a non-empty string."); - } - other => panic!("expected InvalidArgument for empty baseURL, got {other:?}"), - } -} - -// ═════════════════════════════════════════════════════════════════════════════ -// Authentication (TS: anthropic-provider.test.ts → 'authentication') -// ═════════════════════════════════════════════════════════════════════════════ - -/// TS: "sends Authorization Bearer header when authToken is provided". -#[tokio::test] -async fn sends_authorization_bearer_when_auth_token_provided() { - let server = MockServer::start().await; - mock_messages_ok(&server, text_response("Hi")).await; - - let config = AnthropicConfig::builder() - .auth_token("test-auth-token") - .base_url(server.uri()) - .build() - .expect("builder build"); - let model = make_model_with_config(config); - - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - assert_eq!( - requests[0] - .headers - .get("authorization") - .and_then(|v| v.to_str().ok()), - Some("Bearer test-auth-token"), - "Authorization Bearer header missing or wrong" - ); - // When an auth token is set, the `x-api-key` header must NOT be sent. - assert!( - requests[0].headers.get("x-api-key").is_none(), - "x-api-key should not be sent when authToken is provided" - ); -} - -/// TS: "throws error when both apiKey and authToken options are provided". -#[test] -fn throws_when_both_api_key_and_auth_token_provided() { - let result = AnthropicConfig::builder() - .api_key("test-api-key") - .auth_token("test-auth-token") - .build(); - match result { - Err(AiMuxError::InvalidArgument(msg)) => { - assert_eq!( - msg, - "Both apiKey and authToken were provided. \ - Please use only one authentication method." - ); - } - other => { - panic!("expected InvalidArgument for conflicting apiKey + authToken, got {other:?}") - } - } -} - -// ═════════════════════════════════════════════════════════════════════════════ -// Custom provider name (TS: anthropic-provider.test.ts → 'custom provider name') -// ═════════════════════════════════════════════════════════════════════════════ - -/// TS: "should use custom provider name when specified". -#[test] -fn uses_custom_provider_name_when_specified() { - let config = AnthropicConfig::new("test-api-key").with_name("my-claude-proxy"); - let model = make_model_with_config(config); - assert_eq!(model.provider(), "my-claude-proxy"); -} - -/// TS: "should default to anthropic.messages when name not specified". -#[test] -fn defaults_to_anthropic_messages_when_not_specified() { - let config = AnthropicConfig::new("test-api-key"); - let model = make_model_with_config(config); - assert_eq!(model.provider(), "anthropic.messages"); -} - -// ═════════════════════════════════════════════════════════════════════════════ -// supportedUrls (TS: anthropic-provider.test.ts → 'supportedUrls') -// ═════════════════════════════════════════════════════════════════════════════ - -/// TS: "should support image/* URLs". -// TODO: implementation gap — the Rust `LanguageModel` trait has no -// `supported_urls` surface, so there is nothing to assert against. Re-enable -// once `supported_urls` is added to the trait and the Anthropic model exposes -// its image/PDF URL allow-list. -#[tokio::test] -#[ignore] -async fn supports_image_urls() { - let config = AnthropicConfig::new("test-api-key"); - let model = make_model_with_config(config); - // Intended: `model.supported_urls()["image/*"]` matches - // `https://example.com/image.png`. - let _ = model.provider(); -} - -/// TS: "should support application/pdf URLs". -// TODO: implementation gap — see `supports_image_urls`. -#[tokio::test] -#[ignore] -async fn supports_pdf_urls() { - let config = AnthropicConfig::new("test-api-key"); - let model = make_model_with_config(config); - let _ = model.provider(); -} - -// ═════════════════════════════════════════════════════════════════════════════ -// Unknown-model max output tokens (TS: anthropic-unknown-model-max-output-tokens.test.ts) -// ═════════════════════════════════════════════════════════════════════════════ - -/// Build a model with the given model id pointing at the wiremock server. -fn make_model_with_id(server: &MockServer, model_id: &str) -> AnthropicModel { - AnthropicModel::new( - model_id.to_string(), - AnthropicConfig::new("test-api-key").with_base_url(server.uri()), - ) -} - -/// TS: "should warn when using the default max output token limit". -#[tokio::test] -async fn warns_when_using_default_max_output_token_limit() { - let server = MockServer::start().await; - mock_messages_ok(&server, text_response("Hello!")).await; - - let model = make_model_with_id(&server, "future-model"); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert_eq!(body["model"], "future-model"); - assert_eq!(body["max_tokens"], 4096); - - assert_eq!(result.warnings.len(), 1); - match &result.warnings[0] { - Warning::Compatibility { feature, details } => { - assert_eq!(feature, "maxOutputTokens"); - assert_eq!( - details.as_deref(), - Some( - "The model \"future-model\" is unknown. The max output tokens have been \ - limited to 4096. Set maxOutputTokens explicitly to override this limit." - ) - ); - } - other => panic!("expected Compatibility warning, got {other:?}"), - } -} - -/// TS: "should not warn when max output tokens are provided". -#[tokio::test] -async fn does_not_warn_when_max_output_tokens_provided() { - let server = MockServer::start().await; - mock_messages_ok(&server, text_response("Hello!")).await; - - let model = make_model_with_id(&server, "future-model"); - let mut opts = default_options(test_prompt()); - opts.max_output_tokens = Some(123456); - let result = model - .do_generate(&opts) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert_eq!(body["model"], "future-model"); - assert_eq!(body["max_tokens"], 123456); - assert!( - result.warnings.is_empty(), - "warnings: {:?}", - result.warnings - ); -} - -/// TS: "should use the current-generation default and warn for an unknown Claude model". -#[tokio::test] -async fn uses_current_gen_default_and_warns_for_unknown_claude_model() { - let server = MockServer::start().await; - mock_messages_ok(&server, text_response("Hello!")).await; - - let model = make_model_with_id(&server, "claude-future-9"); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert_eq!(body["model"], "claude-future-9"); - assert_eq!(body["max_tokens"], 128000); - - assert_eq!(result.warnings.len(), 1); - match &result.warnings[0] { - Warning::Compatibility { feature, details } => { - assert_eq!(feature, "maxOutputTokens"); - assert_eq!( - details.as_deref(), - Some( - "The model \"claude-future-9\" is unknown. The max output tokens have been \ - limited to 128000. Set maxOutputTokens explicitly to override this limit." - ) - ); - } - other => panic!("expected Compatibility warning, got {other:?}"), - } -} - -// ═════════════════════════════════════════════════════════════════════════════ -// sanitize-json-schema (TS: sanitize-json-schema.test.ts) -// ═════════════════════════════════════════════════════════════════════════════ - -/// TS: "strips unsupported number constraints and adds readable descriptions". -#[test] -fn sanitize_strips_unsupported_number_constraints() { - let schema = json!({ - "type": "object", - "properties": { - "recurringIntervalMinutes": { - "type": "number", - "exclusiveMinimum": 0, - "minimum": 1, - "maximum": 60, - "exclusiveMaximum": 120 - } - }, - "required": ["recurringIntervalMinutes"], - "additionalProperties": false - }); - - let expected = json!({ - "additionalProperties": false, - "properties": { - "recurringIntervalMinutes": { - "description": - "minimum: 1; maximum: 60; exclusive minimum: 0; exclusive maximum: 120.", - "type": "number" - } - }, - "required": ["recurringIntervalMinutes"], - "type": "object" - }); - - assert_eq!(sanitize_json_schema(&schema), expected); -} - -/// TS: "strips unsupported string constraints and unsupported formats". -#[test] -fn sanitize_strips_unsupported_string_constraints_and_formats() { - let schema = json!({ - "type": "object", - "properties": { - "slug": { - "type": "string", - "description": "A URL slug", - "minLength": 1, - "maxLength": 20, - "pattern": "^[a-z0-9-]+$", - "format": "regex" - } - } - }); - - let expected = json!({ - "additionalProperties": false, - "properties": { - "slug": { - "description": - "A URL slug\nmin length: 1; max length: 20; pattern: ^[a-z0-9-]+$; format: regex.", - "type": "string" - } - }, - "type": "object" - }); - - assert_eq!(sanitize_json_schema(&schema), expected); -} - -/// TS: "recursively sanitizes arrays, definitions, and composition schemas". -#[test] -fn sanitize_recursively_handles_arrays_defs_and_composition() { - let schema = json!({ - "type": "object", - "$defs": { - "PositiveInteger": { "type": "integer", "minimum": 1 } - }, - "properties": { - "count": { "$ref": "#/$defs/PositiveInteger" }, - "tags": { - "type": "array", - "minItems": 2, - "maxItems": 4, - "uniqueItems": true, - "items": { - "anyOf": [ - { "type": "string", "minLength": 1 }, - { "type": "number", "maximum": 10 } - ] - } - } - } - }); - - let expected = json!({ - "$defs": { - "PositiveInteger": { - "description": "minimum: 1.", - "type": "integer" - } - }, - "additionalProperties": false, - "properties": { - "count": { "$ref": "#/$defs/PositiveInteger" }, - "tags": { - "description": "min items: 2; max items: 4; unique items: true.", - "items": { - "anyOf": [ - { "description": "min length: 1.", "type": "string" }, - { "description": "maximum: 10.", "type": "number" } - ] - }, - "type": "array" - } - }, - "type": "object" - }); - - assert_eq!(sanitize_json_schema(&schema), expected); -} - -/// TS: "converts oneOf to anyOf". -#[test] -fn sanitize_converts_one_of_to_any_of() { - let schema = json!({ - "oneOf": [ - { "type": "string", "minLength": 1 }, - { "type": "number", "minimum": 0 } - ] - }); - - let expected = json!({ - "anyOf": [ - { "description": "min length: 1.", "type": "string" }, - { "description": "minimum: 0.", "type": "number" } - ] - }); - - assert_eq!(sanitize_json_schema(&schema), expected); -} - -/// TS: "does not mutate the input schema". -#[test] -fn sanitize_does_not_mutate_input_schema() { - let schema = json!({ - "type": "object", - "properties": { - "value": { "type": "number", "exclusiveMinimum": 0 } - } - }); - - let snapshot = schema.clone(); - let _ = sanitize_json_schema(&schema); - assert_eq!(schema, snapshot); -} - -// ═════════════════════════════════════════════════════════════════════════════ -// Mid-conversation tool changes -// (TS: anthropic-language-model.test.ts → 'mid-conversation tool changes') -// ═════════════════════════════════════════════════════════════════════════════ - -/// TS: "should send tool change blocks and the beta header". -/// -/// The TS test sends a system message carrying -/// `providerOptions.anthropic.toolChanges = [{ type: 'tool_removal', toolName: -/// 'get_weather' }]` and asserts: -/// - `messages` contains a `system` message whose content is -/// `[{ type: 'tool_removal', tool: { type: 'tool_reference', name: -/// 'get_weather' } }]`. -/// - the `anthropic-beta` request header contains -/// `mid-conversation-tool-changes-2026-07-01`. -// TODO: implementation gap — the Rust convert does not yet read -// `toolChanges` from provider options, emit `tool_removal`/`tool_addition` -// system blocks, or set the `mid-conversation-tool-changes-2026-07-01` beta -// header. Re-enable once that pipeline is implemented. -#[tokio::test] -#[ignore] -async fn sends_tool_change_blocks_and_beta_header() { - let server = MockServer::start().await; - mock_messages_ok(&server, text_response("OK")).await; - - let config = AnthropicConfig::new("test-api-key").with_base_url(server.uri()); - let model = AnthropicModel::new("claude-opus-4-8".to_string(), config); - - // NOTE: when implemented, `toolChanges` should be carried on the system - // message's provider options. The current `LanguageModelPromptMessage` - // / `CallOptions` shape may not yet expose the required provider-options - // slot; that is part of the implementation gap. - let prompt: LanguageModelPrompt = vec![ - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Say OK.")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::System, - content: vec![], - ..Default::default() - }, - ]; - let opts = default_options(prompt); - - let _ = model - .do_generate(&opts) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - - // The system message should carry the tool_removal block. - let system_msgs: Vec<&Value> = body["messages"] - .as_array() - .expect("messages array") - .iter() - .filter(|m| m["role"] == "system") - .collect(); - assert!( - system_msgs.iter().any(|m| { - m["content"] - .as_array() - .map(|c| { - c.iter().any(|block| { - block["type"] == "tool_removal" - && block["tool"]["type"] == "tool_reference" - && block["tool"]["name"] == "get_weather" - }) - }) - .unwrap_or(false) - }), - "expected a system tool_removal block, got body: {body}" - ); - - // The beta header should advertise mid-conversation tool changes. - let beta = requests[0] - .headers - .get("anthropic-beta") - .and_then(|v| v.to_str().ok()) - .unwrap_or(""); - assert!( - beta.contains("mid-conversation-tool-changes-2026-07-01"), - "missing beta header, got: {beta}" - ); -} diff --git a/aimux-providers/tests/azure_model_test.rs b/aimux-providers/tests/azure_model_test.rs deleted file mode 100644 index daab5af47..000000000 --- a/aimux-providers/tests/azure_model_test.rs +++ /dev/null @@ -1,750 +0,0 @@ -//! Wiremock tests for the Azure OpenAI provider. -//! -//! These cover the Azure-specific differences from the OpenAI provider: -//! - deployment-based URL construction with an `api-version` query parameter -//! - `api-key` header authentication (API key) -//! - `Authorization: Bearer ` authentication (Azure AD token provider) -//! - request/response handling (reusing the OpenAI conversion logic) -//! - streaming -//! - error status codes -//! - config validation (missing resource/base URL, missing auth) -//! -//! Each test spins up a `wiremock` mock server, points an `AzureProvider` at -//! it via `base_url`, and asserts on the issued request and/or the parsed -//! result. - -use aimux_core::tool::RawToolCall; -use std::sync::Arc; -use std::sync::atomic::{AtomicU32, Ordering}; - -use async_trait::async_trait; -use futures::StreamExt; -use serde_json::{Value, json}; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::content::ContentPart; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; -use aimux_core::options::CallOptions; -use aimux_core::result::GenerateContent; -use aimux_core::stream_part::StreamPart; -use aimux_core::types::FinishReasonUnified; - -use aimux_providers::{AzureConfig, AzureProvider, TokenProvider}; - -// ── helpers ───────────────────────────────────────────────────────────────── - -/// The TS `TEST_PROMPT`: a single user text message "Hello". -fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] -} - -/// Build `CallOptions` with everything unset except `prompt`. -fn default_options(prompt: LanguageModelPrompt) -> CallOptions { - CallOptions::new(prompt) -} - -/// Standard mock for a JSON chat completion response on a given path. -async fn mock_json_response(server: &MockServer, request_path: &str, body: Value) { - Mock::given(method("POST")) - .and(path(request_path)) - .respond_with(ResponseTemplate::new(200).set_body_json(body)) - .mount(server) - .await; -} - -/// Mock for a JSON error response on a given path with a status code. -async fn mock_json_error(server: &MockServer, request_path: &str, status: u16, body: Value) { - Mock::given(method("POST")) - .and(path(request_path)) - .respond_with(ResponseTemplate::new(status).set_body_json(body)) - .mount(server) - .await; -} - -/// Standard mock for an SSE streaming response on a given path. -async fn mock_sse_response(server: &MockServer, request_path: &str, sse_body: &str) { - Mock::given(method("POST")) - .and(path(request_path)) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body.to_string()), - ) - .mount(server) - .await; -} - -/// A simple text chat-completion response body. -fn text_completion_response(text: &str) -> Value { - json!({ - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1711115037, - "model": "gpt-4o", - "choices": [{ - "index": 0, - "message": { "role": "assistant", "content": text }, - "finish_reason": "stop" - }], - "usage": { "prompt_tokens": 4, "total_tokens": 34, "completion_tokens": 30 } - }) -} - -/// Concatenate SSE events and append `[DONE]`. -fn sse_body(events: &[&str]) -> String { - let mut body = String::new(); - for event in events { - body.push_str(event); - } - body.push_str("data: [DONE]\n\n"); - body -} - -/// Collect all `StreamPart`s from a `StreamResult` into a `Vec`. -async fn collect_stream(result: aimux_core::result::StreamResult) -> Vec { - let mut parts = Vec::new(); - let mut stream = result.stream; - while let Some(part) = stream.next().await { - match part { - Ok(p) => parts.push(p), - Err(e) => panic!("stream error: {e:?}"), - } - } - parts -} - -/// Extract text deltas from a list of stream parts. -fn text_deltas(parts: &[StreamPart]) -> Vec { - parts - .iter() - .filter_map(|p| match p { - StreamPart::TextDelta { delta, .. } => Some(delta.clone()), - _ => None, - }) - .collect() -} - -/// A token provider that returns a fixed token. -struct StaticToken(String); -#[async_trait] -impl TokenProvider for StaticToken { - async fn get_token(&self) -> Result { - Ok(self.0.clone()) - } -} - -/// A token provider that returns a different token on each call -/// (`token-1`, `token-2`, …) so tests can assert it is invoked per request. -struct CountingToken { - count: AtomicU32, -} -#[async_trait] -impl TokenProvider for CountingToken { - async fn get_token(&self) -> Result { - let n = self.count.fetch_add(1, Ordering::SeqCst) + 1; - Ok(format!("token-{n}")) - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// URL construction -// ════════════════════════════════════════════════════════════════════════════ - -/// The default deployment-based URL form places the deployment in the path and -/// appends `?api-version=2024-10-21` (the default api version). -#[tokio::test] -async fn should_build_deployment_based_url_with_default_api_version() { - let server = MockServer::start().await; - mock_json_response( - &server, - "/deployments/gpt-4o/chat/completions", - text_completion_response("hi"), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - let url = &requests[0].url; - assert_eq!(url.path(), "/deployments/gpt-4o/chat/completions"); - assert_eq!(url.query(), Some("api-version=2024-10-21")); -} - -/// A custom `api_version` is reflected in the `api-version` query parameter. -#[tokio::test] -async fn should_use_custom_api_version() { - let server = MockServer::start().await; - mock_json_response( - &server, - "/deployments/gpt-4o/chat/completions", - text_completion_response("hi"), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key") - .with_api_version("2025-04-01-preview"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.unwrap(); - assert_eq!( - requests[0].url.query(), - Some("api-version=2025-04-01-preview") - ); -} - -/// The v1 URL form uses `/v1/chat/completions` (no deployment in the path). -/// With a non-Azure gateway `base_url`, the gateway owns versioning and -/// `api-version` is omitted — mirroring the TS `useAzureOpenAIEndpoint` gate. -/// (The real-Azure v1 case — `api-version` appended — is covered by the unit -/// tests in `azure/model.rs`, which construct the URL from a `resource_name` -/// without needing a live HTTP endpoint.) -#[tokio::test] -async fn should_use_v1_url_form_on_gateway() { - let server = MockServer::start().await; - mock_json_response( - &server, - "/v1/chat/completions", - text_completion_response("hi"), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key") - .use_v1_urls(); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.unwrap(); - assert_eq!(requests[0].url.path(), "/v1/chat/completions"); - // Non-Azure gateway owns its own versioning — no api-version query param. - assert_eq!(requests[0].url.query(), None); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Authentication -// ════════════════════════════════════════════════════════════════════════════ - -/// API-key auth sends the `api-key` header (not a Bearer token). -#[tokio::test] -async fn should_send_api_key_header() { - let server = MockServer::start().await; - mock_json_response( - &server, - "/deployments/gpt-4o/chat/completions", - text_completion_response("hi"), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.unwrap(); - let headers = &requests[0].headers; - assert_eq!( - headers.get("api-key").and_then(|v| v.to_str().ok()), - Some("test-api-key") - ); - // No Authorization header when using api-key auth. - assert!(headers.get("authorization").is_none()); -} - -/// A token provider supplies a Bearer token via the `Authorization` header. -#[tokio::test] -async fn should_send_bearer_token_from_token_provider() { - let server = MockServer::start().await; - mock_json_response( - &server, - "/deployments/gpt-4o/chat/completions", - text_completion_response("hi"), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_token_provider(Arc::new(StaticToken("test-azure-ad-token".to_string()))); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.unwrap(); - let headers = &requests[0].headers; - assert_eq!( - headers.get("authorization").and_then(|v| v.to_str().ok()), - Some("Bearer test-azure-ad-token") - ); - // No api-key header when using token auth. - assert!(headers.get("api-key").is_none()); -} - -/// The token provider is invoked on every request (tokens differ per call). -#[tokio::test] -async fn should_call_token_provider_per_request() { - let server = MockServer::start().await; - mock_json_response( - &server, - "/deployments/gpt-4o/chat/completions", - text_completion_response("hi"), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_token_provider(Arc::new(CountingToken { - count: AtomicU32::new(0), - })); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("first call"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("second call"); - - let requests = server.received_requests().await.unwrap(); - assert_eq!(requests.len(), 2); - assert_eq!( - requests[0] - .headers - .get("authorization") - .and_then(|v| v.to_str().ok()), - Some("Bearer token-1") - ); - assert_eq!( - requests[1] - .headers - .get("authorization") - .and_then(|v| v.to_str().ok()), - Some("Bearer token-2") - ); -} - -/// Provider-level and per-request headers are merged onto the auth headers. -#[tokio::test] -async fn should_merge_provider_and_request_headers() { - let server = MockServer::start().await; - mock_json_response( - &server, - "/deployments/gpt-4o/chat/completions", - text_completion_response("hi"), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key") - .with_header("Custom-Provider-Header", "provider-header-value"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - - let mut options = default_options(test_prompt()); - let mut req_headers = std::collections::HashMap::new(); - req_headers.insert( - "Custom-Request-Header".to_string(), - "request-header-value".to_string(), - ); - options.headers = Some(req_headers); - - let _ = model.do_generate(&options).await.expect("do_generate"); - - let requests = server.received_requests().await.unwrap(); - let headers = &requests[0].headers; - assert_eq!( - headers.get("api-key").and_then(|v| v.to_str().ok()), - Some("test-api-key") - ); - assert_eq!( - headers - .get("custom-provider-header") - .and_then(|v| v.to_str().ok()), - Some("provider-header-value") - ); - assert_eq!( - headers - .get("custom-request-header") - .and_then(|v| v.to_str().ok()), - Some("request-header-value") - ); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Response parsing (reuses OpenAI conversion logic) -// ════════════════════════════════════════════════════════════════════════════ - -/// `do_generate` maps Azure/OpenAI usage tokens. -#[tokio::test] -async fn should_extract_usage() { - let server = MockServer::start().await; - mock_json_response( - &server, - "/deployments/gpt-4o/chat/completions", - text_completion_response(""), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - assert_eq!(result.usage.input_tokens.total, Some(4)); - assert_eq!(result.usage.output_tokens.total, Some(30)); -} - -/// The request body carries the deployment as `model` plus the user message. -#[tokio::test] -async fn should_send_deployment_and_messages_in_request_body() { - let server = MockServer::start().await; - mock_json_response( - &server, - "/deployments/gpt-4o/chat/completions", - text_completion_response("hi"), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .unwrap(); - - let requests = server.received_requests().await.unwrap(); - let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert_eq!(body["model"], "gpt-4o"); - assert_eq!(body["messages"][0]["role"], "user"); - assert_eq!(body["messages"][0]["content"], "Hello"); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Streaming -// ════════════════════════════════════════════════════════════════════════════ - -/// `do_stream` emits text deltas followed by a Finish part. -#[tokio::test] -async fn should_stream_text_deltas() { - let server = MockServer::start().await; - let sse = sse_body(&[ - "data: {\"id\":\"chatcmpl-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n", - "data: {\"id\":\"chatcmpl-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\", World!\"},\"finish_reason\":null}]}\n\n", - "data: {\"id\":\"chatcmpl-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":4,\"completion_tokens\":10,\"total_tokens\":14}}\n\n", - ]); - mock_sse_response(&server, "/deployments/gpt-4o/chat/completions", &sse).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - - let result = model - .do_stream(&default_options(test_prompt())) - .await - .expect("do_stream should succeed"); - - let parts = collect_stream(result).await; - let deltas = text_deltas(&parts); - assert_eq!(deltas, vec!["Hello".to_string(), ", World!".to_string()]); - - // The stream must end with a Finish part carrying the stop reason + usage. - let finish = parts - .iter() - .rev() - .find(|p| matches!(p, StreamPart::Finish { .. })) - .expect("stream should end with Finish"); - match finish { - StreamPart::Finish { - finish_reason, - usage, - .. - } => { - assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); - assert_eq!(usage.output_tokens.total, Some(10)); - } - _ => unreachable!(), - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// Errors & config validation -// ════════════════════════════════════════════════════════════════════════════ - -/// A 401 response maps to `AiMuxError::ApiCall` (401 in `status_code`). -#[tokio::test] -async fn should_return_auth_error_on_401() { - let server = MockServer::start().await; - mock_json_error( - &server, - "/deployments/gpt-4o/chat/completions", - 401, - json!({ "error": { "message": "invalid api key", "type": "authentication_error" } }), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("bad-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - - let err = model - .do_generate(&default_options(test_prompt())) - .await - .expect_err("401 should error"); - assert!( - matches!(err, ref e if e.status_code() == Some(401)), - "{err:?}" - ); -} - -/// Constructing a provider without `resource_name` or `base_url` is rejected. -#[tokio::test] -async fn should_reject_missing_resource_and_base_url() { - let config = AzureConfig::new().with_api_key("test-api-key"); - assert!(matches!( - AzureProvider::new(config), - Err(AiMuxError::InvalidArgument(_)) - )); -} - -/// Calling a model with no auth configured returns an `Auth` error at request -/// time. -#[tokio::test] -async fn should_error_when_no_auth_configured() { - let server = MockServer::start().await; - // No mock is mounted: the request must never be sent because auth fails - // before the HTTP call. - let config = AzureConfig::new().with_base_url(server.uri()); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - - let err = model - .do_generate(&default_options(test_prompt())) - .await - .expect_err("no auth should error"); - assert!(matches!(err, AiMuxError::InvalidArgument(_)), "{err:?}"); - - // No request should have hit the server. - let requests = server.received_requests().await.unwrap(); - assert!(requests.is_empty()); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Additional cases — tool calls, response format, reasoning, headers, errors. -// ════════════════════════════════════════════════════════════════════════════ - -/// A chat-completion response carrying a single tool call. -fn tool_call_completion_body() -> Value { - json!({ - "id": "chatcmpl-tool", - "object": "chat.completion", - "created": 1711115037, - "model": "gpt-4o", - "choices": [{ - "index": 0, - "message": { - "role": "assistant", - "tool_calls": [{ - "id": "call_abc", - "type": "function", - "function": { - "name": "get-weather", - "arguments": "{\"city\":\"SF\"}" - } - }] - }, - "finish_reason": "tool_calls" - }], - "usage": { "prompt_tokens": 10, "total_tokens": 20, "completion_tokens": 10 } - }) -} - -/// TS: tool-call extraction — `do_generate` surfaces a `ToolCall` content. -#[tokio::test] -async fn should_extract_tool_call() { - let server = MockServer::start().await; - mock_json_response( - &server, - "/deployments/gpt-4o/chat/completions", - tool_call_completion_body(), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - assert_eq!(result.content.len(), 1); - match &result.content[0] { - GenerateContent::ToolCall(RawToolCall { - tool_call_id, - tool_name, - input, - .. - }) => { - assert_eq!(tool_call_id, "call_abc"); - assert_eq!(tool_name, "get-weather"); - assert_eq!(input, r#"{"city":"SF"}"#); - } - other => panic!("expected ToolCall, got {other:?}"), - } - assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); -} - -/// TS: "should stream tool call content" — streaming emits a ToolCall part. -#[tokio::test] -async fn should_stream_tool_call() { - let server = MockServer::start().await; - let sse = sse_body(&[ - "data: {\"id\":\"chatcmpl-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"tool_calls\":[{\"index\":0,\"id\":\"call_abc\",\"type\":\"function\",\"function\":{\"name\":\"get-weather\",\"arguments\":\"\"}}]},\"finish_reason\":null}]}\n\n", - "data: {\"id\":\"chatcmpl-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"city\\\":\\\"SF\\\"}\"}}]},\"finish_reason\":null}]}\n\n", - "data: {\"id\":\"chatcmpl-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"tool_calls\"}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":10,\"total_tokens\":20}}\n\n", - ]); - mock_sse_response(&server, "/deployments/gpt-4o/chat/completions", &sse).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - - let result = model - .do_stream(&default_options(test_prompt())) - .await - .expect("do_stream should succeed"); - - let parts = collect_stream(result).await; - let tool_call = parts.iter().find_map(|p| match p { - StreamPart::ToolCall(RawToolCall { - tool_call_id, - tool_name, - input, - .. - }) => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), - _ => None, - }); - let (id, name, input) = tool_call.expect("should have ToolCall"); - assert_eq!(id, "call_abc"); - assert_eq!(name, "get-weather"); - assert_eq!(input, r#"{"city":"SF"}"#); -} - -/// TS: a `length` finish reason maps to `FinishReasonUnified::Length`. -#[tokio::test] -async fn should_map_length_finish_reason() { - let server = MockServer::start().await; - mock_json_response( - &server, - "/deployments/gpt-4o/chat/completions", - json!({ - "id": "chatcmpl-len", - "object": "chat.completion", - "created": 1711115037, - "model": "gpt-4o", - "choices": [{ - "index": 0, - "message": { "role": "assistant", "content": "truncated" }, - "finish_reason": "length" - }], - "usage": { "prompt_tokens": 4, "total_tokens": 34, "completion_tokens": 30 } - }), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - assert_eq!(result.finish_reason.unified, FinishReasonUnified::Length); - assert_eq!(result.finish_reason.raw.as_deref(), Some("length")); -} - -/// TS: a 429 response maps to `AiMuxError::ApiCall` (429 in `status_code`). -#[tokio::test] -async fn should_return_rate_limited_on_429() { - let server = MockServer::start().await; - mock_json_error( - &server, - "/deployments/gpt-4o/chat/completions", - 429, - json!({ "error": { "message": "Too many requests", "type": "rate_limit" } }), - ) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.deployment("gpt-4o"); - - let err = model - .do_generate(&default_options(test_prompt())) - .await - .expect_err("429 should error"); - assert!(matches!(err, ref e if e.status_code() == Some(429))); -} diff --git a/aimux-providers/tests/azure_responses_test.rs b/aimux-providers/tests/azure_responses_test.rs deleted file mode 100644 index 69eb7971c..000000000 --- a/aimux-providers/tests/azure_responses_test.rs +++ /dev/null @@ -1,981 +0,0 @@ -//! Wiremock tests for the Azure OpenAI Responses API provider. -//! -//! These cover the Azure-specific differences from the OpenAI Responses -//! provider: -//! - deployment-based URL construction with an `api-version` query parameter -//! - v1 URL form on a non-Azure gateway (no `api-version`) -//! - `api-key` header authentication (API key) -//! - `Authorization: Bearer ` authentication (Azure AD token provider) -//! - token provider invoked per request (tokens differ per call) -//! - custom provider/request headers + user-agent suffix -//! - Azure `assistant-` file ID prefix passthrough (file_id vs base64) -//! - doGenerate text extraction, usage, response metadata, provider metadata -//! - doStream text content streaming -//! - error status codes -//! - env-var config (serial) -//! -//! Each test spins up a `wiremock` mock server, points an `AzureProvider` at -//! it via `base_url`, and asserts on the issued request and/or the parsed -//! result. -//! -//! Translated from the TS test suite -//! `reference/ai/packages/azure/src/azure-openai-provider.test.ts` -//! (responses-specific cases). Not all 44 TS cases are translated — the focus -//! is on Azure-specific capabilities. - -use aimux_core::tool::RawToolCall; -use std::sync::Arc; -use std::sync::atomic::{AtomicU32, Ordering}; - -use async_trait::async_trait; -use futures::StreamExt; -use serde_json::{Value, json}; -use wiremock::matchers::{header, method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::content::ContentPart; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; -use aimux_core::options::CallOptions; -use aimux_core::result::GenerateContent; -use aimux_core::stream_part::StreamPart; -use aimux_core::types::{FinishReasonUnified, ResponseMetadata}; - -use aimux_providers::{AzureConfig, AzureProvider, TokenProvider}; - -// ── helpers ───────────────────────────────────────────────────────────────── - -/// The TS `TEST_PROMPT`: a single user text message "Hello". -fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] -} - -/// Build `CallOptions` with everything unset except `prompt`. -fn default_options(prompt: LanguageModelPrompt) -> CallOptions { - CallOptions::new(prompt) -} - -/// The default deployment-based path used by most tests. -const DEPLOYMENT_PATH: &str = "/deployments/test-deployment/responses"; - -/// Standard mock for a JSON responses-api response on the deployment path. -async fn mock_json_response(server: &MockServer, body: Value) { - Mock::given(method("POST")) - .and(path(DEPLOYMENT_PATH)) - .respond_with(ResponseTemplate::new(200).set_body_json(body)) - .mount(server) - .await; -} - -/// Standard mock for an SSE streaming response on the deployment path. -async fn mock_sse_response(server: &MockServer, sse_body: &str) { - Mock::given(method("POST")) - .and(path(DEPLOYMENT_PATH)) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body.to_string()), - ) - .mount(server) - .await; -} - -/// A simple text responses-api response body (single message output item). -fn text_response_body() -> Value { - json!({ - "id": "resp_67c97c0203188190a025beb4a75242bc", - "object": "response", - "created_at": 1741257730, - "status": "completed", - "error": null, - "incomplete_details": null, - "model": "gpt-4o", - "output": [ - { - "id": "msg_67c97c02656c81908e080dfdf4a03cd1", - "type": "message", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "answer text", - "annotations": [] - } - ] - } - ], - "usage": { - "input_tokens": 345, - "input_tokens_details": { - "cached_tokens": 234 - }, - "output_tokens": 538, - "output_tokens_details": { - "reasoning_tokens": 123 - } - }, - "reasoning": { "effort": null, "summary": null, "context": "current_turn" } - }) -} - -/// A responses-api response body with a function_call output item. -fn tool_call_response_body() -> Value { - json!({ - "id": "resp_tool_call_123", - "object": "response", - "created_at": 1741257730, - "status": "completed", - "error": null, - "incomplete_details": null, - "model": "gpt-4o", - "output": [ - { - "id": "fc_001", - "type": "function_call", - "status": "completed", - "call_id": "call_abc123", - "name": "getWeather", - "arguments": "{\"location\": \"San Francisco\"}" - } - ], - "usage": { - "input_tokens": 50, - "output_tokens": 20, - "total_tokens": 70 - } - }) -} - -/// Concatenate SSE events and append `[DONE]`. -fn sse_body(events: &[&str]) -> String { - let mut body = String::new(); - for event in events { - body.push_str(event); - } - body.push_str("data: [DONE]\n\n"); - body -} - -/// Build an SSE event string from a JSON string. -fn sse_event(json_str: &str) -> String { - format!("data: {json_str}\n\n") -} - -/// Collect all `StreamPart`s from a `StreamResult` into a `Vec`. -async fn collect_stream(result: aimux_core::result::StreamResult) -> Vec { - let mut parts = Vec::new(); - let mut stream = result.stream; - while let Some(part) = stream.next().await { - match part { - Ok(p) => parts.push(p), - Err(e) => panic!("stream error: {e:?}"), - } - } - parts -} - -/// Extract text deltas from a list of stream parts. -fn text_deltas(parts: &[StreamPart]) -> Vec { - parts - .iter() - .filter_map(|p| match p { - StreamPart::TextDelta { delta, .. } => Some(delta.clone()), - _ => None, - }) - .collect() -} - -/// Extract the first recorded request body as a JSON value. -async fn first_request_body(server: &MockServer) -> Value { - let requests = server - .received_requests() - .await - .expect("no requests received"); - serde_json::from_slice(&requests[0].body).expect("invalid JSON body") -} - -/// A token provider that returns a fixed token. -struct StaticToken(String); -#[async_trait] -impl TokenProvider for StaticToken { - async fn get_token(&self) -> Result { - Ok(self.0.clone()) - } -} - -/// A token provider that returns a different token on each call -/// (`token-1`, `token-2`, …) so tests can assert it is invoked per request. -struct CountingToken { - count: AtomicU32, -} -#[async_trait] -impl TokenProvider for CountingToken { - async fn get_token(&self) -> Result { - let n = self.count.fetch_add(1, Ordering::SeqCst) + 1; - Ok(format!("token-{n}")) - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// URL construction & api-version -// ════════════════════════════════════════════════════════════════════════════ - -/// The default deployment-based URL form places the deployment in the path and -/// appends `?api-version=2024-10-21` (the Rust AzureConfig default). -#[tokio::test] -async fn should_build_deployment_url_with_default_api_version() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - assert_eq!(requests[0].url.path(), DEPLOYMENT_PATH); - assert_eq!(requests[0].url.query(), Some("api-version=2024-10-21")); -} - -/// A custom `api_version` is reflected in the `api-version` query parameter. -#[tokio::test] -async fn should_use_custom_api_version() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key") - .with_api_version("2025-04-01-preview"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.unwrap(); - assert_eq!( - requests[0].url.query(), - Some("api-version=2025-04-01-preview") - ); -} - -/// The v1 form on a non-Azure gateway baseURL omits `api-version` -/// (the gateway owns its own versioning). -#[tokio::test] -async fn should_omit_api_version_on_gateway_v1_url() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/responses")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_response_body())) - .mount(&server) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key") - .use_v1_urls(); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.unwrap(); - assert_eq!(requests[0].url.path(), "/v1/responses"); - // Non-Azure gateway owns its own versioning — no api-version query param. - assert_eq!(requests[0].url.query(), None); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Authentication -// ════════════════════════════════════════════════════════════════════════════ - -/// API-key auth sends the `api-key` header (not a Bearer token). -#[tokio::test] -async fn should_send_api_key_header() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.unwrap(); - let headers = &requests[0].headers; - assert_eq!( - headers.get("api-key").and_then(|v| v.to_str().ok()), - Some("test-api-key") - ); - // No Authorization header when using api-key auth. - assert!(headers.get("authorization").is_none()); -} - -/// A token provider supplies a Bearer token via the `Authorization` header. -#[tokio::test] -async fn should_send_bearer_token_from_token_provider() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_token_provider(Arc::new(StaticToken("test-azure-ad-token".to_string()))); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - let requests = server.received_requests().await.unwrap(); - let headers = &requests[0].headers; - assert_eq!( - headers.get("authorization").and_then(|v| v.to_str().ok()), - Some("Bearer test-azure-ad-token") - ); - // No api-key header when using token auth. - assert!(headers.get("api-key").is_none()); -} - -/// The token provider is invoked on every request (tokens differ per call). -#[tokio::test] -async fn should_call_token_provider_per_request() { - let server = MockServer::start().await; - // Mount two responses (one per request). - for _ in 0..2 { - mock_json_response(&server, text_response_body()).await; - } - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_token_provider(Arc::new(CountingToken { - count: AtomicU32::new(0), - })); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("first do_generate"); - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("second do_generate"); - - let requests = server.received_requests().await.unwrap(); - assert_eq!( - requests[0] - .headers - .get("authorization") - .and_then(|v| v.to_str().ok()), - Some("Bearer token-1") - ); - assert_eq!( - requests[1] - .headers - .get("authorization") - .and_then(|v| v.to_str().ok()), - Some("Bearer token-2") - ); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Headers -// ════════════════════════════════════════════════════════════════════════════ - -/// Provider-level and per-request headers are merged; user-agent suffix is set. -#[tokio::test] -async fn should_pass_custom_headers() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path(DEPLOYMENT_PATH)) - .and(header("api-key", "test-api-key")) - .and(header("custom-provider-header", "provider-header-value")) - .and(header("custom-request-header", "request-header-value")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_response_body())) - .mount(&server) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key") - .with_header("Custom-Provider-Header", "provider-header-value"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let options = CallOptions { - headers: Some({ - let mut m = std::collections::HashMap::new(); - m.insert( - "Custom-Request-Header".to_string(), - "request-header-value".to_string(), - ); - m - }), - ..default_options(test_prompt()) - }; - - let _ = model - .do_generate(&options) - .await - .expect("do_generate should succeed"); - - // If the mock matched, the headers were all present. - let requests = server.received_requests().await.unwrap(); - assert_eq!(requests.len(), 1); - let ua = requests[0] - .headers - .get("user-agent") - .and_then(|v| v.to_str().ok()) - .unwrap_or(""); - assert!( - ua.contains("ai-sdk/azure"), - "user-agent should contain ai-sdk/azure, got: {ua}" - ); -} - -// ════════════════════════════════════════════════════════════════════════════ -// doGenerate — response parsing -// ════════════════════════════════════════════════════════════════════════════ - -/// Text content is extracted from the `output` array. -#[tokio::test] -async fn should_extract_text_content() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - let text: Vec<&str> = result - .content - .iter() - .filter_map(|c| match c { - GenerateContent::Text { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect(); - assert_eq!(text, vec!["answer text"]); -} - -/// Function-call output items are extracted as ToolCall content. -#[tokio::test] -async fn should_extract_tool_call_content() { - let server = MockServer::start().await; - mock_json_response(&server, tool_call_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - let tool_calls: Vec<_> = result - .content - .iter() - .filter_map(|c| match c { - GenerateContent::ToolCall(RawToolCall { - tool_call_id, - tool_name, - input, - .. - }) => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), - _ => None, - }) - .collect(); - assert_eq!(tool_calls.len(), 1); - assert_eq!(tool_calls[0].0, "call_abc123"); - assert_eq!(tool_calls[0].1, "getWeather"); - assert_eq!(tool_calls[0].2, r#"{"location": "San Francisco"}"#); -} - -/// Usage is extracted from the `usage` field. -#[tokio::test] -async fn should_extract_usage() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - assert_eq!(result.usage.input_tokens.total, Some(345)); - assert_eq!(result.usage.output_tokens.total, Some(538)); -} - -/// Provider metadata uses the `azure` namespace key. -#[tokio::test] -async fn should_use_azure_provider_metadata_namespace() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - let pm = result.provider_metadata.expect("provider_metadata"); - assert!(pm.contains_key("azure"), "should have 'azure' key"); - assert_eq!( - pm["azure"]["responseId"], - "resp_67c97c0203188190a025beb4a75242bc" - ); - // The provider() method returns "azure.responses". - assert_eq!(model.provider(), "azure.responses"); -} - -/// Finish reason is `stop` for a completed text response. -#[tokio::test] -async fn should_map_finish_reason_stop() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); -} - -/// Finish reason is `tool-calls` when the output contains a function_call. -#[tokio::test] -async fn should_map_finish_reason_tool_calls() { - let server = MockServer::start().await; - mock_json_response(&server, tool_call_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Azure file ID prefix (assistant-) -// ════════════════════════════════════════════════════════════════════════════ - -/// When a FileBase64 part's data starts with `assistant-`, it is passed -/// through as a `file_id` instead of being base64-encoded. -#[tokio::test] -async fn should_pass_through_assistant_file_id_for_image() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ - ContentPart::text("Analyze this image"), - ContentPart::FileBase64 { - data: "assistant-abc123".to_string(), - media_type: "image/jpeg".to_string(), - filename: None, - provider_options: None, - }, - ], - ..Default::default() - }]; - - let _ = model - .do_generate(&default_options(prompt)) - .await - .expect("should succeed"); - - let body = first_request_body(&server).await; - assert_eq!(body["input"][0]["content"][1]["type"], "input_image"); - assert_eq!( - body["input"][0]["content"][1]["file_id"], - "assistant-abc123" - ); - // image_url should NOT be present. - assert!(body["input"][0]["content"][1].get("image_url").is_none()); -} - -/// When a FileBase64 part's data starts with `assistant-` for a PDF, -/// it is passed through as a `file_id` with `input_file` type. -#[tokio::test] -async fn should_pass_through_assistant_file_id_for_pdf() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ - ContentPart::text("Analyze this PDF"), - ContentPart::FileBase64 { - data: "assistant-pdf123".to_string(), - media_type: "application/pdf".to_string(), - filename: None, - provider_options: None, - }, - ], - ..Default::default() - }]; - - let _ = model - .do_generate(&default_options(prompt)) - .await - .expect("should succeed"); - - let body = first_request_body(&server).await; - assert_eq!(body["input"][0]["content"][1]["type"], "input_file"); - assert_eq!( - body["input"][0]["content"][1]["file_id"], - "assistant-pdf123" - ); - // file_data and filename should NOT be present. - assert!(body["input"][0]["content"][1].get("file_data").is_none()); - assert!(body["input"][0]["content"][1].get("filename").is_none()); -} - -/// Non-`assistant-` file IDs fall back to base64 encoding (the default -/// OpenAI behavior). The Rust convert produces `input_file` for `FileBase64` -/// regardless of media type — the data is not treated as a file ID. -#[tokio::test] -async fn should_fall_back_to_base64_for_non_assistant_file_ids() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ - ContentPart::text("Analyze this image"), - ContentPart::FileBase64 { - data: "file-abc123".to_string(), - media_type: "image/jpeg".to_string(), - filename: None, - provider_options: None, - }, - ], - ..Default::default() - }]; - - let _ = model - .do_generate(&default_options(prompt)) - .await - .expect("should succeed"); - - let body = first_request_body(&server).await; - // Rust convert produces input_file for FileBase64 (not input_image). - assert_eq!(body["input"][0]["content"][1]["type"], "input_file"); - assert_eq!( - body["input"][0]["content"][1]["file_data"], - "data:image/jpeg;base64,file-abc123" - ); - // file_id should NOT be present (data doesn't start with assistant-). - assert!(body["input"][0]["content"][1].get("file_id").is_none()); -} - -// ════════════════════════════════════════════════════════════════════════════ -// doStream — text content -// ════════════════════════════════════════════════════════════════════════════ - -/// Streaming text content: response.created → output_item.added → -/// output_text.delta → output_item.done → response.completed. -#[tokio::test] -async fn should_stream_text_content() { - let server = MockServer::start().await; - let sse = sse_body(&[ - &sse_event( - r#"{"type":"response.created","response":{"id":"resp_123","object":"response","created_at":1741257730,"status":"in_progress","model":"gpt-4o","output":[]}}"#, - ), - &sse_event( - r#"{"type":"response.output_item.added","output_index":0,"item":{"id":"msg_1","type":"message","status":"in_progress","role":"assistant","content":[]}}"#, - ), - &sse_event( - r#"{"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"Hello"}"#, - ), - &sse_event( - r#"{"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":" world"}"#, - ), - &sse_event( - r#"{"type":"response.output_item.done","output_index":0,"item":{"id":"msg_1","type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"Hello world","annotations":[]}]}}"#, - ), - &sse_event( - r#"{"type":"response.completed","response":{"id":"resp_123","object":"response","created_at":1741257730,"status":"completed","model":"gpt-4o","output":[{"id":"msg_1","type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"Hello world","annotations":[]}]}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}"#, - ), - ]); - mock_sse_response(&server, &sse).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let result = model - .do_stream(&default_options(test_prompt())) - .await - .expect("do_stream should succeed"); - - let parts = collect_stream(result).await; - - // StreamStart - assert!(matches!(&parts[0], StreamPart::StreamStart { .. })); - - // ResponseMetadata - assert!(matches!( - &parts[1], - StreamPart::ResponseMetadata(ResponseMetadata { id, .. }) if id.as_deref() == Some("resp_123") - )); - - // TextStart - assert!(matches!( - &parts[2], - StreamPart::TextStart { id, .. } if id == "msg_1" - )); - - // TextDelta "Hello" + " world" - let deltas = text_deltas(&parts); - assert_eq!(deltas, vec!["Hello", " world"]); - - // TextEnd - assert!(matches!( - parts.iter().find(|p| matches!(p, StreamPart::TextEnd { .. })), - Some(StreamPart::TextEnd { id, .. }) if id == "msg_1" - )); - - // Finish - let finish = parts - .iter() - .find(|p| matches!(p, StreamPart::Finish { .. })); - assert!(finish.is_some()); - if let Some(StreamPart::Finish { - finish_reason, - usage, - .. - }) = finish - { - assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); - assert_eq!(usage.input_tokens.total, Some(10)); - assert_eq!(usage.output_tokens.total, Some(5)); - } else { - panic!("expected Finish part"); - } -} - -/// Streaming: api-key header is sent and api-version query param is present. -#[tokio::test] -async fn should_send_api_key_header_on_stream() { - let server = MockServer::start().await; - let sse = sse_body(&[ - &sse_event( - r#"{"type":"response.created","response":{"id":"resp_1","object":"response","created_at":1,"status":"in_progress","model":"gpt-4o","output":[]}}"#, - ), - &sse_event( - r#"{"type":"response.completed","response":{"id":"resp_1","object":"response","created_at":1,"status":"completed","model":"gpt-4o","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}"#, - ), - ]); - mock_sse_response(&server, &sse).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let result = model - .do_stream(&default_options(test_prompt())) - .await - .expect("do_stream should succeed"); - let _ = collect_stream(result).await; - - let requests = server.received_requests().await.unwrap(); - assert_eq!(requests.len(), 1); - assert_eq!(requests[0].url.query(), Some("api-version=2024-10-21")); - assert_eq!( - requests[0] - .headers - .get("api-key") - .and_then(|v| v.to_str().ok()), - Some("test-api-key") - ); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Error handling -// ════════════════════════════════════════════════════════════════════════════ - -/// A non-2xx response is mapped to a provider error. -#[tokio::test] -async fn should_handle_error_status() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path(DEPLOYMENT_PATH)) - .respond_with(ResponseTemplate::new(429).set_body_json(json!({ - "error": { - "message": "Rate limit exceeded", - "type": "rate_limit_exceeded", - "code": "rate_limit_exceeded" - } - }))) - .mount(&server) - .await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let result = model.do_generate(&default_options(test_prompt())).await; - - assert!(result.is_err()); - let err = result.unwrap_err(); - let msg = err.to_string().to_lowercase(); - assert!( - msg.contains("rate limit"), - "error should mention rate limit, got: {err}" - ); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Environment variable configuration (serial) -// ════════════════════════════════════════════════════════════════════════════ - -/// `AzureConfig::from_env` reads `AZURE_API_KEY` and `AZURE_RESOURCE_NAME`. -#[tokio::test] -#[serial_test::serial] -async fn should_create_config_from_env() { - // Set env vars. - unsafe { - std::env::set_var("AZURE_API_KEY", "env-api-key"); - std::env::set_var("AZURE_RESOURCE_NAME", "env-resource"); - } - - let config = AzureConfig::from_env().expect("from_env"); - assert_eq!(config.resource_name.as_deref(), Some("env-resource")); - - // Clean up. - unsafe { - std::env::remove_var("AZURE_API_KEY"); - std::env::remove_var("AZURE_RESOURCE_NAME"); - } - - // Verify the config produces the right URL. - let model = AzureProvider::new(config) - .expect("provider") - .responses_model("my-deployment"); - assert_eq!(model.model_id(), "my-deployment"); - assert_eq!(model.provider(), "azure.responses"); -} - -/// `AzureConfig::from_env` fails when `AZURE_API_KEY` is not set. -#[tokio::test] -#[serial_test::serial] -async fn should_fail_from_env_without_api_key() { - // Ensure the env var is not set. - unsafe { - std::env::remove_var("AZURE_API_KEY"); - } - - let result = AzureConfig::from_env(); - assert!(result.is_err()); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Request body construction -// ════════════════════════════════════════════════════════════════════════════ - -/// The request body includes the model (deployment) and input array. -#[tokio::test] -async fn should_send_model_and_input_in_request_body() { - let server = MockServer::start().await; - mock_json_response(&server, text_response_body()).await; - - let config = AzureConfig::new() - .with_base_url(server.uri()) - .with_api_key("test-api-key"); - let provider = AzureProvider::new(config).expect("provider"); - let model = provider.responses_model("test-deployment"); - - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - let body = first_request_body(&server).await; - assert_eq!(body["model"], "test-deployment"); - assert_eq!(body["input"][0]["role"], "user"); - assert_eq!(body["input"][0]["content"][0]["type"], "input_text"); - assert_eq!(body["input"][0]["content"][0]["text"], "Hello"); -} diff --git a/aimux-providers/tests/baseten_test.rs b/aimux-providers/tests/baseten_test.rs index a1bb20359..7e902bf24 100644 --- a/aimux-providers/tests/baseten_test.rs +++ b/aimux-providers/tests/baseten_test.rs @@ -26,20 +26,14 @@ use serial_test::serial; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_providers::{ProviderOptions, provider, provider_from_env}; fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn default_options(prompt: LanguageModelPrompt) -> CallOptions { diff --git a/aimux-providers/tests/bedrock_convert_test.rs b/aimux-providers/tests/bedrock_convert_test.rs index f236d552c..76ffbce9d 100644 --- a/aimux-providers/tests/bedrock_convert_test.rs +++ b/aimux-providers/tests/bedrock_convert_test.rs @@ -7,13 +7,8 @@ //! //! Cases the Rust data model cannot express are skipped with an inline //! comment. The main categories of skips: -//! - **Message-level `providerOptions`**: `LanguageModelPromptMessage` has no -//! `provider_options` field, so message-level cache points (system / user / -//! assistant message `bedrock.cachePoint`) cannot be expressed. Part-level -//! cache points (on `ContentPart::Text` `provider_options`) ARE covered. -//! - **System-after-non-system throw / unsupported-mime throws**: the Rust -//! `convert_prompt_to_bedrock` returns `(system, messages)` (no `Result`), -//! so it cannot surface `UnsupportedFunctionalityError`. +//! - **System-after-non-system throw**: the converter currently lifts +//! all system messages into the system array. //! - **S3 URLs / provider references**: `FileUrl` / `FileReference` are not //! converted by the Rust Bedrock path. //! - **Top-level-only mediaType auto-detection from bytes**: the Rust path @@ -21,19 +16,18 @@ //! - **Mistral tool-call-id normalization (`isMistral`)**: the Rust //! `convert_prompt_to_bedrock` has no `isMistral` parameter. The //! non-Mistral (passthrough) cases ARE covered. -//! - **Redacted reasoning / foreign-provider reasoning**: `ContentPart::Reasoning` -//! carries only `signature` (no `redactedData`, no provider distinction). //! - **Provider-defined tools (web_search, anthropic provider tools)** and //! `additionalTools`/`betas`: the Rust `FunctionTool` has no `type`/`id`. //! - **`raw` echo on `Usage`**: the Rust `Usage` type has no `raw` field. use serde_json::{Value, json}; -use aimux_core::shared::{SharedProviderOptions, provider_namespace}; +use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions, provider_namespace}; -use aimux_core::content::ContentPart; -use aimux_core::language_model_message::LanguageModelPromptMessage; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, FilePart, LanguageModelMessage, ReasoningPart, TextPart, ToolCallPart, ToolPart, + ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, +}; use aimux_core::options::ToolChoice; use aimux_core::tool::FunctionTool; @@ -43,54 +37,72 @@ use aimux_providers::bedrock::convert::{ // ── prompt builders ───────────────────────────────────────────────────────── -fn msg(role: Role, content: Vec) -> LanguageModelPromptMessage { - LanguageModelPromptMessage { - role, +fn user(content: Vec) -> LanguageModelMessage { + LanguageModelMessage::User { + content, + provider_options: None, + } +} + +fn assistant(content: Vec) -> LanguageModelMessage { + LanguageModelMessage::Assistant { content, provider_options: None, } } -fn user(content: Vec) -> LanguageModelPromptMessage { - msg(Role::User, content) +fn system_msg(text: &str) -> LanguageModelMessage { + LanguageModelMessage::System { + content: text.to_string(), + provider_options: None, + } } -fn assistant(content: Vec) -> LanguageModelPromptMessage { - msg(Role::Assistant, content) +fn tool_msg(content: Vec) -> LanguageModelMessage { + LanguageModelMessage::Tool { + content, + provider_options: None, + } } -fn system_msg(text: &str) -> LanguageModelPromptMessage { - msg(Role::System, vec![ContentPart::text(text)]) +fn user_text(text: &str) -> UserPart { + UserPart::Text(TextPart { + text: text.to_string(), + provider_options: None, + }) } -fn tool_msg(content: Vec) -> LanguageModelPromptMessage { - msg(Role::Tool, content) +fn assistant_text(text: &str) -> AssistantPart { + AssistantPart::Text(TextPart { + text: text.to_string(), + provider_options: None, + }) } -/// A `file` part whose inline data is an already-base64 string — mirrors the -/// TS `file` part with `data: { type: 'data', data: '' }` (the Rust -/// `ContentPart::FileBase64` holds the base64 string verbatim). +/// A file part with inline base64 data, passed through verbatim. fn file_base64( data: &str, media_type: &str, filename: Option<&str>, provider_options: Option, -) -> ContentPart { - ContentPart::FileBase64 { - data: data.to_string(), +) -> UserPart { + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64(data.to_string()), + }, media_type: media_type.to_string(), filename: filename.map(std::string::ToString::to_string), provider_options, - } + }) } -fn text_with_cache(text: &str, cache_type: &str, ttl: Option<&str>) -> ContentPart { +fn text_with_cache(text: &str, cache_type: &str, ttl: Option<&str>) -> TextPart { let mut cp = serde_json::Map::new(); cp.insert("type".to_string(), json!(cache_type)); if let Some(t) = ttl { cp.insert("ttl".to_string(), json!(t)); } - ContentPart::Text { + TextPart { text: text.to_string(), provider_options: Some( provider_namespace("bedrock", json!({ "cachePoint": Value::Object(cp) })).unwrap(), @@ -98,20 +110,32 @@ fn text_with_cache(text: &str, cache_type: &str, ttl: Option<&str>) -> ContentPa } } -fn reasoning(text: &str, signature: Option<&str>) -> ContentPart { - ContentPart::Reasoning { +fn reasoning(text: &str, signature: Option<&str>) -> AssistantPart { + AssistantPart::Reasoning(ReasoningPart { text: text.to_string(), - signature: signature.map(std::string::ToString::to_string), - provider_options: None, - } + provider_options: signature.map(|signature| { + provider_namespace("amazonBedrock", json!({ "signature": signature })).unwrap() + }), + }) } -fn tool_call(id: &str, name: &str, input: Value) -> ContentPart { - ContentPart::tool_call(id.to_string(), name.to_string(), input) +fn tool_call(id: &str, name: &str, input: Value) -> AssistantPart { + AssistantPart::ToolCall(ToolCallPart { + tool_call_id: id.to_string(), + tool_name: name.to_string(), + input, + provider_executed: None, + provider_options: None, + }) } -fn tool_result(id: &str, output: Value) -> ContentPart { - ContentPart::tool_result(id.to_string(), output) +fn tool_result(id: &str, output: ToolResultOutput) -> ToolPart { + ToolPart::ToolResult(ToolResultPart { + tool_call_id: id.to_string(), + tool_name: "test".to_string(), + output, + provider_options: None, + }) } // ════════════════════════════════════════════════════════════════════════════ @@ -123,7 +147,8 @@ fn tool_result(id: &str, output: Value) -> ContentPart { /// TS: "should combine multiple leading system messages into a single system message" #[test] fn system_combine_multiple_leading() { - let (system, _) = convert_prompt_to_bedrock(&vec![system_msg("Hello"), system_msg("World")]); + let (system, _) = + convert_prompt_to_bedrock(&vec![system_msg("Hello"), system_msg("World")]).unwrap(); assert_eq!( Value::Array(system), json!([{ "text": "Hello" }, { "text": "World" }]) @@ -131,16 +156,13 @@ fn system_combine_multiple_leading() { } // SKIPPED (TS: "should throw an error if a system message is provided after a -// non-system message"): convert_prompt_to_bedrock returns (system, messages) -// with no Result, so it cannot surface UnsupportedFunctionalityError. - -// SKIPPED (TS: system message cache point, 5m, 1h — 3 cases): message-level -// providerOptions are not modelled on LanguageModelPromptMessage. +// non-system message"): the converter currently lifts all system messages +// into the system array. /// TS: "should extract the system message" #[test] fn system_extract_single() { - let (system, _) = convert_prompt_to_bedrock(&vec![system_msg("Hello")]); + let (system, _) = convert_prompt_to_bedrock(&vec![system_msg("Hello")]).unwrap(); assert_eq!(Value::Array(system), json!([{ "text": "Hello" }])); } @@ -150,9 +172,10 @@ fn system_extract_single() { #[test] fn user_convert_image_parts() { let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![ - ContentPart::text("Hello"), + user_text("Hello"), file_base64("AAECAw==", "image/png", None, None), - ])]); + ])]) + .unwrap(); assert_eq!( Value::Array(messages), json!([{ @@ -172,9 +195,10 @@ fn user_convert_image_parts() { #[test] fn user_convert_document_parts() { let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![ - ContentPart::text("Hello"), + user_text("Hello"), file_base64("AAECAw==", "application/pdf", None, None), - ])]); + ])]) + .unwrap(); assert_eq!( messages[0]["content"], json!([ @@ -188,14 +212,15 @@ fn user_convert_document_parts() { #[test] fn user_strip_file_extension() { let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![ - ContentPart::text("Hello"), + user_text("Hello"), file_base64( "AAECAw==", "application/pdf", Some("custom-filename.pdf"), None, ), - ])]); + ])]) + .unwrap(); assert_eq!( messages[0]["content"][1], json!({ "document": { "format": "pdf", "name": "custom-filename", "source": { "bytes": "AAECAw==" } } }) @@ -210,7 +235,8 @@ fn user_preserve_filename_without_extension() { "application/pdf", Some("custom-filename"), None, - )])]); + )])]) + .unwrap(); assert_eq!( messages[0]["content"][0], json!({ "document": { "format": "pdf", "name": "custom-filename", "source": { "bytes": "AAECAw==" } } }) @@ -225,9 +251,10 @@ fn user_consistent_document_names() { file_base64("AAECAw==", "application/pdf", None, None), file_base64("BAUGBw==", "application/pdf", None, None), ]), - assistant(vec![ContentPart::text("OK")]), + assistant(vec![assistant_text("OK")]), user(vec![file_base64("AAECAw==", "application/pdf", None, None)]), - ]); + ]) + .unwrap(); assert_eq!( Value::Array(messages), json!([ @@ -243,20 +270,20 @@ fn user_consistent_document_names() { ); } -// SKIPPED (TS: user message cache point, 5m, 1h — 3 cases): message-level -// providerOptions are not modelled on LanguageModelPromptMessage. +// Message-level cache-point cases remain outside this restored test subset. // SKIPPED (TS: "should throw for file parts with provider references"): -// FileReference is not converted (no Result to surface the throw). +// FileReference is not converted by the user-file path. /// TS: "should add cache point to user content part when specified" #[test] fn user_content_part_cache_point() { let (system, messages) = convert_prompt_to_bedrock(&vec![user(vec![ - ContentPart::text("Hello"), - text_with_cache("cached", "default", Some("5m")), - ContentPart::text("World"), - ])]); + user_text("Hello"), + UserPart::Text(text_with_cache("cached", "default", Some("5m"))), + user_text("World"), + ])]) + .unwrap(); assert!(system.is_empty()); assert_eq!( Value::Array(messages), @@ -278,9 +305,10 @@ fn user_content_part_cache_point() { #[test] fn assistant_trim_trailing_whitespace_last() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("user content")]), - assistant(vec![ContentPart::text("assistant content ")]), - ]); + user(vec![user_text("user content")]), + assistant(vec![assistant_text("assistant content ")]), + ]) + .unwrap(); assert_eq!( Value::Array(messages), json!([ @@ -294,12 +322,13 @@ fn assistant_trim_trailing_whitespace_last() { #[test] fn assistant_trim_trailing_whitespace_multi_part() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("user content")]), + user(vec![user_text("user content")]), assistant(vec![ - ContentPart::text("assistant "), - ContentPart::text("content "), + assistant_text("assistant "), + assistant_text("content "), ]), - ]); + ]) + .unwrap(); assert_eq!( messages[1]["content"], json!([{ "text": "assistant " }, { "text": "content" }]) @@ -310,10 +339,11 @@ fn assistant_trim_trailing_whitespace_multi_part() { #[test] fn assistant_keep_trailing_whitespace_with_further_user() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("user content")]), - assistant(vec![ContentPart::text("assistant content ")]), - user(vec![ContentPart::text("user content 2")]), - ]); + user(vec![user_text("user content")]), + assistant(vec![assistant_text("assistant content ")]), + user(vec![user_text("user content 2")]), + ]) + .unwrap(); assert_eq!( messages[1]["content"], json!([{ "text": "assistant content " }]) @@ -324,11 +354,12 @@ fn assistant_keep_trailing_whitespace_with_further_user() { #[test] fn assistant_combine_sequential() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("Hi!")]), - assistant(vec![ContentPart::text("Hello")]), - assistant(vec![ContentPart::text("World")]), - assistant(vec![ContentPart::text("!")]), - ]); + user(vec![user_text("Hi!")]), + assistant(vec![assistant_text("Hello")]), + assistant(vec![assistant_text("World")]), + assistant(vec![assistant_text("!")]), + ]) + .unwrap(); assert_eq!( Value::Array(messages), json!([ @@ -338,17 +369,17 @@ fn assistant_combine_sequential() { ); } -// SKIPPED (TS: assistant message cache point, 5m, 1h — 3 cases): message-level -// providerOptions are not modelled on LanguageModelPromptMessage. +// Message-level cache-point cases remain outside this restored test subset. /// TS: "should add cache point to assistant content part when specified" #[test] fn assistant_content_part_cache_point() { let (_, messages) = convert_prompt_to_bedrock(&vec![assistant(vec![ - ContentPart::text("Hello"), - text_with_cache("cached", "default", Some("1h")), - ContentPart::text("World"), - ])]); + assistant_text("Hello"), + AssistantPart::Text(text_with_cache("cached", "default", Some("1h"))), + assistant_text("World"), + ])]) + .unwrap(); assert_eq!( Value::Array(messages), json!([{ @@ -367,12 +398,13 @@ fn assistant_content_part_cache_point() { #[test] fn assistant_reasoning_with_signature() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("Explain your reasoning")]), + user(vec![user_text("Explain your reasoning")]), assistant(vec![reasoning( "This is my step-by-step reasoning process", Some("test-signature"), )]), - ]); + ]) + .unwrap(); assert_eq!( messages[1]["content"], json!([{ @@ -386,26 +418,20 @@ fn assistant_reasoning_with_signature() { ); } -// SKIPPED (TS: "should properly convert redacted-reasoning content type"): -// ContentPart::Reasoning has no redactedData field. - -// SKIPPED (TS: "should omit assistant message reasoning parts signed by a -// foreign provider"): ContentPart::Reasoning.signature carries no provider -// distinction, so foreign-provider (anthropic) signatures cannot be detected. - /// TS: "should preserve assistant message reasoning parts with amazonBedrock providerOptions" #[test] fn assistant_preserve_amazon_bedrock_reasoning() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("Explain your reasoning")]), + user(vec![user_text("Explain your reasoning")]), assistant(vec![ reasoning( "Bedrock-signed reasoning round-tripped to Bedrock", Some("bedrock-signature"), ), - ContentPart::text("final answer"), + assistant_text("final answer"), ]), - ]); + ]) + .unwrap(); assert_eq!( messages[1]["content"], json!([ @@ -422,12 +448,13 @@ fn assistant_preserve_amazon_bedrock_reasoning() { #[test] fn assistant_no_trim_reasoning_with_signature() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("Explain your reasoning")]), + user(vec![user_text("Explain your reasoning")]), assistant(vec![reasoning( "This is my reasoning with trailing space ", Some("test-signature"), )]), - ]); + ]) + .unwrap(); assert_eq!( messages[1]["content"][0]["reasoningContent"]["reasoningText"]["text"], json!("This is my reasoning with trailing space ") @@ -438,12 +465,13 @@ fn assistant_no_trim_reasoning_with_signature() { #[test] fn assistant_omit_reasoning_without_signature() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("Explain your reasoning")]), + user(vec![user_text("Explain your reasoning")]), assistant(vec![ reasoning("This is my reasoning with trailing space ", None), - ContentPart::text("final answer"), + assistant_text("final answer"), ]), - ]); + ]) + .unwrap(); assert_eq!(messages[1]["content"], json!([{ "text": "final answer" }])); } @@ -451,13 +479,14 @@ fn assistant_omit_reasoning_without_signature() { #[test] fn assistant_omit_multiple_reasoning_without_signatures() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("Explain your reasoning")]), + user(vec![user_text("Explain your reasoning")]), assistant(vec![ reasoning("First reasoning with trailing space ", None), reasoning("Second reasoning with trailing space ", None), - ContentPart::text("final answer"), + assistant_text("final answer"), ]), - ]); + ]) + .unwrap(); assert_eq!(messages[1]["content"], json!([{ "text": "final answer" }])); } @@ -465,16 +494,20 @@ fn assistant_omit_multiple_reasoning_without_signatures() { #[test] fn assistant_omit_unsigned_reasoning_preserving_tool_calls() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("What is the weather?")]), + user(vec![user_text("What is the weather?")]), assistant(vec![ reasoning("I should call the weather tool.", None), tool_call("call-1", "getWeather", json!({ "city": "SF" })), ]), tool_msg(vec![tool_result( "call-1", - json!({ "type": "text", "value": "Sunny, 72F" }), + ToolResultOutput::Text { + value: "Sunny, 72F".to_string(), + provider_options: None, + }, )]), - ]); + ]) + .unwrap(); assert_eq!( Value::Array(messages), json!([ @@ -493,20 +526,24 @@ fn assistant_omit_unsigned_reasoning_preserving_tool_calls() { #[test] fn assistant_preserve_reasoning_multi_turn() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("What is the weather?")]), + user(vec![user_text("What is the weather?")]), assistant(vec![ reasoning("Let me check the weather API.\n", Some("sig-abc123")), tool_call("call-1", "getWeather", json!({ "city": "SF" })), ]), tool_msg(vec![tool_result( "call-1", - json!({ "type": "text", "value": "Sunny, 72F" }), + ToolResultOutput::Text { + value: "Sunny, 72F".to_string(), + provider_options: None, + }, )]), assistant(vec![ reasoning("The weather is sunny and warm.\n", Some("sig-def456")), - ContentPart::text("It is sunny and 72F in SF."), + assistant_text("It is sunny and 72F in SF."), ]), - ]); + ]) + .unwrap(); assert_eq!( Value::Array(messages), json!([ @@ -534,15 +571,16 @@ fn assistant_preserve_reasoning_multi_turn() { #[test] fn assistant_mix_text_and_reasoning() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("Explain your reasoning")]), + user(vec![user_text("Explain your reasoning")]), assistant(vec![ - ContentPart::text("My answer is 42."), + assistant_text("My answer is 42."), reasoning( "I calculated this by analyzing the meaning of life", Some("reasoning-process"), ), ]), - ]); + ]) + .unwrap(); assert_eq!( messages[1]["content"], json!([ @@ -559,14 +597,15 @@ fn assistant_mix_text_and_reasoning() { #[test] fn assistant_filter_empty_text_blocks() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("Hello")]), + user(vec![user_text("Hello")]), assistant(vec![ - ContentPart::text("\n\n"), + assistant_text("\n\n"), tool_call("call-123", "test", json!({})), - ContentPart::text(" "), - ContentPart::text("actual content"), + assistant_text(" "), + assistant_text("actual content"), ]), - ]); + ]) + .unwrap(); assert_eq!( messages[1]["content"], json!([ @@ -583,7 +622,8 @@ fn assistant_wrap_non_object_tool_input() { "call-1", "cityAttractions", Value::String("{ \"city\": \"San Francisco\", }".to_string()), - )])]); + )])]) + .unwrap(); assert_eq!( messages[0]["content"], json!([{ @@ -607,7 +647,8 @@ fn assistant_strip_invalid_tool_name_chars() { json!({}), ), tool_call("call-3", "$", json!({})), - ])]); + ])]) + .unwrap(); assert_eq!( messages[0]["content"], json!([ @@ -622,15 +663,16 @@ fn assistant_strip_invalid_tool_name_chars() { #[test] fn assistant_preserve_empty_text_with_reasoning() { let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ContentPart::text("Hello")]), + user(vec![user_text("Hello")]), assistant(vec![ reasoning("thinking...", Some("sig-1")), - ContentPart::text(""), + assistant_text(""), reasoning("more thinking...", Some("sig-2")), - ContentPart::text("response text"), + assistant_text("response text"), tool_call("call-123", "test", json!({})), ]), - ]); + ]) + .unwrap(); assert_eq!( messages[1]["content"], json!([ @@ -650,8 +692,14 @@ fn assistant_preserve_empty_text_with_reasoning() { fn tool_result_content_text() { let (_, messages) = convert_prompt_to_bedrock(&vec![tool_msg(vec![tool_result( "call-123", - json!({ "type": "content", "value": [{ "type": "text", "text": "The result is 42" }] }), - )])]); + ToolResultOutput::Content { + value: vec![ToolResultContent::Text(TextPart { + text: "The result is 42".to_string(), + provider_options: None, + })], + }, + )])]) + .unwrap(); assert_eq!( messages[0], json!({ @@ -671,12 +719,18 @@ fn tool_result_content_text() { fn tool_result_content_image() { let (_, messages) = convert_prompt_to_bedrock(&vec![tool_msg(vec![tool_result( "call-123", - json!({ "type": "content", "value": [{ - "type": "file", - "data": { "type": "data", "data": "base64data" }, - "mediaType": "image/jpeg" - }] }), - )])]); + ToolResultOutput::Content { + value: vec![ToolResultContent::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64("base64data".to_string()), + }, + media_type: "image/jpeg".to_string(), + filename: None, + provider_options: None, + })], + }, + )])]) + .unwrap(); assert_eq!( messages[0]["content"][0], json!({ @@ -696,13 +750,18 @@ fn tool_result_content_image() { fn tool_result_content_pdf() { let (_, messages) = convert_prompt_to_bedrock(&vec![tool_msg(vec![tool_result( "call-123", - json!({ "type": "content", "value": [{ - "type": "file", - "data": { "type": "data", "data": "base64data" }, - "mediaType": "application/pdf", - "filename": "tool-result.pdf" - }] }), - )])]); + ToolResultOutput::Content { + value: vec![ToolResultContent::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64("base64data".to_string()), + }, + media_type: "application/pdf".to_string(), + filename: Some("tool-result.pdf".to_string()), + provider_options: None, + })], + }, + )])]) + .unwrap(); assert_eq!( messages[0]["content"][0], json!({ @@ -716,18 +775,21 @@ fn tool_result_content_pdf() { ); } -// SKIPPED (TS: "should throw error for unsupported image format in tool result +// NOT PORTED (TS: "should throw error for unsupported image format in tool result // content" and "should throw error for unsupported mime type in tool result -// file content"): convert_prompt_to_bedrock returns no Result; unknown mimes -// fall back to a default format instead of throwing. +// file content"). /// TS: "should fallback to stringified result when content is undefined" (json output) #[test] fn tool_result_json_output() { let (_, messages) = convert_prompt_to_bedrock(&vec![tool_msg(vec![tool_result( "call-123", - json!({ "type": "json", "value": { "value": 42 } }), - )])]); + ToolResultOutput::Json { + value: json!({ "value": 42 }), + provider_options: None, + }, + )])]) + .unwrap(); assert_eq!( messages[0]["content"][0], json!({ @@ -749,7 +811,8 @@ fn citations_enabled_for_pdf() { "application/pdf", None, Some(provider_namespace("bedrock", json!({ "citations": { "enabled": true } })).unwrap()), - )])]); + )])]) + .unwrap(); assert_eq!( messages[0]["content"][0], json!({ @@ -771,7 +834,8 @@ fn citations_disabled_for_pdf() { "application/pdf", None, Some(provider_namespace("bedrock", json!({ "citations": { "enabled": false } })).unwrap()), - )])]); + )])]) + .unwrap(); assert_eq!( messages[0]["content"][0], json!({ @@ -793,7 +857,8 @@ fn citations_default_for_pdf() { "application/pdf", None, None, - )])]); + )])]) + .unwrap(); assert_eq!( messages[0]["content"][0], json!({ @@ -828,7 +893,8 @@ fn citations_multiple_pdfs() { .unwrap(), ), ), - ])]); + ])]) + .unwrap(); assert_eq!( messages[0]["content"], json!([ @@ -851,7 +917,8 @@ fn file_format_xlsx() { "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", None, None, - )])]); + )])]) + .unwrap(); assert!(system.is_empty()); assert_eq!( messages[0]["content"][0], @@ -867,7 +934,8 @@ fn file_format_docx() { "application/vnd.openxmlformats-officedocument.wordprocessingml.document", None, None, - )])]); + )])]) + .unwrap(); assert_eq!( messages[0]["content"][0], json!({ "document": { "format": "docx", "name": "document-1", "source": { "bytes": "base64data" } } }) @@ -886,8 +954,12 @@ fn mistral_no_normalize_when_false() { let original_id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; let (_, messages) = convert_prompt_to_bedrock(&vec![tool_msg(vec![tool_result( original_id, - json!({ "type": "text", "value": "The result is 42" }), - )])]); + ToolResultOutput::Text { + value: "The result is 42".to_string(), + provider_options: None, + }, + )])]) + .unwrap(); assert_eq!( messages[0]["content"][0]["toolResult"]["toolUseId"], json!(original_id) @@ -900,8 +972,12 @@ fn mistral_default_no_normalize() { let original_id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; let (_, messages) = convert_prompt_to_bedrock(&vec![tool_msg(vec![tool_result( original_id, - json!({ "type": "text", "value": "The result is 42" }), - )])]); + ToolResultOutput::Text { + value: "The result is 42".to_string(), + provider_options: None, + }, + )])]) + .unwrap(); assert_eq!( messages[0]["content"][0]["toolResult"]["toolUseId"], json!(original_id) @@ -918,7 +994,8 @@ fn media_type_pass_through_full_image() { "image/png", None, None, - )])]); + )])]) + .unwrap(); assert_eq!( Value::Array(messages), json!([{ @@ -940,7 +1017,8 @@ fn media_type_route_to_document_text_plain() { "text/plain", None, None, - )])]); + )])]) + .unwrap(); assert_eq!( messages[0]["content"][0], json!({ "document": { "format": "txt", "name": "document-1", "source": { "bytes": "base64data" } } }) diff --git a/aimux-providers/tests/bedrock_reasoning_test.rs b/aimux-providers/tests/bedrock_reasoning_test.rs deleted file mode 100644 index f728a846f..000000000 --- a/aimux-providers/tests/bedrock_reasoning_test.rs +++ /dev/null @@ -1,1215 +0,0 @@ -//! Bedrock reasoning / thinking tests, translated from the TypeScript suite. -//! -//! Translation sources (under `reference/ai/packages/amazon-bedrock/src/`): -//! - `amazon-bedrock-chat-language-model.test.ts` -//! - `describe('reasoning')` under `doGenerate` — "should extract reasoning -//! and text response" (fixture `amazon-bedrock-reasoning.json`). -//! - `describe('reasoning')` under `doStream` — "should stream reasoning and -//! text parts" (fixture `amazon-bedrock-reasoning.chunks.txt`). -//! - "should stream reasoning text deltas" (inline chunks). -//! - "should extract reasoning text with signature" / "…without signature". -//! - "should preserve empty text blocks between reasoning blocks". -//! - "should extract redacted reasoning". -//! - "should handle multiple reasoning blocks". -//! - "should transform reasoningConfig to thinking in stream requests". -//! - "merges user additionalModelRequestFields with derived thinking (stream)". -//! - `convert-to-amazon-bedrock-chat-messages.test.ts` — the reasoning -//! `it(...)` cases (prompt-side replay of signed/unsigned reasoning). -//! -//! Tests that require a data model the Rust crate does not expose are marked -//! `#[ignore]` with an inline reason: -//! - **Redacted reasoning replay** — `ContentPart::Reasoning` has no -//! `redactedData` field and `convert_prompt_to_bedrock` has no -//! `redactedReasoning` branch. -//! - **Foreign-provider reasoning replay** — `ContentPart::Reasoning.signature` -//! is provider-agnostic; the converter cannot distinguish an `anthropic` -//! signature (which must be dropped) from a `bedrock` one (which is kept). - -use futures::StreamExt; -use serde_json::{Value, json}; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::content::ContentPart; -use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; -use aimux_core::options::CallOptions; -use aimux_core::result::{GenerateContent, ReasoningOutput, StreamResult}; -use aimux_core::stream_part::StreamPart; -use aimux_core::types::{FinishReasonUnified, ProviderMetadata}; - -use aimux_providers::bedrock::convert::convert_prompt_to_bedrock; -use aimux_providers::bedrock::event_stream; -use aimux_providers::bedrock::{BedrockAuth, BedrockConfig, BedrockModel}; - -// ── Constants ──────────────────────────────────────────────────────────────── - -/// Model id used by the TS reasoning tests (matches the snapshot `modelId`). -const MODEL_ID: &str = "anthropic.claude-3-haiku-20240307-v1:0"; - -/// Signature on the `amazon-bedrock-reasoning.json` fixture's reasoning block. -const REASONING_SIGNATURE: &str = "EvYBCkgICxABGAIqQA9E9VC377UnbjdfXCw4RwQaaIXsqocZKzI3WwWtXT/wBjzAkBOfIiVkP/mhvsPGb5oQW/2j+E1tZcNFoXIb26YSDBTSbDrbHwXJoAl3/hoMA/Qo4rIHvObP/yBrIjAT+l69mh9k7evFO+w+ransSD5YGq4hYrieOyJH3k/zHv8nBJwiz3Nlrb3jNFi2Ib0qXJH7BIK0hR85yYIoccReoJUpeLEMxeWbsO6f0EpOSMSi3YDLd1U9NAw7Itjeay1fEBPHwvawt2M4e/rn52SsuEz3p2DfOZm+N3bnL63rg0Cb9lFcp3k5kuU1h6hcGAE="; - -/// Signature carried by the streaming reasoning fixture's final delta. -const STREAM_SIGNATURE: &str = "Ep0CCkgICxABGAIqQOWPB6/PmA5SW9jC6FvaNq3E+U9ev4FMWcFWuAho+VGLCtazKc5WDjQ5i0MuxsY0o5pKDSVWVKii8KJDusXH4eASDK7jyzuk8iij7fJNihoMxHO9haYzt48R36HVIjCb/EmIFrJLIXShqN6DN//T6vZBtO9qj1QhNWJa3CGm8VZoq80S2/Ok4U0aVaIDiZcqggHC2b8BHuv8BHZrmsR0wjU1ynansBGMdfjnG+iIv8R5lPpRmYGhSVwNybwP3aQZ6o8Dr48Rau8TJfdsArW+r7bvL7bPs4f5nnlp2vG7WkMzWwABHK3fdM44zZ1GZQaWyECNWR2GfY6dXiklo94vgpFTPuZ97mfiN3LY6uYyBwL8RkDaGAE="; - -// ── Helpers ────────────────────────────────────────────────────────────────── - -/// TS `TEST_PROMPT`: a system message + a user "Hello" message. -fn test_prompt() -> LanguageModelPrompt { - vec![ - LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text("System Prompt")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }, - ] -} - -fn default_options(prompt: LanguageModelPrompt) -> CallOptions { - CallOptions::new(prompt) -} - -fn make_model(server: &MockServer) -> BedrockModel { - BedrockModel::new( - MODEL_ID.to_string(), - BedrockConfig { - base_url: server.uri(), - auth: BedrockAuth::BearerToken("test-token".to_string()), - retry_config: aimux_provider_utils::RetryConfig::default(), - api_key_source: None, - }, - ) -} - -/// Mock the non-streaming `/converse` endpoint with a JSON body (HTTP 200). -async fn mock_converse(server: &MockServer, body: Value) { - let model_path = format!("/model/{MODEL_ID}/converse"); - Mock::given(method("POST")) - .and(path(model_path.as_str())) - .respond_with(ResponseTemplate::new(200).set_body_json(body)) - .mount(server) - .await; -} - -/// Encode a sequence of `(message-type, event-type, payload)` events into the -/// AWS event stream binary format used by `/converse-stream`. -fn encode_events(events: &[(&str, &str, Value)]) -> Vec { - let mut buf = Vec::new(); - for (mt, et, payload) in events { - buf.extend_from_slice(&event_stream::encode_message(mt, et, &payload.to_string())); - } - buf -} - -/// Mock the streaming `/converse-stream` endpoint with the given events. -async fn mock_stream(server: &MockServer, events: &[(&str, &str, Value)]) { - let body_bytes = encode_events(events); - let model_path = format!("/model/{MODEL_ID}/converse-stream"); - Mock::given(method("POST")) - .and(path(model_path.as_str())) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "application/vnd.amazon.eventstream") - .set_body_raw(body_bytes, "application/vnd.amazon.eventstream"), - ) - .mount(server) - .await; -} - -async fn collect_stream(result: StreamResult) -> Vec { - let mut parts = Vec::new(); - let mut stream = result.stream; - while let Some(part) = stream.next().await { - match part { - Ok(p) => parts.push(p), - Err(e) => panic!("stream error: {e:?}"), - } - } - parts -} - -fn as_text(item: &GenerateContent) -> &str { - match item { - GenerateContent::Text { text, .. } => text, - _ => panic!("expected Text content, got {item:?}"), - } -} - -fn as_reasoning(item: &GenerateContent) -> (&str, &Option) { - match item { - GenerateContent::Reasoning(ReasoningOutput { - text, - provider_metadata, - }) => (text.as_str(), provider_metadata), - _ => panic!("expected Reasoning content, got {item:?}"), - } -} - -fn reasoning_deltas(parts: &[StreamPart]) -> Vec { - parts - .iter() - .filter_map(|p| match p { - StreamPart::ReasoningDelta { delta, .. } => Some(delta.clone()), - _ => None, - }) - .collect() -} - -/// Reasoning deltas that carry actual text, i.e. excluding the zero-length -/// deltas that only transport a `signature` / `redactedData` in -/// `provider_metadata` (upstream emits those with `delta: ''`). -fn reasoning_text_deltas(parts: &[StreamPart]) -> Vec { - reasoning_deltas(parts) - .into_iter() - .filter(|d| !d.is_empty()) - .collect() -} - -/// The `signature` values carried on reasoning-block ends, in emission order. -/// The signature arrives on a trailing text-less delta and is attached to the -/// concluding `ReasoningEnd`. -fn reasoning_signatures(parts: &[StreamPart]) -> Vec { - parts - .iter() - .filter_map(|p| match p { - StreamPart::ReasoningEnd { - provider_metadata: Some(pm), - .. - } => pm - .get("amazonBedrock") - .and_then(|b| b.get("signature")) - .and_then(|s| s.as_str()) - .map(std::string::ToString::to_string), - _ => None, - }) - .collect() -} - -fn text_deltas(parts: &[StreamPart]) -> Vec { - parts - .iter() - .filter_map(|p| match p { - StreamPart::TextDelta { delta, .. } => Some(delta.clone()), - _ => None, - }) - .collect() -} - -fn has_reasoning_start(parts: &[StreamPart]) -> bool { - parts - .iter() - .any(|p| matches!(p, StreamPart::ReasoningStart { .. })) -} - -fn has_reasoning_end(parts: &[StreamPart]) -> bool { - parts - .iter() - .any(|p| matches!(p, StreamPart::ReasoningEnd { .. })) -} - -// ── convert helpers (prompt-side) ──────────────────────────────────────────── - -fn msg(role: Role, content: Vec) -> LanguageModelPromptMessage { - LanguageModelPromptMessage { - role, - content, - provider_options: None, - } -} - -fn user(content: Vec) -> LanguageModelPromptMessage { - msg(Role::User, content) -} - -fn assistant(content: Vec) -> LanguageModelPromptMessage { - msg(Role::Assistant, content) -} - -fn tool_msg(content: Vec) -> LanguageModelPromptMessage { - msg(Role::Tool, content) -} - -fn reasoning(text: &str, signature: Option<&str>) -> ContentPart { - ContentPart::Reasoning { - text: text.to_string(), - signature: signature.map(std::string::ToString::to_string), - provider_options: None, - } -} - -fn tool_call(id: &str, name: &str, input: Value) -> ContentPart { - ContentPart::tool_call(id.to_string(), name.to_string(), input) -} - -fn tool_result(id: &str, output: Value) -> ContentPart { - ContentPart::tool_result(id.to_string(), output) -} - -// ════════════════════════════════════════════════════════════════════════════ -// doGenerate — reasoning response extraction (TDD red until implemented) -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: `doGenerate > reasoning > should extract reasoning and text response`. -/// -/// Uses the `amazon-bedrock-reasoning.json` fixture: an assistant message with -/// a `reasoningContent.reasoningText` block (text + signature) followed by a -/// `text` block. The model should emit `Reasoning` then `Text`. -#[tokio::test] -async fn bedrock_generate_reasoning_response() { - let server = MockServer::start().await; - mock_converse( - &server, - json!({ - "metrics": { "latencyMs": 2202 }, - "output": { - "message": { - "role": "assistant", - "content": [ - { - "reasoningContent": { - "reasoningText": { - "signature": REASONING_SIGNATURE, - "text": "Let me count the r's in \"strawberry\":\n\ns-t-r-a-w-b-e-r-r-y\n\nThere are 3 r's." - } - } - }, - { - "text": "There are **3** r's in \"strawberry\":\n\n1. st**r**awbe**r****r**y" - } - ] - } - }, - "stopReason": "end_turn", - "usage": { - "cacheReadInputTokenCount": 0, - "cacheReadInputTokens": 0, - "cacheWriteInputTokenCount": 0, - "cacheWriteInputTokens": 0, - "inputTokens": 51, - "outputTokens": 78, - "serverToolUsage": {}, - "totalTokens": 129 - } - }), - ) - .await; - - let model = make_model(&server); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - // content[0] — reasoning with signature under both `amazonBedrock` and - // `bedrock` provider-metadata keys. - assert_eq!(result.content.len(), 2, "expected reasoning + text content"); - let (rtext, rmeta) = as_reasoning(&result.content[0]); - assert_eq!( - rtext, - "Let me count the r's in \"strawberry\":\n\ns-t-r-a-w-b-e-r-r-y\n\nThere are 3 r's." - ); - let rmeta = rmeta - .as_ref() - .expect("reasoning provider_metadata should be set"); - assert_eq!(rmeta["amazonBedrock"]["signature"], REASONING_SIGNATURE); - assert_eq!(rmeta["bedrock"]["signature"], REASONING_SIGNATURE); - - // content[1] — text. - assert_eq!( - as_text(&result.content[1]), - "There are **3** r's in \"strawberry\":\n\n1. st**r**awbe**r****r**y" - ); - - assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); - assert_eq!(result.finish_reason.raw.as_deref(), Some("end_turn")); - assert_eq!(result.usage.input_tokens.total, Some(51)); - assert_eq!(result.usage.output_tokens.total, Some(78)); -} - -/// TS: "should extract reasoning text with signature". -#[tokio::test] -async fn bedrock_generate_reasoning_text_with_signature() { - let server = MockServer::start().await; - mock_converse( - &server, - json!({ - "output": { - "message": { - "role": "assistant", - "content": [ - { - "reasoningContent": { - "reasoningText": { - "text": "I need to think about this problem carefully...", - "signature": "abc123signature" - } - } - }, - { "type": "text", "text": "The answer is 42." } - ] - } - }, - "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, - "stopReason": "stop_sequence" - }), - ) - .await; - - let model = make_model(&server); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - assert_eq!(result.content.len(), 2); - let (rtext, rmeta) = as_reasoning(&result.content[0]); - assert_eq!(rtext, "I need to think about this problem carefully..."); - let rmeta = rmeta.as_ref().expect("provider_metadata should be set"); - assert_eq!(rmeta["amazonBedrock"]["signature"], "abc123signature"); - assert_eq!(rmeta["bedrock"]["signature"], "abc123signature"); - - assert_eq!(as_text(&result.content[1]), "The answer is 42."); -} - -/// TS: "should preserve empty text blocks between reasoning blocks". -/// -/// The empty `{ text: "" }` block sitting between two reasoning blocks must be -/// preserved as a `Text` content item (not dropped). -#[tokio::test] -async fn bedrock_generate_reasoning_preserve_empty_text_between_reasoning() { - let server = MockServer::start().await; - mock_converse( - &server, - json!({ - "output": { - "message": { - "role": "assistant", - "content": [ - { "reasoningContent": { "reasoningText": { "text": "thinking...", "signature": "sig-1" } } }, - { "text": "" }, - { "reasoningContent": { "reasoningText": { "text": "more thinking...", "signature": "sig-2" } } }, - { "text": "The answer is 42." } - ] - } - }, - "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, - "stopReason": "stop_sequence" - }), - ) - .await; - - let model = make_model(&server); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - assert_eq!(result.content.len(), 4); - let (t0, m0) = as_reasoning(&result.content[0]); - assert_eq!(t0, "thinking..."); - assert_eq!(m0.as_ref().unwrap()["bedrock"]["signature"], "sig-1"); - - assert_eq!(as_text(&result.content[1]), ""); - - let (t2, m2) = as_reasoning(&result.content[2]); - assert_eq!(t2, "more thinking..."); - assert_eq!(m2.as_ref().unwrap()["bedrock"]["signature"], "sig-2"); - - assert_eq!(as_text(&result.content[3]), "The answer is 42."); -} - -/// TS: "should extract reasoning text without signature". -/// -/// A `reasoningText` block with no `signature` yields a `Reasoning` content item -/// with no provider metadata. -#[tokio::test] -async fn bedrock_generate_reasoning_text_without_signature() { - let server = MockServer::start().await; - mock_converse( - &server, - json!({ - "output": { - "message": { - "role": "assistant", - "content": [ - { "reasoningContent": { "reasoningText": { "text": "I need to think about this problem carefully..." } } }, - { "type": "text", "text": "The answer is 42." } - ] - } - }, - "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, - "stopReason": "stop_sequence" - }), - ) - .await; - - let model = make_model(&server); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - assert_eq!(result.content.len(), 2); - let (rtext, rmeta) = as_reasoning(&result.content[0]); - assert_eq!(rtext, "I need to think about this problem carefully..."); - assert!( - rmeta.is_none(), - "reasoning without signature should have no provider_metadata" - ); - assert_eq!(as_text(&result.content[1]), "The answer is 42."); -} - -/// TS: "should extract redacted reasoning". -/// -/// A `redactedReasoning` block yields a `Reasoning` content item with empty text -/// and `redactedData` under both `amazonBedrock` and `bedrock` provider metadata. -#[tokio::test] -async fn bedrock_generate_reasoning_redacted() { - let server = MockServer::start().await; - mock_converse( - &server, - json!({ - "output": { - "message": { - "role": "assistant", - "content": [ - { "reasoningContent": { "redactedReasoning": { "data": "redacted-reasoning-data" } } }, - { "type": "text", "text": "The answer is 42." } - ] - } - }, - "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, - "stopReason": "stop_sequence" - }), - ) - .await; - - let model = make_model(&server); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - assert_eq!(result.content.len(), 2); - let (rtext, rmeta) = as_reasoning(&result.content[0]); - assert_eq!(rtext, ""); - let rmeta = rmeta.as_ref().expect("provider_metadata should be set"); - assert_eq!( - rmeta["amazonBedrock"]["redactedData"], - "redacted-reasoning-data" - ); - assert_eq!(rmeta["bedrock"]["redactedData"], "redacted-reasoning-data"); - assert_eq!(as_text(&result.content[1]), "The answer is 42."); -} - -/// TS: "should handle multiple reasoning blocks". -/// -/// A signed `reasoningText` block followed by a `redactedReasoning` block then -/// text — three content items in order. -#[tokio::test] -async fn bedrock_generate_reasoning_multiple_blocks() { - let server = MockServer::start().await; - mock_converse( - &server, - json!({ - "output": { - "message": { - "role": "assistant", - "content": [ - { "reasoningContent": { "reasoningText": { "text": "First reasoning block", "signature": "sig1" } } }, - { "reasoningContent": { "redactedReasoning": { "data": "redacted-data" } } }, - { "type": "text", "text": "The answer is 42." } - ] - } - }, - "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, - "stopReason": "stop_sequence" - }), - ) - .await; - - let model = make_model(&server); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - assert_eq!(result.content.len(), 3); - - let (t0, m0) = as_reasoning(&result.content[0]); - assert_eq!(t0, "First reasoning block"); - assert_eq!(m0.as_ref().unwrap()["bedrock"]["signature"], "sig1"); - - let (t1, m1) = as_reasoning(&result.content[1]); - assert_eq!(t1, ""); - assert_eq!( - m1.as_ref().unwrap()["bedrock"]["redactedData"], - "redacted-data" - ); - - assert_eq!(as_text(&result.content[2]), "The answer is 42."); -} - -// ════════════════════════════════════════════════════════════════════════════ -// doStream — reasoning streaming (TDD red until implemented) -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: `doStream > reasoning > should stream reasoning and text parts`. -/// -/// Uses the `amazon-bedrock-reasoning.chunks.txt` fixture: a reasoning block -/// (index 0) streamed as `reasoningContent.text` deltas + a final -/// `reasoningContent.signature` delta, then a text block (index 1). -#[tokio::test] -async fn bedrock_stream_reasoning_and_text() { - let server = MockServer::start().await; - - let events: Vec<(&str, &str, Value)> = vec![ - ("event", "messageStart", json!({ "role": "assistant" })), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": "Let me count the r" } } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": "'s in \"" } } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": "strawberry\":" } } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": "\n\ns-t-r-a" } } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": "-w-b-e-r" } } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": "-r-y\n\nr" } } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": " appears at positions " } } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": "3, 8" } } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": ", and 9.\n\nSo there" } } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": " are 3 r's." } } }), - ), - // Empty-text delta — TS skips emitting a reasoning-delta for this. - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": "" } } }), - ), - // Signature-only delta — emitted as a reasoning-delta with delta="" - // carrying providerMetadata.signature, matching TS. - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "signature": STREAM_SIGNATURE } } }), - ), - ( - "event", - "contentBlockStop", - json!({ "contentBlockIndex": 0 }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 1, "delta": { "text": "There" } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 1, "delta": { "text": " are **3** r's in \"" } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 1, "delta": { "text": "strawberry\":" } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 1, "delta": { "text": "\n\n1" } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 1, "delta": { "text": ". st" } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 1, "delta": { "text": "**" } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 1, "delta": { "text": "r**awbe" } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 1, "delta": { "text": "**r****" } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 1, "delta": { "text": "r**y" } }), - ), - ( - "event", - "contentBlockStop", - json!({ "contentBlockIndex": 1 }), - ), - ( - "event", - "messageStop", - json!({ "additionalModelResponseFields": { "delta": { "stop_sequence": null } }, "stopReason": "end_turn" }), - ), - ( - "event", - "metadata", - json!({ "metrics": { "latencyMs": 2281 }, "usage": { "inputTokens": 51, "outputTokens": 94, "serverToolUsage": {}, "totalTokens": 145 } }), - ), - ]; - mock_stream(&server, &events).await; - - let model = make_model(&server); - let result = model - .do_stream(&default_options(test_prompt())) - .await - .expect("do_stream should succeed"); - let parts = collect_stream(result).await; - - // Reasoning block (id "0"): start, non-empty text deltas, end. - assert!(has_reasoning_start(&parts), "expected a ReasoningStart"); - assert_eq!( - reasoning_text_deltas(&parts), - vec![ - "Let me count the r", - "'s in \"", - "strawberry\":", - "\n\ns-t-r-a", - "-w-b-e-r", - "-r-y\n\nr", - " appears at positions ", - "3, 8", - ", and 9.\n\nSo there", - " are 3 r's.", - ] - ); - assert!(has_reasoning_end(&parts), "expected a ReasoningEnd"); - // The signature on the final (text-less) reasoning delta is attached to - // the concluding ReasoningEnd via provider_metadata, so extended-thinking - // multi-turn can echo it back (#113). - let reasoning_end_meta = parts.iter().find_map(|p| match p { - StreamPart::ReasoningEnd { - provider_metadata, .. - } => provider_metadata.clone(), - _ => None, - }); - let reasoning_end_meta = reasoning_end_meta.expect("ReasoningEnd with provider_metadata"); - assert_eq!( - reasoning_end_meta["bedrock"]["signature"].as_str(), - Some(STREAM_SIGNATURE) - ); - // Dual-key shape matches the non-streaming path. - assert_eq!( - reasoning_end_meta["amazonBedrock"]["signature"].as_str(), - Some(STREAM_SIGNATURE) - ); - - // Text block (id "1"). - assert_eq!( - text_deltas(&parts), - vec![ - "There", - " are **3** r's in \"", - "strawberry\":", - "\n\n1", - ". st", - "**", - "r**awbe", - "**r****", - "r**y", - ] - ); - - // Finish. - let finish = parts.iter().find_map(|p| match p { - StreamPart::Finish { finish_reason, .. } => Some(finish_reason.clone()), - _ => None, - }); - let finish = finish.expect("expected a Finish part"); - assert_eq!(finish.unified, FinishReasonUnified::Stop); - assert_eq!(finish.raw.as_deref(), Some("end_turn")); -} - -/// TS: "should stream reasoning text deltas". -/// -/// Inline chunks: two reasoning text deltas, a signature-only delta, then a -/// text delta and `messageStop`. -#[tokio::test] -async fn bedrock_stream_reasoning_text_deltas() { - let server = MockServer::start().await; - let events: Vec<(&str, &str, Value)> = vec![ - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": "I am thinking" } } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "text": " about this problem..." } } }), - ), - // Signature-only delta — emitted as a reasoning-delta with delta="" - // carrying providerMetadata.signature, matching TS. - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 0, "delta": { "reasoningContent": { "signature": "abc123signature" } } }), - ), - ( - "event", - "contentBlockDelta", - json!({ "contentBlockIndex": 1, "delta": { "text": "Based on my reasoning, the answer is 42." } }), - ), - ( - "event", - "messageStop", - json!({ "stopReason": "stop_sequence" }), - ), - ]; - mock_stream(&server, &events).await; - - let model = make_model(&server); - let result = model - .do_stream(&default_options(test_prompt())) - .await - .expect("do_stream should succeed"); - let parts = collect_stream(result).await; - - assert!(has_reasoning_start(&parts), "expected a ReasoningStart"); - assert_eq!( - reasoning_text_deltas(&parts), - vec!["I am thinking", " about this problem..."] - ); - assert_eq!( - reasoning_signatures(&parts), - vec!["abc123signature"], - "the thinking-block signature must reach the caller" - ); - assert!(has_reasoning_end(&parts), "expected a ReasoningEnd"); - - assert_eq!( - text_deltas(&parts), - vec!["Based on my reasoning, the answer is 42."] - ); - - let finish = parts.iter().find_map(|p| match p { - StreamPart::Finish { finish_reason, .. } => Some(finish_reason.clone()), - _ => None, - }); - let finish = finish.expect("expected a Finish part"); - assert_eq!(finish.unified, FinishReasonUnified::Stop); - assert_eq!(finish.raw.as_deref(), Some("stop_sequence")); -} - -// ════════════════════════════════════════════════════════════════════════════ -// doStream — request-side reasoningConfig → thinking (TDD red until implemented) -// ════════════════════════════════════════════════════════════════════════════ - -// ════════════════════════════════════════════════════════════════════════════ -// convert_prompt_to_bedrock — prompt-side reasoning replay -// (green for signed/unsigned reasoning; #[ignore] where the data model cannot -// express redacted / foreign-provider reasoning) -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: "should properly convert reasoning content type". -/// -/// An assistant reasoning part carrying a `bedrock` signature is replayed as a -/// `reasoningContent.reasoningText { text, signature }` block. -#[test] -fn convert_reasoning_content_type() { - let prompt = vec![ - user(vec![ContentPart::text("Explain your reasoning")]), - assistant(vec![reasoning( - "This is my step-by-step reasoning process", - Some("test-signature"), - )]), - ]; - let (system, messages) = convert_prompt_to_bedrock(&prompt); - - assert_eq!(Value::Array(system), json!([])); - assert_eq!( - Value::Array(messages), - json!([ - { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, - { - "role": "assistant", - "content": [ - { - "reasoningContent": { - "reasoningText": { - "text": "This is my step-by-step reasoning process", - "signature": "test-signature" - } - } - } - ] - } - ]) - ); -} - -/// TS: "should omit assistant message reasoning parts signed by a foreign -/// provider". -/// -/// SKIPPED: `ContentPart::Reasoning.signature` is provider-agnostic, so the -/// converter cannot distinguish a foreign-provider signature (which must be -/// dropped) from a Bedrock one (which is kept). The intended behaviour is that -/// the reasoning part is omitted and only the trailing text remains. -#[test] -#[ignore = "ContentPart::Reasoning cannot express provider distinction (foreign vs bedrock signature)"] -fn convert_reasoning_foreign_provider_omitted() { - let prompt = vec![ - user(vec![ContentPart::text("Explain your reasoning")]), - assistant(vec![ - // In TS this carries `providerOptions.anthropic.signature` and is - // dropped because it is not a Bedrock signature. - reasoning( - "Anthropic-signed reasoning replayed to Bedrock", - Some("anthropic-signature"), - ), - ContentPart::text("final answer"), - ]), - ]; - let (system, messages) = convert_prompt_to_bedrock(&prompt); - - assert_eq!(Value::Array(system), json!([])); - assert_eq!( - Value::Array(messages), - json!([ - { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, - { "role": "assistant", "content": [{ "text": "final answer" }] } - ]) - ); -} - -/// TS: "should preserve assistant message reasoning parts with amazonBedrock -/// providerOptions". -/// -/// A reasoning part with a Bedrock signature is preserved. (The Rust data model -/// does not distinguish `amazonBedrock` from `bedrock` provider options — the -/// signature is a direct field — but the replay behaviour matches.) -#[test] -fn convert_reasoning_amazon_bedrock_provider_options_preserved() { - let prompt = vec![ - user(vec![ContentPart::text("Explain your reasoning")]), - assistant(vec![ - reasoning( - "Bedrock-signed reasoning round-tripped to Bedrock", - Some("bedrock-signature"), - ), - ContentPart::text("final answer"), - ]), - ]; - let (system, messages) = convert_prompt_to_bedrock(&prompt); - - assert_eq!(Value::Array(system), json!([])); - assert_eq!( - Value::Array(messages), - json!([ - { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, - { - "role": "assistant", - "content": [ - { - "reasoningContent": { - "reasoningText": { - "text": "Bedrock-signed reasoning round-tripped to Bedrock", - "signature": "bedrock-signature" - } - } - }, - { "text": "final answer" } - ] - } - ]) - ); -} - -/// TS: "should not trim reasoning text when a signature is present". -/// -/// Trailing whitespace on reasoning text must be preserved verbatim (only the -/// last *text* part of the last block is trimmed, never reasoning). -#[test] -fn convert_reasoning_no_trim_with_signature() { - let prompt = vec![ - user(vec![ContentPart::text("Explain your reasoning")]), - assistant(vec![reasoning( - "This is my reasoning with trailing space ", - Some("test-signature"), - )]), - ]; - let (system, messages) = convert_prompt_to_bedrock(&prompt); - - assert_eq!(Value::Array(system), json!([])); - assert_eq!( - Value::Array(messages), - json!([ - { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, - { - "role": "assistant", - "content": [ - { - "reasoningContent": { - "reasoningText": { - "text": "This is my reasoning with trailing space ", - "signature": "test-signature" - } - } - } - ] - } - ]) - ); -} - -/// TS: "should omit reasoning content without signature". -/// -/// Unsigned reasoning is dropped; the trailing text is kept. -#[test] -fn convert_reasoning_omit_without_signature() { - let prompt = vec![ - user(vec![ContentPart::text("Explain your reasoning")]), - assistant(vec![ - reasoning("This is my reasoning with trailing space ", None), - ContentPart::text("final answer"), - ]), - ]; - let (system, messages) = convert_prompt_to_bedrock(&prompt); - - assert_eq!(Value::Array(system), json!([])); - assert_eq!( - Value::Array(messages), - json!([ - { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, - { "role": "assistant", "content": [{ "text": "final answer" }] } - ]) - ); -} - -/// TS: "should omit multiple reasoning parts without signatures". -#[test] -fn convert_reasoning_omit_multiple_without_signatures() { - let prompt = vec![ - user(vec![ContentPart::text("Explain your reasoning")]), - assistant(vec![ - reasoning("First reasoning with trailing space ", None), - reasoning("Second reasoning with trailing space ", None), - ContentPart::text("final answer"), - ]), - ]; - let (system, messages) = convert_prompt_to_bedrock(&prompt); - - assert_eq!(Value::Array(system), json!([])); - assert_eq!( - Value::Array(messages), - json!([ - { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, - { "role": "assistant", "content": [{ "text": "final answer" }] } - ]) - ); -} - -/// TS: "should omit unsigned reasoning while preserving tool calls in -/// multi-turn tool use". -#[test] -fn convert_reasoning_omit_unsigned_preserve_tool_calls() { - let prompt = vec![ - user(vec![ContentPart::text("What is the weather?")]), - assistant(vec![ - reasoning("I should call the weather tool.", None), - tool_call("call-1", "getWeather", json!({ "city": "SF" })), - ]), - tool_msg(vec![tool_result( - "call-1", - json!({ "type": "text", "value": "Sunny, 72F" }), - )]), - ]; - let (system, messages) = convert_prompt_to_bedrock(&prompt); - - assert_eq!(Value::Array(system), json!([])); - assert_eq!( - Value::Array(messages), - json!([ - { "role": "user", "content": [{ "text": "What is the weather?" }] }, - { - "role": "assistant", - "content": [ - { - "toolUse": { - "input": { "city": "SF" }, - "name": "getWeather", - "toolUseId": "call-1" - } - } - ] - }, - { - "role": "user", - "content": [ - { - "toolResult": { - "content": [{ "text": "Sunny, 72F" }], - "toolUseId": "call-1" - } - } - ] - } - ]) - ); -} - -/// TS: "should preserve reasoning text with signature in multi-turn tool use". -#[test] -fn convert_reasoning_preserve_signed_in_multi_turn_tool_use() { - let prompt = vec![ - user(vec![ContentPart::text("What is the weather?")]), - assistant(vec![ - reasoning("Let me check the weather API.\n", Some("sig-abc123")), - tool_call("call-1", "getWeather", json!({ "city": "SF" })), - ]), - tool_msg(vec![tool_result( - "call-1", - json!({ "type": "text", "value": "Sunny, 72F" }), - )]), - assistant(vec![ - reasoning("The weather is sunny and warm.\n", Some("sig-def456")), - ContentPart::text("It is sunny and 72F in SF."), - ]), - ]; - let (system, messages) = convert_prompt_to_bedrock(&prompt); - - assert_eq!(Value::Array(system), json!([])); - assert_eq!( - Value::Array(messages), - json!([ - { "role": "user", "content": [{ "text": "What is the weather?" }] }, - { - "role": "assistant", - "content": [ - { - "reasoningContent": { - "reasoningText": { - "text": "Let me check the weather API.\n", - "signature": "sig-abc123" - } - } - }, - { - "toolUse": { - "input": { "city": "SF" }, - "name": "getWeather", - "toolUseId": "call-1" - } - } - ] - }, - { - "role": "user", - "content": [ - { - "toolResult": { - "content": [{ "text": "Sunny, 72F" }], - "toolUseId": "call-1" - } - } - ] - }, - { - "role": "assistant", - "content": [ - { - "reasoningContent": { - "reasoningText": { - "text": "The weather is sunny and warm.\n", - "signature": "sig-def456" - } - } - }, - { "text": "It is sunny and 72F in SF." } - ] - } - ]) - ); -} - -/// TS: "should handle a mix of text and reasoning content types". -#[test] -fn convert_reasoning_mix_text_and_reasoning() { - let prompt = vec![ - user(vec![ContentPart::text("Explain your reasoning")]), - assistant(vec![ - ContentPart::text("My answer is 42."), - reasoning( - "I calculated this by analyzing the meaning of life", - Some("reasoning-process"), - ), - ]), - ]; - let (system, messages) = convert_prompt_to_bedrock(&prompt); - - assert_eq!(Value::Array(system), json!([])); - assert_eq!( - Value::Array(messages), - json!([ - { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, - { - "role": "assistant", - "content": [ - { "text": "My answer is 42." }, - { - "reasoningContent": { - "reasoningText": { - "text": "I calculated this by analyzing the meaning of life", - "signature": "reasoning-process" - } - } - } - ] - } - ]) - ); -} diff --git a/aimux-providers/tests/bedrock_stream_replay_test.rs b/aimux-providers/tests/bedrock_stream_replay_test.rs index 880761d37..5fe4eafa6 100644 --- a/aimux-providers/tests/bedrock_stream_replay_test.rs +++ b/aimux-providers/tests/bedrock_stream_replay_test.rs @@ -24,10 +24,8 @@ use common::replay::mount_cassettes; use futures::StreamExt; use wiremock::MockServer; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::stream_part::StreamPart; use aimux_providers::{BedrockProvider, BedrockProviderConfig}; @@ -38,11 +36,9 @@ use aimux_providers::{BedrockProvider, BedrockProviderConfig}; const MODEL_ID: &str = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"; fn prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("What is the capital of France?")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text( + "What is the capital of France?", + )] } /// Decoded frames must reach the parser and come out as text plus a finish. diff --git a/aimux-providers/tests/body_overrides_test.rs b/aimux-providers/tests/body_overrides_test.rs deleted file mode 100644 index f8b392dcb..000000000 --- a/aimux-providers/tests/body_overrides_test.rs +++ /dev/null @@ -1,286 +0,0 @@ -//! Tests for RFC-0017 phase 1: `body_overrides` (JSON deep-merge) and -//! `max_retries` (per-call retry override). -//! -//! `body_overrides` is a per-call JSON object deep-merged into the -//! provider-built request body after built-in vendor overrides. It lets users -//! inject/override arbitrary request fields (e.g. `enable_thinking`, -//! `thinking_budget`) without closure bridging — critical for aimux's -//! multi-language C ABI architecture. `null` values delete keys. -//! -//! stage2-001 (RFC-0017 phase 2) additions: `max_tokens_key` branch and direct -//! `reasoning` → `reasoning_effort` passthrough. The old "reasoning no-mapping" -//! warning block was removed (F6): under v3 direct-passthrough semantics it is -//! unreachable dead code (is_custom_reasoning ⇒ resolved effort is always Some), -//! and the passthrough tests below now guard against re-adding it. - -use aimux_core::content::ContentPart; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; -use aimux_core::options::CallOptions; -use aimux_core::types::{ReasoningEffort, Warning}; -use aimux_providers::openai::convert::build_request_body; -use aimux_providers::openai::convert::build_request_body_with_warnings; -use serde_json::{Value, json}; - -fn user_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] -} - -fn opts_with_overrides(prompt: LanguageModelPrompt, body_overrides: Value) -> CallOptions { - CallOptions { - prompt, - body_overrides: Some(body_overrides), - ..CallOptions::default() - } -} - -// ── body_overrides: inject ─────────────────────────────────────────────────── - -/// A top-level key in body_overrides is injected into the request body. -#[test] -fn body_overrides_injects_new_field() { - let opts = opts_with_overrides(user_prompt(), json!({ "enable_thinking": false })); - let body = build_request_body("gpt-4o", &opts, false).unwrap(); - assert_eq!(body["enable_thinking"], json!(false)); - assert_eq!(body["model"], json!("gpt-4o")); -} - -/// body_overrides can inject nested objects. -#[test] -fn body_overrides_injects_nested_object() { - let opts = opts_with_overrides(user_prompt(), json!({ "thinking": { "type": "disabled" } })); - let body = build_request_body("gpt-4o", &opts, false).unwrap(); - assert_eq!(body["thinking"], json!({ "type": "disabled" })); -} - -// ── body_overrides: override ───────────────────────────────────────────────── - -/// body_overrides overwrites an existing field (e.g. temperature). -#[test] -fn body_overrides_overwrites_existing_field() { - let mut opts = opts_with_overrides(user_prompt(), json!({ "temperature": 0.1 })); - opts.temperature = Some(0.9); // set by standard option - let body = build_request_body("gpt-4o", &opts, false).unwrap(); - // body_overrides wins over standard option - assert_eq!(body["temperature"], json!(0.1)); -} - -/// body_overrides injects a field into the request body (RFC-0017). -/// stage2-001（RFC-0017 阶段 2）后内置 vendor override 已退役——此前 DeepSeek -/// 会从 `reasoning:none` 注入 `thinking:{type:"disabled"}`;现在 thinking 注入 -/// 完全由 body_overrides 定义（此处直接注入 enabled）。 -#[test] -fn body_overrides_overwrites_vendor_override_field() { - let opts = CallOptions { - prompt: user_prompt(), - reasoning: Some(aimux_core::types::ReasoningEffort::None), - body_overrides: Some(json!({ "thinking": { "type": "enabled" } })), - ..CallOptions::default() - }; - // DeepSeek profile 已回归 full()（特化退役）,body_overrides 注入 thinking。 - let body = aimux_providers::openai::convert::build_request_body_with_warnings( - "deepseek-v4-flash", - &opts, - false, - "deepseek", - &aimux_providers::openai::OpenAICompatProfile::deepseek(), - ) - .unwrap() - .body; - assert_eq!(body["thinking"], json!({ "type": "enabled" })); -} - -// ── body_overrides: deep merge ─────────────────────────────────────────────── - -/// Nested objects are merged recursively, not replaced wholesale. -#[test] -fn body_overrides_deep_merges_nested_objects() { - // The standard body has stream_options: { include_usage: true } for streams. - // body_overrides adds another key to stream_options without clobbering - // include_usage. - let opts = opts_with_overrides( - user_prompt(), - json!({ "stream_options": { "include_usage": false } }), - ); - let body = build_request_body("gpt-4o", &opts, true).unwrap(); - assert_eq!(body["stream_options"]["include_usage"], json!(false)); -} - -// ── body_overrides: null = delete ──────────────────────────────────────────── - -/// A `null` value in body_overrides deletes the corresponding key from the -/// request body. -#[test] -fn body_overrides_null_deletes_key() { - let opts = opts_with_overrides( - user_prompt(), - json!({ "stream_options": null, "temperature": null }), - ); - let body = build_request_body("gpt-4o", &opts, true).unwrap(); - assert!( - body.get("stream_options").is_none(), - "stream_options should be deleted by null" - ); - assert!( - body.get("temperature").is_none(), - "temperature should be deleted by null" - ); -} - -/// null in a nested object deletes the nested key. -#[test] -fn body_overrides_null_deletes_nested_key() { - let opts = opts_with_overrides( - user_prompt(), - json!({ "stream_options": { "include_usage": null } }), - ); - let body = build_request_body("gpt-4o", &opts, true).unwrap(); - // stream_options still exists but include_usage is gone - assert!(body.get("stream_options").is_some()); - assert!(body["stream_options"].get("include_usage").is_none()); -} - -// ── body_overrides: no overrides = unchanged ───────────────────────────────── - -/// When body_overrides is None, the request body is identical to the standard -/// build (no merge happens). -#[test] -fn no_body_overrides_leaves_body_unchanged() { - let opts = CallOptions::new(user_prompt()); - let body = build_request_body("gpt-4o", &opts, false).unwrap(); - assert_eq!(body["model"], json!("gpt-4o")); - assert!(body.get("enable_thinking").is_none()); -} - -// ── max_retries ────────────────────────────────────────────────────────────── - -/// max_retries is stored in CallOptions and defaults to None. -#[test] -fn max_retries_defaults_to_none() { - let opts = CallOptions::new(user_prompt()); - assert!(opts.max_retries.is_none()); -} - -/// max_retries can be set (e.g. Some(0) to disable retries). -#[test] -fn max_retries_can_be_set() { - let opts = CallOptions { - prompt: user_prompt(), - max_retries: Some(0), - ..CallOptions::default() - }; - assert_eq!(opts.max_retries, Some(0)); -} - -// ── max_tokens_key (stage2-001, RFC-0017 phase 2 §2.3) ────────────────────── - -/// max_tokens_key=Some("max_tokens") + 推理模型 → 发 `max_tokens`, -/// 不含 `max_completion_tokens`（修推理模型推断 bug）。 -#[test] -fn max_tokens_key_max_tokens_reasoning_model() { - let opts = CallOptions { - prompt: user_prompt(), - max_output_tokens: Some(100), - ..CallOptions::default() - }; - let profile = aimux_providers::openai::OpenAICompatProfile { - max_tokens_key: Some("max_tokens"), - ..aimux_providers::openai::OpenAICompatProfile::full() - }; - let body = build_request_body_with_warnings("o4-mini", &opts, false, "openai", &profile) - .unwrap() - .body; - assert_eq!(body["max_tokens"], json!(100)); - assert!( - body.get("max_completion_tokens").is_none(), - "只认 max_tokens 的厂商不应收到 max_completion_tokens" - ); -} - -/// max_tokens_key=Some("max_completion_tokens") → 非推理分支也发 mct -///（groq/heroku:max_tokens 弃用）。 -#[test] -fn max_tokens_key_max_completion_tokens_non_reasoning() { - let opts = CallOptions { - prompt: user_prompt(), - max_output_tokens: Some(100), - ..CallOptions::default() - }; - let profile = aimux_providers::openai::OpenAICompatProfile { - max_tokens_key: Some("max_completion_tokens"), - ..aimux_providers::openai::OpenAICompatProfile::full() - }; - let body = build_request_body_with_warnings("gpt-4o", &opts, false, "openai", &profile) - .unwrap() - .body; - assert_eq!(body["max_completion_tokens"], json!(100)); - assert!( - body.get("max_tokens").is_none(), - "max_tokens_key=mct 时不应再发 max_tokens" - ); -} - -// ── reasoning 直传: 无 warning（F6: 旧"无映射提示"warning 块已删除）──────────── -// -// v3 直传语义下 `reasoning` 一律映射为 `reasoning_effort`,warning 块不可达已被 -// 删除。以下两个用例断言"无 reasoning warning",作为未来误加 warning 的回归护栏。 - -/// 直传路径: groq 不再特化归一化(`none` 原样透传 'none')→ 已发 effort,不 warning。 -#[test] -fn groq_none_passthrough_no_warning() { - let opts = CallOptions { - prompt: user_prompt(), - reasoning: Some(ReasoningEffort::None), - ..CallOptions::default() - }; - let result = build_request_body_with_warnings( - "llama-3.3-70b-versatile", - &opts, - false, - "groq", - &aimux_providers::openai::OpenAICompatProfile::groq(), - ) - .unwrap(); - assert_eq!(result.body["reasoning_effort"], json!("none")); - let reasoning_warning = result - .warnings - .iter() - .find(|w| matches!(w, Warning::Compatibility { feature, .. } if feature == "reasoning")); - assert!( - reasoning_warning.is_none(), - "已发 effort 时不应 warning(防误报): {:?}", - result.warnings - ); -} - -/// 已发 effort 路径（'none' 透传,OpenAI 有效）→ 不 warning（防误报）。 -#[test] -fn no_warning_when_reasoning_translated() { - let opts = CallOptions { - prompt: user_prompt(), - reasoning: Some(ReasoningEffort::None), - ..CallOptions::default() - }; - let result = build_request_body_with_warnings( - "deepseek-reasoner", - &opts, - false, - "openai", - &aimux_providers::openai::OpenAICompatProfile::full(), - ) - .unwrap(); - assert_eq!(result.body["reasoning_effort"], json!("none")); - let reasoning_warning = result - .warnings - .iter() - .find(|w| matches!(w, Warning::Compatibility { feature, .. } if feature == "reasoning")); - assert!( - reasoning_warning.is_none(), - "已发 effort 时不应 warning(防误报): {:?}", - result.warnings - ); -} diff --git a/aimux-providers/tests/cerebras_test.rs b/aimux-providers/tests/cerebras_test.rs index 77188c825..b58d5ad16 100644 --- a/aimux-providers/tests/cerebras_test.rs +++ b/aimux-providers/tests/cerebras_test.rs @@ -25,20 +25,14 @@ use serial_test::serial; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_providers::{ProviderOptions, provider, provider_from_env}; fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn default_options(prompt: LanguageModelPrompt) -> CallOptions { diff --git a/aimux-providers/tests/codex_test.rs b/aimux-providers/tests/codex_test.rs index cd957ac24..481f92096 100644 --- a/aimux-providers/tests/codex_test.rs +++ b/aimux-providers/tests/codex_test.rs @@ -13,11 +13,9 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, StreamResult}; use aimux_core::stream_part::StreamPart; @@ -27,11 +25,7 @@ use aimux_providers::{CODEX_OAUTH_TOKEN_URL, CodexConfig, CodexProvider, codex_r // helpers fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn default_options(prompt: LanguageModelPrompt) -> CallOptions { diff --git a/aimux-providers/tests/conformance_test.rs b/aimux-providers/tests/conformance_test.rs index 42f571ffe..5fe52f920 100644 --- a/aimux-providers/tests/conformance_test.rs +++ b/aimux-providers/tests/conformance_test.rs @@ -15,10 +15,8 @@ use common::replay::mount_cassettes; use futures::StreamExt; use wiremock::MockServer; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; @@ -35,11 +33,7 @@ use aimux_providers::{ // 閳光偓閳光偓 helpers 閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓 fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn default_options(prompt: LanguageModelPrompt) -> CallOptions { diff --git a/aimux-providers/tests/data_loss_regression_test.rs b/aimux-providers/tests/data_loss_regression_test.rs index e9533443a..c3bb88c49 100644 --- a/aimux-providers/tests/data_loss_regression_test.rs +++ b/aimux-providers/tests/data_loss_regression_test.rs @@ -32,10 +32,8 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path as path_matcher}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, ReasoningOutput, Source, StreamResult}; use aimux_core::stream_part::StreamPart; @@ -120,11 +118,7 @@ fn cassette_json(rel: &str) -> Value { // ═══════════════════════════════════════════════════════════════════════════ fn user_prompt(text: &str) -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text(text)], - ..Default::default() - }] + vec![LanguageModelMessage::user_text(text)] } fn opts() -> CallOptions { diff --git a/aimux-providers/tests/deepinfra_test.rs b/aimux-providers/tests/deepinfra_test.rs index abdbce70c..a1d6e3799 100644 --- a/aimux-providers/tests/deepinfra_test.rs +++ b/aimux-providers/tests/deepinfra_test.rs @@ -24,20 +24,14 @@ use serial_test::serial; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_providers::{ProviderOptions, provider, provider_from_env}; fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn default_options(prompt: LanguageModelPrompt) -> CallOptions { diff --git a/aimux-providers/tests/deepseek_chat_test.rs b/aimux-providers/tests/deepseek_chat_test.rs index 5b16ace99..0ff6132c9 100644 --- a/aimux-providers/tests/deepseek_chat_test.rs +++ b/aimux-providers/tests/deepseek_chat_test.rs @@ -31,12 +31,13 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + FilePart, LanguageModelMessage, LanguageModelPrompt, TextPart, UserPart, +}; use aimux_core::options::{CallOptions, Tool}; use aimux_core::result::GenerateContent; +use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; use aimux_core::tool::FunctionTool; use aimux_core::types::FinishReasonUnified; @@ -51,11 +52,7 @@ use aimux_providers::{ProviderOptions, provider}; /// The TS `TEST_PROMPT`: a single user text message "Hello". fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } /// `CallOptions` with only `prompt` set (everything else default/None). @@ -346,13 +343,22 @@ async fn should_convert_user_text_part_to_string_content() { /// TS: file parts with image media type are converted to image_url (OpenAI format). #[tokio::test] async fn should_convert_image_file_parts_to_image_url() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, + let prompt = vec![LanguageModelMessage::User { content: vec![ - ContentPart::text("Hello"), - ContentPart::file_base64("AAECAw==", "image/png"), + UserPart::Text(TextPart { + text: "Hello".into(), + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64("AAECAw==".into()), + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + }), ], - ..Default::default() + provider_options: None, }]; let options = default_options(prompt); let result = build_request_body_with_warnings( @@ -384,13 +390,22 @@ async fn should_convert_image_file_parts_to_image_url() { /// NOT include any image/mediaType data in the output. #[tokio::test] async fn should_accept_top_level_only_mediatype_without_error() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, + let prompt = vec![LanguageModelMessage::User { content: vec![ - ContentPart::text("Hello"), - ContentPart::file_base64("AAECAw==", "image"), + UserPart::Text(TextPart { + text: "Hello".into(), + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64("AAECAw==".into()), + }, + media_type: "image".into(), + filename: None, + provider_options: None, + }), ], - ..Default::default() + provider_options: None, }]; let options = default_options(prompt); let result = build_request_body_with_warnings( @@ -410,127 +425,6 @@ async fn should_accept_top_level_only_mediatype_without_error() { ); } -/// TS: tool calls › "should stringify arguments to tool calls" (line ~100). -/// -/// An assistant tool-call message followed by a tool-result message should -/// produce: -/// - assistant: `{ role: "assistant", content: "", tool_calls: [...] }` -/// - tool: `{ role: "tool", content: "...", tool_call_id: "..." }` -#[tokio::test] -async fn should_stringify_arguments_to_tool_calls() { - let prompt = vec![ - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::tool_call( - "quux", - "thwomp", - json!({ "foo": "bar123" }), - )], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result("quux", json!({ "oof": "321rab" }))], - ..Default::default() - }, - ]; - let options = default_options(prompt); - let result = build_request_body_with_warnings( - "deepseek-chat", - &options, - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - - let messages = result.body["messages"].as_array().expect("messages array"); - assert_eq!(messages.len(), 2); - - // Assistant message with tool_calls. - assert_eq!(messages[0]["role"], json!("assistant")); - // OpenAI converter sets content to null when empty. - assert!(messages[0]["content"].is_null()); - assert_eq!( - messages[0]["tool_calls"], - json!([{ - "id": "quux", - "type": "function", - "function": { - "name": "thwomp", - "arguments": "{\"foo\":\"bar123\"}" - } - }]) - ); - - // Tool result message. - assert_eq!(messages[1]["role"], json!("tool")); - assert_eq!(messages[1]["content"], json!("{\"oof\":\"321rab\"}")); - assert_eq!(messages[1]["tool_call_id"], json!("quux")); - - assert!(result.warnings.is_empty(), "expected no warnings"); -} - -/// TS: tool calls › "should handle text output type in tool results" (line ~159). -/// -/// A tool result with a text output should produce the text value directly as -/// the tool message content. -#[tokio::test] -async fn should_handle_text_output_type_in_tool_results() { - let prompt = vec![ - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::tool_call( - "call-1", - "getWeather", - json!({ "query": "weather" }), - )], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result( - "call-1", - json!("It is sunny today"), - )], - ..Default::default() - }, - ]; - let options = default_options(prompt); - let result = build_request_body_with_warnings( - "deepseek-chat", - &options, - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - - let messages = result.body["messages"].as_array().expect("messages array"); - assert_eq!(messages.len(), 2); - - // Assistant message. - assert_eq!(messages[0]["role"], json!("assistant")); - // OpenAI converter sets content to null when empty. - assert!(messages[0]["content"].is_null()); - assert_eq!( - messages[0]["tool_calls"], - json!([{ - "id": "call-1", - "type": "function", - "function": { - "name": "getWeather", - "arguments": "{\"query\":\"weather\"}" - } - }]) - ); - - // Tool result message with text output. - assert_eq!(messages[1]["role"], json!("tool")); - assert_eq!(messages[1]["content"], json!("It is sunny today")); - assert_eq!(messages[1]["tool_call_id"], json!("call-1")); -} - // ════════════════════════════════════════════════════════════════════════════ // Tool preparation — deepseek-prepare-tools.test.ts // diff --git a/aimux-providers/tests/fireworks_test.rs b/aimux-providers/tests/fireworks_test.rs index f927626f1..7261fd73c 100644 --- a/aimux-providers/tests/fireworks_test.rs +++ b/aimux-providers/tests/fireworks_test.rs @@ -28,20 +28,14 @@ use serial_test::serial; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_providers::{ProviderOptions, provider, provider_from_env}; fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn default_options(prompt: LanguageModelPrompt) -> CallOptions { diff --git a/aimux-providers/tests/google_files_test.rs b/aimux-providers/tests/google_files_test.rs index f8cf2234e..e18b5e667 100644 --- a/aimux-providers/tests/google_files_test.rs +++ b/aimux-providers/tests/google_files_test.rs @@ -614,3 +614,29 @@ async fn transient_upload_failure_is_retried_without_re_initiating() { ); assert_eq!(upload_attempts.load(std::sync::atomic::Ordering::SeqCst), 2); } + +// ════════════════════════════════════════════════════════════════════════════ +// get_model_path (TS: get-model-path.test.ts) +// ════════════════════════════════════════════════════════════════════════════ + +mod get_model_path_tests { + use aimux_providers::google::utils::get_model_path; + + #[test] + fn pass_through_for_models_slash() { + assert_eq!(get_model_path("models/some-model"), "models/some-model"); + } + + #[test] + fn pass_through_for_tuned_models_slash() { + assert_eq!( + get_model_path("tunedModels/some-model"), + "tunedModels/some-model" + ); + } + + #[test] + fn add_prefix_to_models_without_slash() { + assert_eq!(get_model_path("some-model"), "models/some-model"); + } +} diff --git a/aimux-providers/tests/google_remaining_test.rs b/aimux-providers/tests/google_remaining_test.rs deleted file mode 100644 index 30023ac97..000000000 --- a/aimux-providers/tests/google_remaining_test.rs +++ /dev/null @@ -1,1440 +0,0 @@ -//! Remaining Google provider tests — ported from the TS SDK test suite. -//! -//! Covers the pure-function and config tests that were not already exercised -//! by `google_model_test.rs` / `google_provider_tools_test.rs`: -//! -//! - get-model-path.test.ts -> get_model_path tests -//! - google-supported-file-url.test.ts -> is_supported_file_url tests -//! - google-model-capabilities.test.ts -> get_google_model_capabilities -//! - convert-json-schema-to-openapi-schema.test.ts -> schema conversion -//! - google-json-accumulator.test.ts -> GoogleJsonAccumulator -//! - google-provider.test.ts -> provider config (env var, headers, base URL) -//! -//! Reference: reference/ai/packages/google/src/. - -use serde_json::{Value, json}; -use serial_test::serial; -use wiremock::matchers::{header, method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::content::ContentPart; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; -use aimux_core::options::CallOptions; -use aimux_providers::google::convert::{ - convert_json_schema_to_openapi_schema, get_google_model_capabilities, -}; -use aimux_providers::google::utils::{ - GoogleJsonAccumulator, PartialArg, get_model_path, is_supported_file_url, -}; -use aimux_providers::{GoogleConfig, GoogleProvider}; - -// ════════════════════════════════════════════════════════════════════════════ -// get_model_path (TS: get-model-path.test.ts) -// ════════════════════════════════════════════════════════════════════════════ - -mod get_model_path_tests { - use super::*; - - #[test] - fn pass_through_for_models_slash() { - assert_eq!(get_model_path("models/some-model"), "models/some-model"); - } - - #[test] - fn pass_through_for_tuned_models_slash() { - assert_eq!( - get_model_path("tunedModels/some-model"), - "tunedModels/some-model" - ); - } - - #[test] - fn add_prefix_to_models_without_slash() { - assert_eq!(get_model_path("some-model"), "models/some-model"); - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// is_supported_file_url (TS: google-supported-file-url.test.ts) -// ════════════════════════════════════════════════════════════════════════════ - -mod is_supported_file_url_tests { - use super::*; - - #[test] - fn valid_google_file_urls() { - assert!(is_supported_file_url( - "https://generativelanguage.googleapis.com/v1beta/files/00000000-00000000-00000000-00000000" - )); - assert!(is_supported_file_url( - "https://generativelanguage.googleapis.com/v1beta/files/test123" - )); - } - - #[test] - fn valid_youtube_urls() { - let valid = [ - "https://www.youtube.com/watch?v=dQw4w9WgXcQ", - "https://youtube.com/watch?v=dQw4w9WgXcQ", - "https://youtu.be/dQw4w9WgXcQ", - "https://www.youtube.com/watch?v=dQw4w9WgXcQ&feature=youtu.be", - "https://youtu.be/dQw4w9WgXcQ?t=42", - ]; - for url in valid { - assert!(is_supported_file_url(url), "should be supported: {url}"); - } - } - - #[test] - fn invalid_youtube_urls() { - let invalid = [ - "https://youtube.com/channel/UCdQw4w9WgXcQ", - "https://youtube.com/playlist?list=PLdQw4w9WgXcQ", - "https://m.youtube.com/watch?v=dQw4w9WgXcQ", - "http://youtube.com/watch?v=dQw4w9WgXcQ", - "https://vimeo.com/123456789", - ]; - for url in invalid { - assert!( - !is_supported_file_url(url), - "should NOT be supported: {url}" - ); - } - } - - #[test] - fn non_google_file_urls() { - let cases = [ - "https://example.com", - "https://example.com/foo/bar", - "https://generativelanguage.googleapis.com", - "https://generativelanguage.googleapis.com/v1/other", - "http://generativelanguage.googleapis.com/v1beta/files/test", - "https://api.googleapis.com/v1beta/files/test", - ]; - for url in cases { - assert!( - !is_supported_file_url(url), - "should NOT be supported: {url}" - ); - } - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// get_google_model_capabilities (TS: google-model-capabilities.test.ts) -// ════════════════════════════════════════════════════════════════════════════ - -mod model_capabilities_tests { - use super::*; - - #[test] - fn gemini_pro_is_legacy() { - let caps = get_google_model_capabilities("gemini-pro"); - assert!(!caps.supports_gemini_2_tools); - assert!(!caps.supports_file_search); - assert!(!caps.uses_gemini_3_features); - } - - #[test] - fn gemini_pro_vision_is_legacy() { - let caps = get_google_model_capabilities("gemini-pro-vision"); - assert!(!caps.supports_gemini_2_tools); - assert!(!caps.supports_file_search); - assert!(!caps.uses_gemini_3_features); - } - - #[test] - fn gemini_1_5_flash_is_legacy() { - let caps = get_google_model_capabilities("gemini-1.5-flash"); - assert!(!caps.supports_gemini_2_tools); - assert!(!caps.supports_file_search); - assert!(!caps.uses_gemini_3_features); - } - - #[test] - fn gemini_robotics_er_1_5_is_legacy() { - let caps = get_google_model_capabilities("gemini-robotics-er-1.5-preview"); - assert!(!caps.supports_gemini_2_tools); - assert!(!caps.supports_file_search); - assert!(!caps.uses_gemini_3_features); - } - - #[test] - fn gemini_2_0_flash_supports_gemini_2_tools() { - let caps = get_google_model_capabilities("gemini-2.0-flash"); - assert!(caps.supports_gemini_2_tools); - assert!(!caps.supports_file_search); - assert!(!caps.uses_gemini_3_features); - } - - #[test] - fn gemini_2_5_flash_supports_file_search() { - let caps = get_google_model_capabilities("gemini-2.5-flash"); - assert!(caps.supports_gemini_2_tools); - assert!(caps.supports_file_search); - assert!(!caps.uses_gemini_3_features); - } - - #[test] - fn gemini_3_1_pro_uses_gemini_3_features() { - let caps = get_google_model_capabilities("gemini-3.1-pro-preview"); - assert!(caps.supports_gemini_2_tools); - assert!(caps.supports_file_search); - assert!(caps.uses_gemini_3_features); - } - - #[test] - fn unknown_future_gemini_inherits_newest() { - let caps = get_google_model_capabilities("gemini-99-pro-preview"); - assert!(caps.supports_gemini_2_tools); - assert!(caps.supports_file_search); - assert!(caps.uses_gemini_3_features); - } - - #[test] - fn gemini_ultra_latest_uses_gemini_3_features() { - let caps = get_google_model_capabilities("gemini-ultra-latest"); - assert!(caps.supports_gemini_2_tools); - assert!(caps.supports_file_search); - assert!(caps.uses_gemini_3_features); - } - - #[test] - fn nano_banana_supports_gemini_2_tools_only() { - let caps = get_google_model_capabilities("nano-banana-pro-preview"); - assert!(caps.supports_gemini_2_tools); - assert!(!caps.supports_file_search); - assert!(!caps.uses_gemini_3_features); - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// convert_json_schema_to_openapi_schema (TS: convert-json-schema-to-openapi-schema.test.ts) -// ════════════════════════════════════════════════════════════════════════════ - -mod schema_conversion_tests { - use super::*; - - fn convert_root(schema: &Value) -> Value { - convert_json_schema_to_openapi_schema(schema, true) - } - - fn convert_nested(schema: &Value) -> Value { - convert_json_schema_to_openapi_schema(schema, false) - } - - #[test] - fn removes_additional_properties_and_schema() { - let input = json!({ - "$schema": "http://json-schema.org/draft-07/schema#", - "type": "object", - "properties": { - "name": { "type": "string" }, - "age": { "type": "number" } - }, - "additionalProperties": false - }); - assert_eq!( - convert_root(&input), - json!({ - "type": "object", - "properties": { - "name": { "type": "string" }, - "age": { "type": "number" } - } - }) - ); - } - - #[test] - fn removes_additional_properties_from_nested_objects() { - let input = json!({ - "type": "object", - "properties": { - "keys": { - "type": "object", - "additionalProperties": { "type": "string" }, - "description": "Description for the key" - } - }, - "additionalProperties": false, - "$schema": "http://json-schema.org/draft-07/schema#" - }); - assert_eq!( - convert_root(&input), - json!({ - "type": "object", - "properties": { - "keys": { - "type": "object", - "description": "Description for the key" - } - } - }) - ); - } - - #[test] - fn handles_nested_objects_and_arrays() { - let input = json!({ - "type": "object", - "properties": { - "users": { - "type": "array", - "items": { - "type": "object", - "properties": { - "id": { "type": "number" }, - "name": { "type": "string" } - }, - "additionalProperties": false - } - } - }, - "additionalProperties": false - }); - assert_eq!( - convert_root(&input), - json!({ - "type": "object", - "properties": { - "users": { - "type": "array", - "items": { - "type": "object", - "properties": { - "id": { "type": "number" }, - "name": { "type": "string" } - } - } - } - } - }) - ); - } - - #[test] - fn converts_const_to_enum_single_value() { - let input = json!({ - "type": "object", - "properties": { "status": { "const": "active" } } - }); - assert_eq!( - convert_root(&input), - json!({ - "type": "object", - "properties": { "status": { "enum": ["active"] } } - }) - ); - } - - #[test] - fn handles_all_of_any_of_one_of() { - let input = json!({ - "type": "object", - "properties": { - "allOfProp": { "allOf": [{ "type": "string" }, { "minLength": 5 }] }, - "anyOfProp": { "anyOf": [{ "type": "string" }, { "type": "number" }] }, - "oneOfProp": { "oneOf": [{ "type": "boolean" }, { "type": "null" }] } - } - }); - assert_eq!( - convert_root(&input), - json!({ - "type": "object", - "properties": { - "allOfProp": { "allOf": [{ "type": "string" }, { "minLength": 5 }] }, - "anyOfProp": { "anyOf": [{ "type": "string" }, { "type": "number" }] }, - "oneOfProp": { "oneOf": [{ "type": "boolean" }, { "type": "null" }] } - } - }) - ); - } - - #[test] - fn preserves_format() { - let input = json!({ - "type": "object", - "properties": { "timestamp": { "type": "string", "format": "date-time" } } - }); - assert_eq!( - convert_root(&input), - json!({ - "type": "object", - "properties": { "timestamp": { "type": "string", "format": "date-time" } } - }) - ); - } - - #[test] - fn handles_required_properties() { - let input = json!({ - "type": "object", - "properties": { "id": { "type": "number" }, "name": { "type": "string" } }, - "required": ["id"] - }); - assert_eq!( - convert_root(&input), - json!({ - "type": "object", - "properties": { "id": { "type": "number" }, "name": { "type": "string" } }, - "required": ["id"] - }) - ); - } - - #[test] - fn deeply_nested_const_becomes_enum() { - let input = json!({ - "type": "object", - "properties": { - "nested": { - "type": "object", - "properties": { - "deeplyNested": { - "anyOf": [ - { "type": "object", "properties": { "value": { "const": "specific value" } } }, - { "type": "string" } - ] - } - } - } - } - }); - assert_eq!( - convert_root(&input), - json!({ - "type": "object", - "properties": { - "nested": { - "type": "object", - "properties": { - "deeplyNested": { - "anyOf": [ - { "type": "object", "properties": { "value": { "enum": ["specific value"] } } }, - { "type": "string" } - ] - } - } - } - } - }) - ); - } - - #[test] - fn complex_schema_with_nested_const_and_anyof() { - let input = json!({ - "type": "object", - "properties": { - "name": { "type": "string" }, - "age": { "type": "number" }, - "contact": { - "anyOf": [ - { - "type": "object", - "properties": { - "type": { "type": "string", "const": "email" }, - "value": { "type": "string" } - }, - "required": ["type", "value"], - "additionalProperties": false - }, - { - "type": "object", - "properties": { - "type": { "type": "string", "const": "phone" }, - "value": { "type": "string" } - }, - "required": ["type", "value"], - "additionalProperties": false - } - ] - } - }, - "required": ["name", "age", "contact"], - "additionalProperties": false, - "$schema": "http://json-schema.org/draft-07/schema#" - }); - let result = convert_root(&input); - assert_eq!(result["type"], "object"); - assert_eq!(result["required"], json!(["name", "age", "contact"])); - let contact_anyof = &result["properties"]["contact"]["anyOf"]; - assert_eq!( - contact_anyof[0]["properties"]["type"]["enum"], - json!(["email"]) - ); - assert_eq!( - contact_anyof[1]["properties"]["type"]["enum"], - json!(["phone"]) - ); - } - - #[test] - fn handles_null_type_correctly() { - let input = json!({ - "type": "object", - "properties": { - "nullableField": { "type": ["string", "null"] }, - "explicitNullField": { "type": "null" } - } - }); - let result = convert_root(&input); - assert_eq!( - result["properties"]["nullableField"], - json!({ "anyOf": [{ "type": "string" }], "nullable": true }) - ); - assert_eq!( - result["properties"]["explicitNullField"], - json!({ "type": "null" }) - ); - } - - #[test] - fn handles_descriptions() { - let input = json!({ - "type": "object", - "description": "A user object", - "properties": { - "id": { "type": "number", "description": "The user ID" }, - "name": { "type": "string", "description": "The user's full name" }, - "email": { "type": "string", "format": "email", "description": "The user's email address" } - }, - "required": ["id", "name"] - }); - assert_eq!(convert_root(&input), input); - } - - #[test] - fn returns_null_for_empty_object_schemas_at_root() { - let cases = [ - json!({ "type": "object" }), - json!({ "type": "object", "properties": {} }), - ]; - for schema in &cases { - assert!(convert_root(schema).is_null(), "expected null for {schema}"); - } - } - - #[test] - fn preserves_nested_empty_object_with_description() { - let input = json!({ - "type": "object", - "properties": { - "url": { "type": "string", "description": "URL to navigate to" }, - "launchOptions": { "type": "object", "description": "PuppeteerJS LaunchOptions" }, - "allowDangerous": { "type": "boolean", "description": "Allow dangerous options" } - }, - "required": ["url", "launchOptions"] - }); - assert_eq!(convert_root(&input), input); - } - - #[test] - fn preserves_nested_empty_object_without_description() { - let input = json!({ - "type": "object", - "properties": { "options": { "type": "object" } }, - "required": ["options"] - }); - assert_eq!(convert_root(&input), input); - } - - #[test] - fn handles_non_empty_object_schemas() { - let input = json!({ - "type": "object", - "properties": { "name": { "type": "string" } } - }); - assert_eq!(convert_root(&input), input); - } - - #[test] - fn converts_string_enum_properties() { - let input = json!({ - "type": "object", - "properties": { - "kind": { "type": "string", "enum": ["text", "code", "image"] } - }, - "required": ["kind"], - "additionalProperties": false, - "$schema": "http://json-schema.org/draft-07/schema#" - }); - assert_eq!( - convert_root(&input), - json!({ - "type": "object", - "properties": { - "kind": { "type": "string", "enum": ["text", "code", "image"] } - }, - "required": ["kind"] - }) - ); - } - - #[test] - fn converts_nullable_string_enum() { - let input = json!({ - "type": "object", - "properties": { - "fieldD": { - "anyOf": [ - { "type": "string", "enum": ["a", "b", "c"] }, - { "type": "null" } - ] - } - }, - "required": ["fieldD"], - "additionalProperties": false, - "$schema": "http://json-schema.org/draft-07/schema#" - }); - let result = convert_root(&input); - // The anyOf with a null branch collapses into nullable + the non-null - // schema's fields. - assert_eq!( - result["properties"]["fieldD"], - json!({ "nullable": true, "type": "string", "enum": ["a", "b", "c"] }) - ); - } - - #[test] - fn type_array_multiple_non_null_plus_null() { - let input = json!({ - "type": "object", - "properties": { - "multiTypeField": { "type": ["string", "number", "null"] } - } - }); - assert_eq!( - convert_root(&input), - json!({ - "type": "object", - "properties": { - "multiTypeField": { - "anyOf": [{ "type": "string" }, { "type": "number" }], - "nullable": true - } - } - }) - ); - } - - #[test] - fn type_array_without_null_becomes_anyof() { - let input = json!({ - "type": "object", - "properties": { - "multiTypeField": { "type": ["string", "number"] } - } - }); - assert_eq!( - convert_root(&input), - json!({ - "type": "object", - "properties": { - "multiTypeField": { "anyOf": [{ "type": "string" }, { "type": "number" }] } - } - }) - ); - } - - #[test] - fn boolean_schema_becomes_object() { - let input = json!(true); - assert_eq!( - convert_root(&input), - json!({ "type": "boolean", "properties": {} }) - ); - } - - #[test] - fn nested_empty_object_root_returns_null_nested_returns_type_object() { - let empty = json!({ "type": "object", "properties": {} }); - assert!(convert_root(&empty).is_null()); - // When nested, it becomes { "type": "object" }. - assert_eq!(convert_nested(&empty), json!({ "type": "object" })); - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// GoogleJsonAccumulator (TS: google-json-accumulator.test.ts) -// ════════════════════════════════════════════════════════════════════════════ - -mod json_accumulator_tests { - use super::*; - - fn s(path: &str, value: &str) -> PartialArg { - PartialArg { - json_path: path.to_string(), - string_value: Some(value.to_string()), - ..Default::default() - } - } - - fn s_continue(path: &str, value: &str) -> PartialArg { - PartialArg { - json_path: path.to_string(), - string_value: Some(value.to_string()), - will_continue: Some(true), - ..Default::default() - } - } - - fn n(path: &str, value: f64) -> PartialArg { - PartialArg { - json_path: path.to_string(), - number_value: Some(value), - ..Default::default() - } - } - - fn b(path: &str, value: bool) -> PartialArg { - PartialArg { - json_path: path.to_string(), - bool_value: Some(value), - ..Default::default() - } - } - - fn null_arg(path: &str) -> PartialArg { - PartialArg { - json_path: path.to_string(), - null_value: Some(()), - ..Default::default() - } - } - - // ── flat paths ───────────────────────────────────────────────────────── - - #[test] - fn accumulate_simple_string_with_will_continue() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc - .process_partial_args(&[s_continue("$.location", "Boston")]) - .unwrap(); - assert_eq!(r.current_json, json!({ "location": "Boston" })); - assert_eq!(r.text_delta, "{\"location\":\"Boston"); - } - - #[test] - fn continue_string_across_chunks() { - let mut acc = GoogleJsonAccumulator::new(); - acc.process_partial_args(&[s_continue("$.location", "Boston")]) - .unwrap(); - let r = acc - .process_partial_args(&[s("$.location", ", MA")]) - .unwrap(); - assert_eq!(r.current_json, json!({ "location": "Boston, MA" })); - assert_eq!(r.text_delta, ", MA"); - } - - #[test] - fn accumulate_complete_string() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc - .process_partial_args(&[s("$.location", "Boston")]) - .unwrap(); - assert_eq!(r.current_json, json!({ "location": "Boston" })); - assert_eq!(r.text_delta, "{\"location\":\"Boston\""); - } - - #[test] - fn accumulate_number() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc - .process_partial_args(&[n("$.brightness", 50.0)]) - .unwrap(); - assert_eq!(r.current_json, json!({ "brightness": 50 })); - assert_eq!(r.text_delta, "{\"brightness\":50"); - } - - #[test] - fn accumulate_boolean() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc.process_partial_args(&[b("$.enabled", true)]).unwrap(); - assert_eq!(r.current_json, json!({ "enabled": true })); - assert_eq!(r.text_delta, "{\"enabled\":true"); - } - - #[test] - fn accumulate_null() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc.process_partial_args(&[null_arg("$.nickname")]).unwrap(); - assert_eq!(r.current_json, json!({ "nickname": null })); - assert_eq!(r.text_delta, "{\"nickname\":null"); - } - - #[test] - fn accumulate_multiple_args_with_commas() { - let mut acc = GoogleJsonAccumulator::new(); - let first = acc - .process_partial_args(&[n("$.brightness", 50.0)]) - .unwrap(); - assert_eq!(first.text_delta, "{\"brightness\":50"); - - let second = acc.process_partial_args(&[b("$.enabled", true)]).unwrap(); - assert_eq!( - second.current_json, - json!({ "brightness": 50, "enabled": true }) - ); - assert_eq!(second.text_delta, ",\"enabled\":true"); - } - - #[test] - fn accumulate_multiple_args_single_call() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc - .process_partial_args(&[ - n("$.brightness", 50.0), - b("$.enabled", false), - null_arg("$.nickname"), - ]) - .unwrap(); - assert_eq!( - r.current_json, - json!({ "brightness": 50, "enabled": false, "nickname": null }) - ); - assert_eq!( - r.text_delta, - "{\"brightness\":50,\"enabled\":false,\"nickname\":null" - ); - } - - #[test] - fn escape_special_chars_in_continued_strings() { - let mut acc = GoogleJsonAccumulator::new(); - acc.process_partial_args(&[s_continue("$.query", "Boston \"Lo")]) - .unwrap(); - let r = acc.process_partial_args(&[s("$.query", "gan\"")]).unwrap(); - assert_eq!(r.current_json, json!({ "query": "Boston \"Logan\"" })); - // The continuation delta is the escaped inner content of `gan"`. - assert_eq!(r.text_delta, "gan\\\""); - } - - #[test] - fn skip_args_with_empty_path() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc.process_partial_args(&[s("$.", "ignored")]).unwrap(); - assert_eq!(r.current_json, json!({})); - assert_eq!(r.text_delta, ""); - } - - #[test] - fn skip_args_with_no_resolvable_value() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc - .process_partial_args(&[PartialArg { - json_path: "$.something".to_string(), - ..Default::default() - }]) - .unwrap(); - assert_eq!(r.current_json, json!({})); - assert_eq!(r.text_delta, ""); - } - - /// Regression (audit finding on H3): a malformed empty path like `$.[]` - /// parses to zero segments; it must be skipped instead of panicking with a - /// slice underflow in `emit_navigation_to`. - #[test] - fn malformed_empty_path_is_skipped_not_panicked() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc - .process_partial_args(&[ - PartialArg { - json_path: "$.[]".to_string(), - string_value: Some("x".to_string()), - ..Default::default() - }, - PartialArg { - json_path: "$.ok".to_string(), - string_value: Some("y".to_string()), - ..Default::default() - }, - ]) - .unwrap(); - // The malformed arg is ignored; the well-formed one still lands. - assert_eq!(r.current_json, json!({ "ok": "y" })); - } - - /// Regression (audit round 2): an oversized array index must produce an - /// error instead of expanding an array with ~1e6+ Null entries (OOM). - #[test] - fn oversized_array_index_is_rejected() { - let mut acc = GoogleJsonAccumulator::new(); - let err = acc - .process_partial_args(&[PartialArg { - json_path: "$.a[1000000]".to_string(), - string_value: Some("x".to_string()), - ..Default::default() - }]) - .unwrap_err(); - assert!(matches!(err, AiMuxError::JsonParse(_))); - } - - /// Regression (audit round 3, A1): an oversized **intermediate** index - /// (not the last segment) must also be rejected — previously only the - /// leaf segment was checked, so `$.a[1000000000].b` could still force - /// ~1 GiB of array allocation before failing. - #[test] - fn intermediate_oversized_index_is_rejected_atomically() { - let mut acc = GoogleJsonAccumulator::new(); - let err = acc - .process_partial_args(&[PartialArg { - json_path: "$.a[1000000].b".to_string(), - string_value: Some("x".to_string()), - ..Default::default() - }]) - .unwrap_err(); - assert!(matches!(err, AiMuxError::JsonParse(_))); - - // Error atomicity: the rejected arg must not leave partially-built - // containers (`a: []`) behind — a subsequent valid write lands cleanly. - let r = acc - .process_partial_args(&[PartialArg { - json_path: "$.ok".to_string(), - string_value: Some("y".to_string()), - ..Default::default() - }]) - .unwrap(); - assert_eq!(r.current_json, json!({ "ok": "y" })); - } - - /// Regression (audit round 2): an overly deep path must be rejected. - #[test] - fn overdeep_path_is_rejected() { - let mut acc = GoogleJsonAccumulator::new(); - let deep = format!("${}", ".a".repeat(100)); - let err = acc - .process_partial_args(&[PartialArg { - json_path: deep, - string_value: Some("x".to_string()), - ..Default::default() - }]) - .unwrap_err(); - assert!(matches!(err, AiMuxError::JsonParse(_))); - } - - #[test] - fn empty_partial_args_returns_empty_delta() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc.process_partial_args(&[]).unwrap(); - assert_eq!(r.current_json, json!({})); - assert_eq!(r.text_delta, ""); - } - - // ── nested paths ─────────────────────────────────────────────────────── - - #[test] - fn build_nested_object_from_dotted_path() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc - .process_partial_args(&[s("$.recipe.name", "Lasagna")]) - .unwrap(); - assert_eq!(r.current_json, json!({ "recipe": { "name": "Lasagna" } })); - assert_eq!(r.text_delta, "{\"recipe\":{\"name\":\"Lasagna\""); - } - - #[test] - fn build_nested_object_with_array_from_indexed_path() { - let mut acc = GoogleJsonAccumulator::new(); - let amount = acc - .process_partial_args(&[s("$.recipe.ingredients[0].amount", "16 oz")]) - .unwrap(); - assert_eq!( - amount.current_json, - json!({ "recipe": { "ingredients": [{ "amount": "16 oz" }] } }) - ); - assert_eq!( - amount.text_delta, - "{\"recipe\":{\"ingredients\":[{\"amount\":\"16 oz\"" - ); - - let name = acc - .process_partial_args(&[s("$.recipe.ingredients[0].name", "Lasagna noodles")]) - .unwrap(); - assert_eq!( - name.current_json, - json!({ "recipe": { "ingredients": [{ "amount": "16 oz", "name": "Lasagna noodles" }] } }) - ); - assert_eq!(name.text_delta, ",\"name\":\"Lasagna noodles\""); - } - - #[test] - fn accumulate_multiple_array_elements_across_chunks() { - let mut acc = GoogleJsonAccumulator::new(); - let mut deltas = Vec::new(); - - let r = acc - .process_partial_args(&[s("$.recipe.ingredients[0].amount", "16 oz")]) - .unwrap(); - deltas.push(r.text_delta.clone()); - - let r = acc - .process_partial_args(&[s("$.recipe.ingredients[0].name", "Noodles")]) - .unwrap(); - deltas.push(r.text_delta.clone()); - - let r = acc - .process_partial_args(&[s("$.recipe.ingredients[1].amount", "1 lb")]) - .unwrap(); - deltas.push(r.text_delta.clone()); - assert_eq!(r.text_delta, "},{\"amount\":\"1 lb\""); - - let r = acc - .process_partial_args(&[s("$.recipe.ingredients[1].name", "Beef")]) - .unwrap(); - deltas.push(r.text_delta.clone()); - assert_eq!(r.text_delta, ",\"name\":\"Beef\""); - - assert_eq!( - r.current_json, - json!({ - "recipe": { - "ingredients": [ - { "amount": "16 oz", "name": "Noodles" }, - { "amount": "1 lb", "name": "Beef" } - ] - } - }) - ); - - let f = acc.finalize(); - deltas.push(f.closing_delta); - assert_eq!(deltas.join(""), f.final_json); - } - - #[test] - fn string_continuation_on_nested_paths() { - let mut acc = GoogleJsonAccumulator::new(); - let start = acc - .process_partial_args(&[s_continue("$.recipe.steps[0]", "Preheat oven")]) - .unwrap(); - assert_eq!(start.text_delta, "{\"recipe\":{\"steps\":[\"Preheat oven"); - - let cont = acc - .process_partial_args(&[s("$.recipe.steps[0]", " to 375°F.")]) - .unwrap(); - assert_eq!( - cont.current_json, - json!({ "recipe": { "steps": ["Preheat oven to 375°F."] } }) - ); - assert_eq!(cont.text_delta, " to 375°F."); - } - - #[test] - fn mixed_nested_and_flat_paths() { - let mut acc = GoogleJsonAccumulator::new(); - let loc = acc - .process_partial_args(&[s("$.location", "Boston")]) - .unwrap(); - assert_eq!(loc.text_delta, "{\"location\":\"Boston\""); - - let details = acc - .process_partial_args(&[s("$.details.zip", "02101")]) - .unwrap(); - assert_eq!( - details.current_json, - json!({ "details": { "zip": "02101" }, "location": "Boston" }) - ); - assert_eq!(details.text_delta, ",\"details\":{\"zip\":\"02101\""); - - let f = acc.finalize(); - assert_eq!(f.closing_delta, "}}"); - assert_eq!( - f.final_json, - "{\"location\":\"Boston\",\"details\":{\"zip\":\"02101\"}}" - ); - } - - #[test] - fn array_elements_direct_string_values() { - let mut acc = GoogleJsonAccumulator::new(); - let first = acc - .process_partial_args(&[s("$.steps[0]", "Step one")]) - .unwrap(); - assert_eq!(first.text_delta, "{\"steps\":[\"Step one\""); - - let second = acc - .process_partial_args(&[s("$.steps[1]", "Step two")]) - .unwrap(); - assert_eq!( - second.current_json, - json!({ "steps": ["Step one", "Step two"] }) - ); - assert_eq!(second.text_delta, ",\"Step two\""); - } - - #[test] - fn deeply_nested_paths() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc.process_partial_args(&[s("$.a.b.c.d", "deep")]).unwrap(); - assert_eq!( - r.current_json, - json!({ "a": { "b": { "c": { "d": "deep" } } } }) - ); - assert_eq!(r.text_delta, "{\"a\":{\"b\":{\"c\":{\"d\":\"deep\""); - - let f = acc.finalize(); - assert_eq!(f.closing_delta, "}}}}"); - assert_eq!(f.final_json, "{\"a\":{\"b\":{\"c\":{\"d\":\"deep\"}}}}"); - } - - // ── finalize ─────────────────────────────────────────────────────────── - - #[test] - fn finalize_continued_string() { - let mut acc = GoogleJsonAccumulator::new(); - acc.process_partial_args(&[s_continue("$.location", "Boston")]) - .unwrap(); - let f = acc.finalize(); - assert_eq!(f.closing_delta, "\"}"); - assert_eq!(f.final_json, "{\"location\":\"Boston\"}"); - } - - #[test] - fn finalize_complete_string() { - let mut acc = GoogleJsonAccumulator::new(); - acc.process_partial_args(&[s("$.location", "Boston")]) - .unwrap(); - let f = acc.finalize(); - assert_eq!(f.closing_delta, "}"); - assert_eq!(f.final_json, "{\"location\":\"Boston\"}"); - } - - #[test] - fn finalize_multiple_args() { - let mut acc = GoogleJsonAccumulator::new(); - acc.process_partial_args(&[n("$.brightness", 50.0), b("$.enabled", true)]) - .unwrap(); - let f = acc.finalize(); - assert_eq!(f.closing_delta, "}"); - assert_eq!(f.final_json, "{\"brightness\":50,\"enabled\":true}"); - } - - #[test] - fn finalize_continued_string_with_continuation() { - let mut acc = GoogleJsonAccumulator::new(); - acc.process_partial_args(&[s_continue("$.location", "Boston")]) - .unwrap(); - acc.process_partial_args(&[s("$.location", ", MA")]) - .unwrap(); - let f = acc.finalize(); - assert_eq!(f.closing_delta, "\"}"); - assert_eq!(f.final_json, "{\"location\":\"Boston, MA\"}"); - } - - #[test] - fn finalize_empty_accumulator() { - let acc = GoogleJsonAccumulator::new(); - let f = acc.finalize(); - assert_eq!(f.closing_delta, "{}"); - assert_eq!(f.final_json, "{}"); - } - - #[test] - fn finalize_nested_structure() { - let mut acc = GoogleJsonAccumulator::new(); - acc.process_partial_args(&[s("$.recipe.ingredients[0].name", "Noodles")]) - .unwrap(); - acc.process_partial_args(&[s("$.recipe.name", "Lasagna")]) - .unwrap(); - let f = acc.finalize(); - let parsed: Value = serde_json::from_str(&f.final_json).unwrap(); - assert_eq!( - parsed, - json!({ - "recipe": { - "ingredients": [{ "name": "Noodles" }], - "name": "Lasagna" - } - }) - ); - } - - #[test] - fn finalize_nested_arrays_with_string_continuation() { - let mut acc = GoogleJsonAccumulator::new(); - acc.process_partial_args(&[s_continue("$.recipe.steps[0]", "Preheat")]) - .unwrap(); - acc.process_partial_args(&[s("$.recipe.steps[0]", " oven.")]) - .unwrap(); - acc.process_partial_args(&[s("$.recipe.steps[1]", "Cook.")]) - .unwrap(); - let f = acc.finalize(); - let parsed: Value = serde_json::from_str(&f.final_json).unwrap(); - assert_eq!( - parsed, - json!({ - "recipe": { "steps": ["Preheat oven.", "Cook."] } - }) - ); - } - - // ── concatenation invariant ──────────────────────────────────────────── - - #[test] - fn flat_args_concat_invariant() { - let mut acc = GoogleJsonAccumulator::new(); - let mut deltas = Vec::new(); - - let r = acc - .process_partial_args(&[n("$.brightness", 50.0)]) - .unwrap(); - deltas.push(r.text_delta); - - let r = acc.process_partial_args(&[b("$.enabled", true)]).unwrap(); - deltas.push(r.text_delta); - - let r = acc.process_partial_args(&[s("$.name", "test")]).unwrap(); - deltas.push(r.text_delta); - - let f = acc.finalize(); - deltas.push(f.closing_delta); - assert_eq!(deltas.join(""), f.final_json); - let parsed: Value = serde_json::from_str(&f.final_json).unwrap(); - assert_eq!( - parsed, - json!({ "brightness": 50, "enabled": true, "name": "test" }) - ); - } - - #[test] - fn nested_args_concat_invariant() { - let mut acc = GoogleJsonAccumulator::new(); - let mut deltas = Vec::new(); - - let r = acc - .process_partial_args(&[s("$.recipe.ingredients[0].amount", "16 oz")]) - .unwrap(); - deltas.push(r.text_delta); - let r = acc - .process_partial_args(&[s("$.recipe.ingredients[0].name", "Noodles")]) - .unwrap(); - deltas.push(r.text_delta); - let r = acc - .process_partial_args(&[s("$.recipe.ingredients[1].amount", "1 lb")]) - .unwrap(); - deltas.push(r.text_delta); - let r = acc - .process_partial_args(&[s("$.recipe.ingredients[1].name", "Beef")]) - .unwrap(); - deltas.push(r.text_delta); - let r = acc - .process_partial_args(&[s("$.recipe.name", "Lasagna")]) - .unwrap(); - deltas.push(r.text_delta); - let r = acc - .process_partial_args(&[s_continue("$.recipe.steps[0]", "Preheat")]) - .unwrap(); - deltas.push(r.text_delta); - let r = acc - .process_partial_args(&[s("$.recipe.steps[0]", " oven.")]) - .unwrap(); - deltas.push(r.text_delta); - let r = acc - .process_partial_args(&[s("$.recipe.steps[1]", "Cook.")]) - .unwrap(); - deltas.push(r.text_delta); - - let f = acc.finalize(); - deltas.push(f.closing_delta); - assert_eq!(deltas.join(""), f.final_json); - } - - #[test] - fn will_continue_strings_concat_invariant() { - let mut acc = GoogleJsonAccumulator::new(); - let mut deltas = Vec::new(); - - let r = acc - .process_partial_args(&[s_continue("$.location", "Bos")]) - .unwrap(); - deltas.push(r.text_delta); - let r = acc.process_partial_args(&[s("$.location", "ton")]).unwrap(); - deltas.push(r.text_delta); - let r = acc.process_partial_args(&[n("$.count", 42.0)]).unwrap(); - deltas.push(r.text_delta); - - let f = acc.finalize(); - deltas.push(f.closing_delta); - assert_eq!(deltas.join(""), f.final_json); - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// Provider config (TS: google-provider.test.ts — config subset) -// ════════════════════════════════════════════════════════════════════════════ - -mod provider_config_tests { - use super::*; - - fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] - } - - fn default_options(prompt: LanguageModelPrompt) -> CallOptions { - CallOptions::new(prompt) - } - - async fn mock_ok(server: &MockServer) { - Mock::given(method("POST")) - .and(path("/models/gemini-2.0-flash:generateContent")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "candidates": [{ - "content": { "parts": [{ "text": "hi" }], "role": "model" }, - "finishReason": "STOP", "index": 0 - }] - }))) - .mount(server) - .await; - } - - #[tokio::test] - async fn default_base_url_is_generativelanguage() { - let config = GoogleConfig::new("test-api-key"); - assert_eq!( - config.base_url, - "https://generativelanguage.googleapis.com/v1beta" - ); - } - - #[tokio::test] - async fn custom_base_url_is_used() { - let server = MockServer::start().await; - mock_ok(&server).await; - - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.model("gemini-2.0-flash"); - - model - .do_generate(&default_options(test_prompt())) - .await - .expect("ok"); - } - - #[tokio::test] - async fn provider_name_is_google() { - use aimux_core::provider::Provider; - let config = GoogleConfig::new("test-api-key"); - let provider = GoogleProvider::new(config); - assert_eq!(provider.name(), "google"); - } - - #[tokio::test] - async fn model_provider_is_google_generative_ai() { - let server = MockServer::start().await; - mock_ok(&server).await; - - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.model("gemini-2.0-flash"); - - assert_eq!(model.provider(), "google.generative-ai"); - assert_eq!(model.model_id(), "gemini-2.0-flash"); - } - - #[tokio::test] - async fn with_base_url_strips_trailing_slash() { - let config = GoogleConfig::new("key").with_base_url("https://example.com/v1beta/"); - assert_eq!(config.base_url, "https://example.com/v1beta"); - } - - #[tokio::test] - #[serial] - async fn from_env_uses_google_generative_ai_api_key() { - unsafe { - std::env::set_var("GOOGLE_GENERATIVE_AI_API_KEY", "env-api-key"); - } - let config = GoogleConfig::from_env().expect("from_env"); - assert_eq!(config.api_key, "env-api-key"); - unsafe { - std::env::remove_var("GOOGLE_GENERATIVE_AI_API_KEY"); - } - } - - #[tokio::test] - #[serial] - async fn from_env_errors_when_missing() { - unsafe { - std::env::remove_var("GOOGLE_GENERATIVE_AI_API_KEY"); - } - let result = GoogleConfig::from_env(); - assert!(result.is_err()); - } - - #[tokio::test] - async fn api_key_sent_via_x_goog_api_key_header() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/models/gemini-2.0-flash:generateContent")) - .and(header("x-goog-api-key", "my-secret-key")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "candidates": [{ - "content": { "parts": [{ "text": "hi" }], "role": "model" }, - "finishReason": "STOP", "index": 0 - }] - }))) - .mount(&server) - .await; - - let config = GoogleConfig::new("my-secret-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.model("gemini-2.0-flash"); - - model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed — the mock requires the x-goog-api-key header"); - } - - #[tokio::test] - async fn custom_headers_are_sent_alongside_api_key() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/models/gemini-2.0-flash:generateContent")) - .and(header("x-goog-api-key", "test-api-key")) - .and(header("custom-header", "custom-value")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "candidates": [{ - "content": { "parts": [{ "text": "hi" }], "role": "model" }, - "finishReason": "STOP", "index": 0 - }] - }))) - .mount(&server) - .await; - - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.model("gemini-2.0-flash"); - - let mut opts = default_options(test_prompt()); - let mut headers = std::collections::HashMap::new(); - headers.insert("Custom-Header".to_string(), "custom-value".to_string()); - opts.headers = Some(headers); - - model - .do_generate(&opts) - .await - .expect("should succeed — mock requires both headers"); - } - - #[tokio::test] - async fn language_model_trait_method_returns_boxed_model() { - use aimux_core::provider::Provider; - - let config = GoogleConfig::new("test-api-key"); - let provider = GoogleProvider::new(config); - let model = provider.language_model("gemini-2.0-flash").unwrap(); - assert_eq!(model.model_id(), "gemini-2.0-flash"); - assert_eq!(model.provider(), "google.generative-ai"); - } -} diff --git a/aimux-providers/tests/groq_test.rs b/aimux-providers/tests/groq_test.rs index 1a425cd1f..a206caf91 100644 --- a/aimux-providers/tests/groq_test.rs +++ b/aimux-providers/tests/groq_test.rs @@ -24,12 +24,14 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultOutput, ToolResultPart, UserPart, +}; use aimux_core::options::{CallOptions, ResponseFormat, Tool, ToolChoice}; use aimux_core::result::{GenerateContent, ReasoningOutput}; +use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; use aimux_core::tool::FunctionTool; use aimux_core::types::{FinishReasonUnified, ReasoningEffort}; @@ -40,11 +42,7 @@ use aimux_providers::{ProviderOptions, provider}; /// The TS `TEST_PROMPT`: a single user text message "Hello". fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } /// `CallOptions` with only `prompt` set. @@ -236,13 +234,22 @@ mod convert_messages { let model = make_provider(&server); - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::User, + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { content: vec![ - ContentPart::text("Hello"), - ContentPart::file_base64("AAECAw==", "image/png"), + UserPart::Text(TextPart { + text: "Hello".into(), + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64("AAECAw==".into()), + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + }), ], - ..Default::default() + provider_options: None, }]; model.do_generate(&default_options(prompt)).await.unwrap(); @@ -287,19 +294,27 @@ mod convert_messages { let model = make_provider(&server); let prompt: LanguageModelPrompt = vec![ - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::tool_call( - "quux", - "thwomp", - json!({"foo":"bar123"}), - )], - ..Default::default() + LanguageModelMessage::Assistant { + content: vec![AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "quux".into(), + tool_name: "thwomp".into(), + input: json!({"foo":"bar123"}), + provider_executed: None, + provider_options: None, + })], + provider_options: None, }, - LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result("quux", json!({"oof":"321rab"}))], - ..Default::default() + LanguageModelMessage::Tool { + content: vec![ToolPart::ToolResult(ToolResultPart { + tool_call_id: "quux".into(), + tool_name: "thwomp".into(), + output: ToolResultOutput::Json { + value: json!({"oof":"321rab"}), + provider_options: None, + }, + provider_options: None, + })], + provider_options: None, }, ]; model.do_generate(&default_options(prompt)).await.unwrap(); @@ -332,13 +347,21 @@ mod convert_messages { let model = make_provider(&server); - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::Assistant { content: vec![ - ContentPart::reasoning("I think the tool will return the correct value."), - ContentPart::tool_call("quux", "thwomp", json!({"foo":"bar123"})), + AssistantPart::Reasoning(ReasoningPart { + text: "I think the tool will return the correct value.".into(), + provider_options: None, + }), + AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "quux".into(), + tool_name: "thwomp".into(), + input: json!({"foo":"bar123"}), + provider_executed: None, + provider_options: None, + }), ], - ..Default::default() + provider_options: None, }]; model.do_generate(&default_options(prompt)).await.unwrap(); @@ -361,10 +384,12 @@ mod convert_messages { let model = make_provider(&server); - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::text("Hello, how can I help you?")], - ..Default::default() + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: "Hello, how can I help you?".into(), + provider_options: None, + })], + provider_options: None, }]; model.do_generate(&default_options(prompt)).await.unwrap(); @@ -384,22 +409,25 @@ mod convert_messages { let model = make_provider(&server); - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::file_reference( - "image/png", - json!({"groq": "file-ref-123"}), - )], - ..Default::default() + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Reference { + reference: [("groq".into(), "file-ref-123".into())].into(), + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + })], + provider_options: None, }]; - // Groq references don't contain an "openai" key, so the converter - // returns an InvalidArgument error instead of panicking. + // Upstream rejects file parts with provider references outright. let result = model.do_generate(&default_options(prompt)).await; assert!(result.is_err()); let err = result.unwrap_err(); assert!( - err.to_string().contains("No provider reference found"), - "error should mention 'No provider reference found': {err}" + err.to_string() + .contains("file parts with provider references"), + "error should mention 'file parts with provider references': {err}" ); } } diff --git a/aimux-providers/tests/huggingface_responses_test.rs b/aimux-providers/tests/huggingface_responses_test.rs index bd43b9875..81fda04de 100644 --- a/aimux-providers/tests/huggingface_responses_test.rs +++ b/aimux-providers/tests/huggingface_responses_test.rs @@ -23,14 +23,15 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; use aimux_core::generate::{GenerateTextOptions, generate_text}; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + FilePart, LanguageModelMessage, LanguageModelPrompt, UserPart, +}; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, Source}; +use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReasonUnified, ResponseMetadata}; @@ -43,11 +44,7 @@ use aimux_providers::{HuggingFaceConfig, HuggingFaceProvider}; /// The TS `TEST_PROMPT`: a single user text message "Hello". fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } /// `CallOptions` with only `prompt` set (everything else default/None). @@ -807,13 +804,16 @@ async fn should_throw_for_file_parts_with_provider_references() { let provider = make_provider(&server); let model = provider.responses_model("Qwen/Qwen2.5-VL-32B-Instruct"); - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::file_reference( - "image/jpeg", - json!({ "huggingface": "file-ref-123" }), - )], - ..Default::default() + let prompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Reference { + reference: [("huggingface".into(), "file-ref-123".into())].into(), + }, + media_type: "image/jpeg".into(), + filename: None, + provider_options: None, + })], + provider_options: None, }]; let result = model.do_generate(&default_options(prompt)).await; @@ -827,76 +827,6 @@ async fn should_throw_for_file_parts_with_provider_references() { ); } -/// TS: message conversion › "should warn about unsupported assistant content -/// types" — tool-call, tool-result, and reasoning in assistant messages -/// produce NO warnings (they are silently skipped). -#[tokio::test] -async fn should_not_warn_about_assistant_content_types() { - let server = MockServer::start().await; - stub_empty_response(&server).await; - - let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); - - let prompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ - ContentPart::tool_call("test", "test", json!({})), - ContentPart::tool_result("test", json!({ "type": "text", "value": "test" })), - ContentPart::reasoning("thinking..."), - ], - ..Default::default() - }]; - - let result = model - .do_generate(&default_options(prompt)) - .await - .expect("ok"); - - // TS expects no warnings (tool calls/results/reasoning are silently - // skipped, reasoning is included as output_text). - assert!( - result.warnings.is_empty(), - "expected no warnings, got: {:?}", - result.warnings - ); -} - -/// TS: message conversion › "should warn about tool messages" -#[tokio::test] -async fn should_warn_about_tool_messages() { - let server = MockServer::start().await; - stub_empty_response(&server).await; - - let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); - - let prompt = vec![LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result( - "test", - json!({ "type": "text", "value": "test" }), - )], - ..Default::default() - }]; - - let result = model - .do_generate(&default_options(prompt)) - .await - .expect("ok"); - - let features: Vec = result - .warnings - .iter() - .filter_map(|w| match w { - aimux_core::types::Warning::Unsupported { feature, .. } => Some(feature.clone()), - _ => None, - }) - .collect(); - - assert_eq!(features, vec!["tool messages"]); -} - // ════════════════════════════════════════════════════════════════════════════ // Tool calls // ════════════════════════════════════════════════════════════════════════════ @@ -1103,10 +1033,16 @@ const PNG_BASE64: &str = "iVBORw0KGgo="; /// TS: "passes full image/png through unchanged for inline data" #[test] fn passes_full_image_png_through_unchanged_for_inline_data() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::file_base64(PNG_BASE64, "image/png")], - ..Default::default() + let prompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64(PNG_BASE64.into()), + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + })], + provider_options: None, }]; let (input, warnings) = @@ -1126,10 +1062,16 @@ fn passes_full_image_png_through_unchanged_for_inline_data() { /// TS: "detects image subtype from inline bytes for top-level 'image'" #[test] fn detects_image_subtype_from_inline_bytes_for_top_level_image() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::file_base64(PNG_BASE64, "image")], - ..Default::default() + let prompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64(PNG_BASE64.into()), + }, + media_type: "image".into(), + filename: None, + provider_options: None, + })], + provider_options: None, }]; let (input, warnings) = @@ -1146,36 +1088,19 @@ fn detects_image_subtype_from_inline_bytes_for_top_level_image() { ); } -/// TS: "passes through URL source for top-level-only image" -#[test] -fn passes_through_url_source_for_top_level_only_image() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::file_url("https://example.com/x.png", "image")], - ..Default::default() - }]; - - let (input, warnings) = - convert_to_huggingface_responses_messages(&prompt).expect("should succeed"); - - assert!(warnings.is_empty()); - let content = &input[0]["content"][0]; - assert_eq!( - content, - &json!({ - "type": "input_image", - "image_url": "https://example.com/x.png" - }) - ); -} - /// TS: "normalizes image/* wildcard via detection" #[test] fn normalizes_image_wildcard_via_detection() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::file_base64(PNG_BASE64, "image/*")], - ..Default::default() + let prompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64(PNG_BASE64.into()), + }, + media_type: "image/*".into(), + filename: None, + provider_options: None, + })], + provider_options: None, }]; let (input, warnings) = diff --git a/aimux-providers/tests/mistral_embedding_test.rs b/aimux-providers/tests/mistral_embedding_test.rs index 7d5940758..66eed80be 100644 --- a/aimux-providers/tests/mistral_embedding_test.rs +++ b/aimux-providers/tests/mistral_embedding_test.rs @@ -10,6 +10,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::tool::FunctionTool; +use aimux_providers::mistral::convert::prepare_tools; use aimux_providers::{MistralConfig, MistralProvider}; const TEST_VALUES: &[&str] = &["sunny day at the beach", "rainy day in the city"]; @@ -211,3 +213,13 @@ async fn should_pass_headers() { Some("request-header-value") ); } + +/// TS: "should include strict field when set" +#[test] +fn prepare_tools_function_tool_with_strict() { + let tool = FunctionTool::new("strictTool", json!({ "type": "object", "properties": {} })) + .with_strict(true); + let result = prepare_tools(&Some(vec![tool]), None); + let tools = result.tools.expect("tools should be Some"); + assert_eq!(tools[0]["function"]["strict"], json!(true)); +} diff --git a/aimux-providers/tests/mistral_remaining_test.rs b/aimux-providers/tests/mistral_remaining_test.rs deleted file mode 100644 index b5a66188d..000000000 --- a/aimux-providers/tests/mistral_remaining_test.rs +++ /dev/null @@ -1,892 +0,0 @@ -//! Remaining Mistral tests translated from TS that are NOT already covered by -//! `mistral_model_test.rs`. Covers: -//! -//! - `mistral-prepare-tools.test.ts` — all 8 cases (prepare_tools unit tests) -//! - `convert-to-mistral-chat-messages.test.ts` — 12 cases (message conversion -//! unit tests): system, user text/image/file, assistant text/tool-call, -//! prefix continuation, tool results -//! - `convert-mistral-usage.test.ts` — 4 cases (usage conversion via doGenerate): -//! alternate cached-token spellings, zero cached, missing tokens -//! - `mistral-chat-language-model.test.ts` — uncovered scenarios: -//! parse_finish_reason unit tests (6), build_request_body unit tests (7): -//! stream flag, provider tool silently dropped, tool with description/strict, -//! json_schema with name+description - -use base64::Engine; -use serde_json::{Value, json}; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::content::ContentPart; -use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; -use aimux_core::options::{CallOptions, ResponseFormat, Tool, ToolChoice}; -use aimux_core::tool::FunctionTool; -use aimux_core::types::FinishReasonUnified; - -use aimux_providers::mistral::convert::{ - build_request_body, convert_prompt_to_mistral_messages, parse_finish_reason, prepare_tools, -}; -use aimux_providers::{MistralConfig, MistralProvider}; - -// ── helpers ───────────────────────────────────────────────────────────────── - -fn b64(data: &[u8]) -> String { - base64::engine::general_purpose::STANDARD.encode(data) -} - -fn default_options(prompt: LanguageModelPrompt) -> CallOptions { - CallOptions::new(prompt) -} - -fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] -} - -fn basic_function_tool() -> FunctionTool { - FunctionTool::new( - "testFunction", - json!({ "type": "object", "properties": {} }), - ) - .with_description("test description") -} - -fn test_tool() -> Tool { - Tool::Function(FunctionTool::new( - "test-tool", - json!({ - "type": "object", - "properties": { "value": { "type": "string" } }, - "required": ["value"], - "additionalProperties": false, - "$schema": "http://json-schema.org/draft-07/schema#" - }), - )) -} - -async fn mock_json_response(server: &MockServer, body: Value) { - Mock::given(method("POST")) - .and(path("/chat/completions")) - .respond_with(ResponseTemplate::new(200).set_body_json(body)) - .mount(server) - .await; -} - -// ════════════════════════════════════════════════════════════════════════════ -// mistral-prepare-tools.test.ts -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: "should return undefined tools when no tools are provided" -#[test] -fn prepare_tools_empty_returns_none() { - let result = prepare_tools(&Some(vec![]), None); - assert!(result.tools.is_none()); - assert!(result.tool_choice.is_none()); - assert!(result.tool_warnings.is_empty()); -} - -/// TS: "should process function tools correctly" -#[test] -fn prepare_tools_function_tool() { - let result = prepare_tools(&Some(vec![basic_function_tool()]), None); - assert_eq!( - result.tools, - Some(vec![json!({ - "type": "function", - "function": { - "name": "testFunction", - "description": "test description", - "parameters": { "type": "object", "properties": {} } - } - })]) - ); - assert!(result.tool_choice.is_none()); - assert!(result.tool_warnings.is_empty()); -} - -/// TS: "should include strict field when set" -#[test] -fn prepare_tools_function_tool_with_strict() { - let tool = FunctionTool::new("strictTool", json!({ "type": "object", "properties": {} })) - .with_strict(true); - let result = prepare_tools(&Some(vec![tool]), None); - let tools = result.tools.expect("tools should be Some"); - assert_eq!(tools[0]["function"]["strict"], json!(true)); -} - -/// TS: "should handle auto tool choice" -#[test] -fn prepare_tools_auto_choice() { - let result = prepare_tools(&Some(vec![basic_function_tool()]), Some(&ToolChoice::Auto)); - assert_eq!(result.tool_choice, Some(json!("auto"))); - assert!(result.tools.is_some()); -} - -/// TS: "should handle none tool choice" -#[test] -fn prepare_tools_none_choice() { - let result = prepare_tools(&Some(vec![basic_function_tool()]), Some(&ToolChoice::None)); - assert_eq!(result.tool_choice, Some(json!("none"))); - assert!(result.tools.is_some()); -} - -/// TS: "should handle required tool choice" — Mistral maps Required → "any". -#[test] -fn prepare_tools_required_choice() { - let result = prepare_tools( - &Some(vec![basic_function_tool()]), - Some(&ToolChoice::Required), - ); - assert_eq!(result.tool_choice, Some(json!("any"))); - assert!(result.tools.is_some()); -} - -/// TS: "should handle tool type tool choice by filtering tools" — uses "any". -#[test] -fn prepare_tools_tool_choice_filters() { - let result = prepare_tools( - &Some(vec![basic_function_tool()]), - Some(&ToolChoice::Tool { - tool_name: "testFunction".to_string(), - }), - ); - assert_eq!(result.tool_choice, Some(json!("any"))); - let tools = result.tools.expect("tools should be Some"); - assert_eq!(tools.len(), 1); - assert_eq!(tools[0]["function"]["name"], json!("testFunction")); -} - -/// TS: "should return no tools when tool_choice references a non-existent tool" -#[test] -fn prepare_tools_tool_choice_filters_to_empty() { - let result = prepare_tools( - &Some(vec![basic_function_tool()]), - Some(&ToolChoice::Tool { - tool_name: "nonExistent".to_string(), - }), - ); - assert_eq!(result.tool_choice, Some(json!("any"))); - let tools = result.tools.expect("tools should be Some"); - assert!(tools.is_empty()); -} - -// ════════════════════════════════════════════════════════════════════════════ -// convert-to-mistral-chat-messages.test.ts — text messages -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: "should convert a system message to a string content" -#[test] -fn convert_system_message() { - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text("You are a helpful assistant.")], - ..Default::default() - }]; - let messages = convert_prompt_to_mistral_messages(&prompt); - assert_eq!( - messages, - vec![json!({ "role": "system", "content": "You are a helpful assistant." })] - ); -} - -/// TS: "should convert a user message to an array content" -#[test] -fn convert_user_text_message() { - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }]; - let messages = convert_prompt_to_mistral_messages(&prompt); - assert_eq!( - messages, - vec![json!({ - "role": "user", - "content": [{ "type": "text", "text": "Hello" }] - })] - ); -} - -/// TS: "should convert an assistant message to a string content" -#[test] -fn convert_assistant_message() { - let prompt: LanguageModelPrompt = vec![ - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hi")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::text("Hello!")], - ..Default::default() - }, - ]; - let messages = convert_prompt_to_mistral_messages(&prompt); - // The assistant message is last → prefix: true. - let last = messages.last().expect("last message"); - assert_eq!(last["role"], json!("assistant")); - assert_eq!(last["content"], json!("Hello!")); - assert_eq!(last["prefix"], json!(true)); -} - -/// TS: "should not set prefix on a non-last assistant message" -#[test] -fn convert_assistant_message_no_prefix_when_not_last() { - let prompt: LanguageModelPrompt = vec![ - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::text("I am thinking")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Continue")], - ..Default::default() - }, - ]; - let messages = convert_prompt_to_mistral_messages(&prompt); - let assistant_msg = &messages[0]; - assert_eq!(assistant_msg["role"], json!("assistant")); - assert_eq!(assistant_msg["content"], json!("I am thinking")); - assert!( - assistant_msg.get("prefix").is_none(), - "prefix should be absent on non-last assistant" - ); -} - -// ════════════════════════════════════════════════════════════════════════════ -// convert-to-mistral-chat-messages.test.ts — image / file processing -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: "should convert an image part to an image_url data URI" -#[test] -fn convert_user_image_message() { - let data = vec![0, 1, 2, 3]; - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ - ContentPart::text("What is this?"), - ContentPart::Image { - image: data.clone(), - media_type: "image/png".to_string(), - provider_options: None, - }, - ], - ..Default::default() - }]; - let messages = convert_prompt_to_mistral_messages(&prompt); - let expected_url = format!("data:image/png;base64,{}", b64(&data)); - assert_eq!( - messages, - vec![json!({ - "role": "user", - "content": [ - { "type": "text", "text": "What is this?" }, - { "type": "image_url", "image_url": expected_url } - ] - })] - ); -} - -/// TS: "should convert an image file to an image_url data URI" -#[test] -fn convert_user_image_file_message() { - let data = vec![0, 1, 2, 3]; - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::File { - data: data.clone(), - media_type: "image/png".to_string(), - filename: None, - provider_options: None, - }], - ..Default::default() - }]; - let messages = convert_prompt_to_mistral_messages(&prompt); - let expected_url = format!("data:image/png;base64,{}", b64(&data)); - assert_eq!( - messages, - vec![json!({ - "role": "user", - "content": [{ "type": "image_url", "image_url": expected_url }] - })] - ); -} - -/// TS: "should convert a non-image file to a document_url data URI" -#[test] -fn convert_user_document_file_message() { - let data = b"Document content".to_vec(); - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::File { - data: data.clone(), - media_type: "text/plain".to_string(), - filename: Some("doc.txt".to_string()), - provider_options: None, - }], - ..Default::default() - }]; - let messages = convert_prompt_to_mistral_messages(&prompt); - let expected_url = format!("data:text/plain;base64,{}", b64(&data)); - assert_eq!( - messages, - vec![json!({ - "role": "user", - "content": [{ "type": "document_url", "document_url": expected_url }] - })] - ); -} - -/// TS: "should accept top-level 'image' media type and still route to image_url" -#[test] -fn convert_user_image_file_top_level_mediatype() { - let data = vec![0x89, 0x50, 0x4e, 0x47]; - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::File { - data: data.clone(), - media_type: "image".to_string(), - filename: None, - provider_options: None, - }], - ..Default::default() - }]; - let messages = convert_prompt_to_mistral_messages(&prompt); - let content = messages[0]["content"].as_array().expect("content array"); - assert_eq!(content[0]["type"], json!("image_url")); - let url = content[0]["image_url"].as_str().unwrap(); - assert!( - url.starts_with("data:image;base64,"), - "url should start with data:image: {url}" - ); -} - -// ════════════════════════════════════════════════════════════════════════════ -// convert-to-mistral-chat-messages.test.ts — tool messages -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: "should convert an assistant tool call message" -#[test] -fn convert_assistant_tool_call_message() { - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ - ContentPart::text("Calling a tool"), - ContentPart::ToolCall { - tool_call_id: "tool-call-1".to_string(), - tool_name: "tool-1".to_string(), - input: json!({ "test": "This is a tool message" }), - provider_executed: None, - provider_options: None, - }, - ], - ..Default::default() - }]; - let messages = convert_prompt_to_mistral_messages(&prompt); - // The assistant message is last → prefix: true. - assert_eq!( - messages, - vec![json!({ - "role": "assistant", - "content": "Calling a tool", - "tool_calls": [{ - "id": "tool-call-1", - "type": "function", - "function": { - "name": "tool-1", - "arguments": "{\"test\":\"This is a tool message\"}" - } - }], - "prefix": true - })] - ); -} - -/// TS: "should convert a tool result message" -#[test] -fn convert_tool_result_message() { - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::ToolResult { - tool_call_id: "tool-call-1".to_string(), - result: json!({ "test": "This is a tool message" }), - tool_name: None, - is_error: None, - preliminary: None, - dynamic: None, - provider_options: None, - }], - ..Default::default() - }]; - let messages = convert_prompt_to_mistral_messages(&prompt); - assert_eq!( - messages, - vec![json!({ - "role": "tool", - "tool_call_id": "tool-call-1", - "content": "{\"test\":\"This is a tool message\"}" - })] - ); -} - -/// TS: "should convert a tool result with a string output" -#[test] -fn convert_tool_result_with_string_output() { - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::ToolResult { - tool_call_id: "tool-call-1".to_string(), - result: json!("plain string result"), - tool_name: None, - is_error: None, - preliminary: None, - dynamic: None, - provider_options: None, - }], - ..Default::default() - }]; - let messages = convert_prompt_to_mistral_messages(&prompt); - assert_eq!( - messages, - vec![json!({ - "role": "tool", - "tool_call_id": "tool-call-1", - "content": "plain string result" - })] - ); -} - -/// TS: "should convert multiple tool results into separate tool messages" -#[test] -fn convert_multiple_tool_results() { - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ - ContentPart::ToolResult { - tool_call_id: "tool-call-1".to_string(), - result: json!({ "test": "result 1" }), - tool_name: None, - is_error: None, - preliminary: None, - dynamic: None, - provider_options: None, - }, - ContentPart::ToolResult { - tool_call_id: "tool-call-2".to_string(), - result: json!({ "test": "result 2" }), - tool_name: None, - is_error: None, - preliminary: None, - dynamic: None, - provider_options: None, - }, - ], - ..Default::default() - }]; - let messages = convert_prompt_to_mistral_messages(&prompt); - assert_eq!(messages.len(), 2); - assert_eq!(messages[0]["role"], json!("tool")); - assert_eq!(messages[0]["tool_call_id"], json!("tool-call-1")); - assert_eq!(messages[0]["content"], json!("{\"test\":\"result 1\"}")); - assert_eq!(messages[1]["role"], json!("tool")); - assert_eq!(messages[1]["tool_call_id"], json!("tool-call-2")); - assert_eq!(messages[1]["content"], json!("{\"test\":\"result 2\"}")); -} - -// ════════════════════════════════════════════════════════════════════════════ -// convert-mistral-usage.test.ts — usage conversion via doGenerate -// -// `convert_usage` is private in model.rs, so these tests exercise it through -// `do_generate` and assert on `result.usage`. -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: "should extract usage with prompt_tokens_details.cached_tokens" -#[tokio::test] -async fn usage_with_prompt_tokens_details_cached() { - let server = MockServer::start().await; - mock_json_response( - &server, - json!({ - "id": "test-id", - "model": "mistral-small-latest", - "usage": { - "prompt_tokens": 100, - "total_tokens": 200, - "completion_tokens": 100, - "prompt_tokens_details": { "cached_tokens": 25 } - }, - "object": "chat.completion", - "choices": [{ - "index": 0, - "finish_reason": "stop", - "message": { "role": "assistant", "content": "ok" } - }] - }), - ) - .await; - - let config = MistralConfig::new("test-api-key").with_base_url(server.uri()); - let provider = MistralProvider::new(config); - let model = provider.model("mistral-small-latest"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - assert_eq!(result.usage.input_tokens.total, Some(100)); - assert_eq!(result.usage.input_tokens.no_cache, Some(75)); - assert_eq!(result.usage.input_tokens.cache_read, Some(25)); - assert_eq!(result.usage.output_tokens.total, Some(100)); -} - -/// TS: "should extract usage with prompt_token_details.cached_tokens" — the -/// alternate (singular "token") spelling used by Mistral at different times. -#[tokio::test] -async fn usage_with_prompt_token_details_cached() { - let server = MockServer::start().await; - mock_json_response( - &server, - json!({ - "id": "test-id", - "model": "mistral-small-latest", - "usage": { - "prompt_tokens": 100, - "total_tokens": 200, - "completion_tokens": 100, - "prompt_token_details": { "cached_tokens": 40 } - }, - "object": "chat.completion", - "choices": [{ - "index": 0, - "finish_reason": "stop", - "message": { "role": "assistant", "content": "ok" } - }] - }), - ) - .await; - - let config = MistralConfig::new("test-api-key").with_base_url(server.uri()); - let provider = MistralProvider::new(config); - let model = provider.model("mistral-small-latest"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - assert_eq!(result.usage.input_tokens.total, Some(100)); - assert_eq!(result.usage.input_tokens.no_cache, Some(60)); - assert_eq!(result.usage.input_tokens.cache_read, Some(40)); -} - -/// TS: "should not set cache_read when cached_tokens is 0" — the impl omits -/// `cache_read` when the value is 0 (uses `if cache_read > 0`). -#[tokio::test] -async fn usage_with_zero_cached_tokens() { - let server = MockServer::start().await; - mock_json_response( - &server, - json!({ - "id": "test-id", - "model": "mistral-small-latest", - "usage": { - "prompt_tokens": 50, - "total_tokens": 100, - "completion_tokens": 50, - "num_cached_tokens": 0 - }, - "object": "chat.completion", - "choices": [{ - "index": 0, - "finish_reason": "stop", - "message": { "role": "assistant", "content": "ok" } - }] - }), - ) - .await; - - let config = MistralConfig::new("test-api-key").with_base_url(server.uri()); - let provider = MistralProvider::new(config); - let model = provider.model("mistral-small-latest"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - assert_eq!(result.usage.input_tokens.total, Some(50)); - assert_eq!(result.usage.input_tokens.no_cache, Some(50)); - assert_eq!(result.usage.input_tokens.cache_read, None); -} - -/// TS: "should default missing token counts to 0" -#[tokio::test] -async fn usage_with_missing_tokens_defaults_to_zero() { - let server = MockServer::start().await; - mock_json_response( - &server, - json!({ - "id": "test-id", - "model": "mistral-small-latest", - "usage": {}, - "object": "chat.completion", - "choices": [{ - "index": 0, - "finish_reason": "stop", - "message": { "role": "assistant", "content": "ok" } - }] - }), - ) - .await; - - let config = MistralConfig::new("test-api-key").with_base_url(server.uri()); - let provider = MistralProvider::new(config); - let model = provider.model("mistral-small-latest"); - - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - assert_eq!(result.usage.input_tokens.total, Some(0)); - assert_eq!(result.usage.input_tokens.no_cache, Some(0)); - assert_eq!(result.usage.input_tokens.cache_read, None); - assert_eq!(result.usage.output_tokens.total, Some(0)); -} - -// ════════════════════════════════════════════════════════════════════════════ -// mistral-chat-language-model.test.ts — parse_finish_reason unit tests -// ════════════════════════════════════════════════════════════════════════════ - -#[test] -fn parse_finish_reason_stop() { - let fr = parse_finish_reason("stop"); - assert_eq!(fr.unified, FinishReasonUnified::Stop); - assert_eq!(fr.raw.as_deref(), Some("stop")); -} - -#[test] -fn parse_finish_reason_length() { - let fr = parse_finish_reason("length"); - assert_eq!(fr.unified, FinishReasonUnified::Length); - assert_eq!(fr.raw.as_deref(), Some("length")); -} - -/// TS: "model_length" is a Mistral-specific finish reason also mapped to Length. -#[test] -fn parse_finish_reason_model_length() { - let fr = parse_finish_reason("model_length"); - assert_eq!(fr.unified, FinishReasonUnified::Length); - assert_eq!(fr.raw.as_deref(), Some("model_length")); -} - -#[test] -fn parse_finish_reason_tool_calls() { - let fr = parse_finish_reason("tool_calls"); - assert_eq!(fr.unified, FinishReasonUnified::ToolCalls); -} - -#[test] -fn parse_finish_reason_content_filter() { - let fr = parse_finish_reason("content_filter"); - assert_eq!(fr.unified, FinishReasonUnified::ContentFilter); -} - -#[test] -fn parse_finish_reason_unknown() { - let fr = parse_finish_reason("something_unexpected"); - assert_eq!(fr.unified, FinishReasonUnified::Other); - assert_eq!(fr.raw.as_deref(), Some("something_unexpected")); -} - -// ════════════════════════════════════════════════════════════════════════════ -// build_request_body unit tests — stream flag, provider tools, tool fields, -// response format with name + description. -// ════════════════════════════════════════════════════════════════════════════ - -#[test] -fn build_request_body_stream_flag() { - let options = default_options(test_prompt()); - let body = build_request_body("mistral-small-latest", &options, true).unwrap(); - assert_eq!(body["stream"], json!(true)); -} - -#[test] -fn build_request_body_no_stream_flag() { - let options = default_options(test_prompt()); - let body = build_request_body("mistral-small-latest", &options, false).unwrap(); - assert!(body.get("stream").is_none()); -} - -/// TS: a function tool with a description includes it in the request body. -#[test] -fn build_request_body_tool_with_description() { - let tool = FunctionTool::new("weather", json!({ "type": "object", "properties": {} })) - .with_description("Get the weather"); - let options = CallOptions { - prompt: test_prompt(), - tools: Some(vec![Tool::Function(tool)]), - ..default_options(Vec::new()) - }; - let body = build_request_body("mistral-small-latest", &options, false).unwrap(); - assert_eq!(body["tools"][0]["function"]["name"], json!("weather")); - assert_eq!( - body["tools"][0]["function"]["description"], - json!("Get the weather") - ); -} - -/// TS: a function tool with strict=true includes it in the request body. -#[test] -fn build_request_body_tool_with_strict() { - let tool = FunctionTool::new("weather", json!({ "type": "object", "properties": {} })) - .with_strict(true); - let options = CallOptions { - prompt: test_prompt(), - tools: Some(vec![Tool::Function(tool)]), - ..default_options(Vec::new()) - }; - let body = build_request_body("mistral-small-latest", &options, false).unwrap(); - assert_eq!(body["tools"][0]["function"]["strict"], json!(true)); -} - -/// TS: json_schema response format with name + description + strict=false. -#[test] -fn build_request_body_json_schema_with_name_and_description() { - let schema = json!({ "type": "object", "properties": { "name": { "type": "string" } } }); - let options = CallOptions { - prompt: test_prompt(), - response_format: Some(ResponseFormat::Json { - schema: Some(schema.clone()), - name: Some("person".to_string()), - description: Some("A person object".to_string()), - }), - ..default_options(Vec::new()) - }; - let body = build_request_body("mistral-small-latest", &options, false).unwrap(); - assert_eq!(body["response_format"]["type"], json!("json_schema")); - assert_eq!(body["response_format"]["json_schema"]["schema"], schema); - assert_eq!( - body["response_format"]["json_schema"]["name"], - json!("person") - ); - assert_eq!( - body["response_format"]["json_schema"]["description"], - json!("A person object") - ); - assert_eq!( - body["response_format"]["json_schema"]["strict"], - json!(false) - ); -} - -/// TS: json_object response format when schema is None. -#[test] -fn build_request_body_json_object_when_no_schema() { - let options = CallOptions { - prompt: test_prompt(), - response_format: Some(ResponseFormat::Json { - schema: None, - name: Some("ignored".to_string()), - description: None, - }), - ..default_options(Vec::new()) - }; - let body = build_request_body("mistral-small-latest", &options, false).unwrap(); - assert_eq!(body["response_format"], json!({ "type": "json_object" })); -} - -// ════════════════════════════════════════════════════════════════════════════ -// doGenerate — request body via the model (verifies build_request_body wiring) -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: "should pass tools and toolChoice" — Mistral maps `required` → `"any"`. -#[tokio::test] -async fn should_pass_tools_with_tool_choice_required() { - let server = MockServer::start().await; - mock_json_response( - &server, - json!({ - "id": "test-id", - "model": "test-model", - "usage": { "prompt_tokens": 4, "total_tokens": 34, "completion_tokens": 30 }, - "object": "chat.completion", - "choices": [{ - "index": 0, - "finish_reason": "stop", - "message": { "role": "assistant", "content": "ok" } - }] - }), - ) - .await; - - let config = MistralConfig::new("test-api-key").with_base_url(server.uri()); - let provider = MistralProvider::new(config); - let model = provider.model("test-model"); - - let options = CallOptions { - prompt: test_prompt(), - tools: Some(vec![test_tool()]), - tool_choice: Some(ToolChoice::Required), - ..default_options(Vec::new()) - }; - - let result = model.do_generate(&options).await.expect("should succeed"); - let body = result.request.expect("request").body.expect("body"); - - assert_eq!( - body["tool_choice"], - json!("any"), - "Mistral uses 'any' for required tool choice" - ); - assert_eq!(body["tools"][0]["function"]["name"], "test-tool"); -} - -/// TS: "should pass tools and toolChoice" — `tool` choice filters tools. -#[tokio::test] -async fn should_pass_tools_with_tool_choice_tool() { - let server = MockServer::start().await; - mock_json_response( - &server, - json!({ - "id": "test-id", - "model": "test-model", - "usage": { "prompt_tokens": 4, "total_tokens": 34, "completion_tokens": 30 }, - "object": "chat.completion", - "choices": [{ - "index": 0, - "finish_reason": "stop", - "message": { "role": "assistant", "content": "ok" } - }] - }), - ) - .await; - - let config = MistralConfig::new("test-api-key").with_base_url(server.uri()); - let provider = MistralProvider::new(config); - let model = provider.model("test-model"); - - let options = CallOptions { - prompt: test_prompt(), - tools: Some(vec![test_tool()]), - tool_choice: Some(ToolChoice::Tool { - tool_name: "test-tool".to_string(), - }), - ..default_options(Vec::new()) - }; - - let result = model.do_generate(&options).await.expect("should succeed"); - let body = result.request.expect("request").body.expect("body"); - - assert_eq!(body["tool_choice"], json!("any")); - assert_eq!(body["tools"].as_array().unwrap().len(), 1); - assert_eq!(body["tools"][0]["function"]["name"], "test-tool"); -} diff --git a/aimux-providers/tests/moonshotai_test.rs b/aimux-providers/tests/moonshotai_test.rs index 8b6bbb8ad..0760eb048 100644 --- a/aimux-providers/tests/moonshotai_test.rs +++ b/aimux-providers/tests/moonshotai_test.rs @@ -23,10 +23,8 @@ use serial_test::serial; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; @@ -37,11 +35,7 @@ use aimux_providers::{ProviderOptions, provider, provider_from_env}; // ── shared helpers ─────────────────────────────────────────────────────────── fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn default_options(prompt: LanguageModelPrompt) -> CallOptions { diff --git a/aimux-providers/tests/open_responses_test.rs b/aimux-providers/tests/open_responses_test.rs index c07300eb5..8cf842c39 100644 --- a/aimux-providers/tests/open_responses_test.rs +++ b/aimux-providers/tests/open_responses_test.rs @@ -17,12 +17,13 @@ use serde_json::{Value, json}; use wiremock::matchers::method; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{ + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, TextPart, UserPart, +}; use aimux_core::options::{CallOptions, ResponseFormat, Tool}; use aimux_core::result::{GenerateContent, ReasoningOutput}; +use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; use aimux_core::tool::FunctionTool; use aimux_core::types::{FinishReasonUnified, ReasoningEffort}; @@ -38,11 +39,7 @@ use aimux_providers::open_responses::{ /// The TS `TEST_PROMPT`: a single user text message "Hello". fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } /// `CallOptions` with only `prompt` set (everything else default/None). @@ -292,10 +289,9 @@ mod convert_tests { #[test] fn convert_single_system_message_to_instructions() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text("You are a helpful assistant.")], - ..Default::default() + let prompt = vec![LanguageModelMessage::System { + content: "You are a helpful assistant.".into(), + provider_options: None, }]; let (input, instructions, warnings) = convert_to_open_responses_input(&prompt); assert_eq!( @@ -309,15 +305,13 @@ mod convert_tests { #[test] fn convert_multiple_system_messages_joined_with_newlines() { let prompt = vec![ - LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text("You are a helpful assistant.")], - ..Default::default() + LanguageModelMessage::System { + content: "You are a helpful assistant.".into(), + provider_options: None, }, - LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text("Always be concise.")], - ..Default::default() + LanguageModelMessage::System { + content: "Always be concise.".into(), + provider_options: None, }, ]; let (_, instructions, _) = convert_to_open_responses_input(&prompt); @@ -337,20 +331,17 @@ mod convert_tests { #[test] fn convert_system_message_with_user_and_assistant() { let prompt = vec![ - LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text("You are a helpful assistant.")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() + LanguageModelMessage::System { + content: "You are a helpful assistant.".into(), + provider_options: None, }, - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::text("Hi there!")], - ..Default::default() + LanguageModelMessage::user_text("Hello"), + LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: "Hi there!".into(), + provider_options: None, + })], + provider_options: None, }, ]; let (input, instructions, _) = convert_to_open_responses_input(&prompt); @@ -393,10 +384,16 @@ mod convert_tests { #[test] fn convert_image_file_base64_to_input_image() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::file_base64("ZmFrZS1kYXRh", "image/png")], - ..Default::default() + let prompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64("ZmFrZS1kYXRh".into()), + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + })], + provider_options: None, }]; let (input, _, warnings) = convert_to_open_responses_input(&prompt); assert_eq!( @@ -410,36 +407,24 @@ mod convert_tests { assert!(warnings.is_empty()); } - #[test] - fn convert_image_file_url_to_input_image() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::file_url( - "https://example.com/image.png", - "image/png", - )], - ..Default::default() - }]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([{ - "type": "message", - "role": "user", - "content": [{"type": "input_image", "image_url": "https://example.com/image.png"}] - }]) - ); - } - #[test] fn convert_pdf_file_base64_to_input_file() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, + let prompt = vec![LanguageModelMessage::User { content: vec![ - ContentPart::text("What does this PDF say?"), - ContentPart::file_base64("UERGREFUQQ==", "application/pdf"), + UserPart::Text(TextPart { + text: "What does this PDF say?".into(), + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64("UERGREFUQQ==".into()), + }, + media_type: "application/pdf".into(), + filename: None, + provider_options: None, + }), ], - ..Default::default() + provider_options: None, }]; let (input, _, warnings) = convert_to_open_responses_input(&prompt); assert_eq!( @@ -456,39 +441,18 @@ mod convert_tests { assert!(warnings.is_empty()); } - #[test] - fn convert_pdf_file_url_to_input_file() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::file_url( - "https://example.com/document.pdf", - "application/pdf", - )], - ..Default::default() - }]; - let (input, _, warnings) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([{ - "type": "message", - "role": "user", - "content": [{"type": "input_file", "file_url": "https://example.com/document.pdf"}] - }]) - ); - assert!(warnings.is_empty()); - } - #[test] fn convert_file_with_custom_filename() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::FileBase64 { - data: "UERGREFUQQ==".to_string(), + let prompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64("UERGREFUQQ==".to_string()), + }, media_type: "application/pdf".to_string(), filename: Some("report.pdf".to_string()), provider_options: None, - }], - ..Default::default() + })], + provider_options: None, }]; let (input, _, _) = convert_to_open_responses_input(&prompt); assert_eq!( @@ -509,10 +473,12 @@ mod convert_tests { #[test] fn convert_assistant_text_part() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::text("Hello from assistant")], - ..Default::default() + let prompt = vec![LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: "Hello from assistant".into(), + provider_options: None, + })], + provider_options: None, }]; let (input, _, _) = convert_to_open_responses_input(&prompt); assert_eq!( @@ -527,13 +493,18 @@ mod convert_tests { #[test] fn convert_assistant_multiple_text_parts() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, + let prompt = vec![LanguageModelMessage::Assistant { content: vec![ - ContentPart::text("First part"), - ContentPart::text("Second part"), + AssistantPart::Text(TextPart { + text: "First part".into(), + provider_options: None, + }), + AssistantPart::Text(TextPart { + text: "Second part".into(), + provider_options: None, + }), ], - ..Default::default() + provider_options: None, }]; let (input, _, _) = convert_to_open_responses_input(&prompt); assert_eq!( @@ -551,278 +522,22 @@ mod convert_tests { // -- Assistant messages with tool calls -- - #[test] - fn convert_assistant_single_tool_call() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::tool_call( - "call_123", - "get_weather", - json!({"location": "San Francisco"}), - )], - ..Default::default() - }]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([{ - "type": "function_call", - "call_id": "call_123", - "name": "get_weather", - "arguments": "{\"location\":\"San Francisco\"}" - }]) - ); - } - - #[test] - fn convert_assistant_tool_call_string_input() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::tool_call( - "call_124", - "get_weather", - Value::String("{\"location\":\"Berlin\"}".to_string()), - )], - ..Default::default() - }]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([{ - "type": "function_call", - "call_id": "call_124", - "name": "get_weather", - "arguments": "{\"location\":\"Berlin\"}" - }]) - ); - } - - #[test] - fn convert_assistant_text_and_tool_call() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ - ContentPart::text("Let me check the weather for you."), - ContentPart::tool_call("call_456", "get_weather", json!({"location": "New York"})), - ], - ..Default::default() - }]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([ - { - "type": "message", - "role": "assistant", - "content": [{"type": "output_text", "text": "Let me check the weather for you."}] - }, - { - "type": "function_call", - "call_id": "call_456", - "name": "get_weather", - "arguments": "{\"location\":\"New York\"}" - } - ]) - ); - } - - #[test] - fn convert_assistant_multiple_tool_calls() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ - ContentPart::tool_call("call_001", "get_weather", json!({"location": "Paris"})), - ContentPart::tool_call("call_002", "get_time", json!({"timezone": "Europe/Paris"})), - ], - ..Default::default() - }]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([ - { - "type": "function_call", - "call_id": "call_001", - "name": "get_weather", - "arguments": "{\"location\":\"Paris\"}" - }, - { - "type": "function_call", - "call_id": "call_002", - "name": "get_time", - "arguments": "{\"timezone\":\"Europe/Paris\"}" - } - ]) - ); - } - // -- Tool messages -- - #[test] - fn convert_tool_message_json_output() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result( - "call_123", - json!({"type": "json", "value": {"temperature": 72, "condition": "sunny"}}), - )], - ..Default::default() - }]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([{ - "type": "function_call_output", - "call_id": "call_123", - "output": "{\"temperature\":72,\"condition\":\"sunny\"}" - }]) - ); - } - - #[test] - fn convert_tool_message_text_output() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result( - "call_456", - json!({"type": "text", "value": "Search results: Found 5 items"}), - )], - ..Default::default() - }]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([{ - "type": "function_call_output", - "call_id": "call_456", - "output": "Search results: Found 5 items" - }]) - ); - } - - #[test] - fn convert_tool_message_error_text_output() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result( - "call_789", - json!({"type": "error-text", "value": "API request failed: timeout"}), - )], - ..Default::default() - }]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([{ - "type": "function_call_output", - "call_id": "call_789", - "output": "API request failed: timeout" - }]) - ); - } - - #[test] - fn convert_tool_message_execution_denied_output() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result( - "call_denied", - json!({"type": "execution-denied", "reason": "User declined the action"}), - )], - ..Default::default() - }]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([{ - "type": "function_call_output", - "call_id": "call_denied", - "output": "User declined the action" - }]) - ); - } - - #[test] - fn convert_tool_message_content_output_text() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result( - "call_content", - json!({ - "type": "content", - "value": [ - {"type": "text", "text": "First result"}, - {"type": "text", "text": "Second result"} - ] - }), - )], - ..Default::default() - }]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([{ - "type": "function_call_output", - "call_id": "call_content", - "output": [ - {"type": "input_text", "text": "First result"}, - {"type": "input_text", "text": "Second result"} - ] - }]) - ); - } - - #[test] - fn convert_tool_message_multiple_results() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ - ContentPart::tool_result( - "call_001", - json!({"type": "json", "value": {"temp": 72}}), - ), - ContentPart::tool_result("call_002", json!({"type": "text", "value": "3:00 PM"})), - ], - ..Default::default() - }]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([ - { - "type": "function_call_output", - "call_id": "call_001", - "output": "{\"temp\":72}" - }, - { - "type": "function_call_output", - "call_id": "call_002", - "output": "3:00 PM" - } - ]) - ); - } - // -- Message chains -- #[test] fn convert_user_assistant_user_chain() { let prompt = vec![ - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("What is the capital of France?")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::text("The capital of France is Paris.")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("And what about Germany?")], - ..Default::default() + LanguageModelMessage::user_text("What is the capital of France?"), + LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: "The capital of France is Paris.".into(), + provider_options: None, + })], + provider_options: None, }, + LanguageModelMessage::user_text("And what about Germany?"), ]; let (input, _, _) = convert_to_open_responses_input(&prompt); assert_eq!( @@ -835,97 +550,22 @@ mod convert_tests { ); } - #[test] - fn convert_user_assistant_tool_tool_chain() { - let prompt = vec![ - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("What is the weather in Tokyo?")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::tool_call( - "call_weather", - "get_weather", - json!({"location": "Tokyo"}), - )], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result( - "call_weather", - json!({"type": "json", "value": {"temperature": 25, "condition": "cloudy"}}), - )], - ..Default::default() - }, - ]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([ - {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is the weather in Tokyo?"}]}, - {"type": "function_call", "call_id": "call_weather", "name": "get_weather", "arguments": "{\"location\":\"Tokyo\"}"}, - {"type": "function_call_output", "call_id": "call_weather", "output": "{\"temperature\":25,\"condition\":\"cloudy\"}"} - ]) - ); - } - - #[test] - fn convert_tool_roundtrip_with_followup_assistant() { - let prompt = vec![ - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("What is the weather in Tokyo?")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::tool_call( - "call_weather", - "get_weather", - Value::String("{\"location\":\"Tokyo\"}".to_string()), - )], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result( - "call_weather", - json!({"type": "json", "value": {"temperature": 25, "condition": "cloudy"}}), - )], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::text("It is 25 C and cloudy in Tokyo.")], - ..Default::default() - }, - ]; - let (input, _, _) = convert_to_open_responses_input(&prompt); - assert_eq!( - input, - json!([ - {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is the weather in Tokyo?"}]}, - {"type": "function_call", "call_id": "call_weather", "name": "get_weather", "arguments": "{\"location\":\"Tokyo\"}"}, - {"type": "function_call_output", "call_id": "call_weather", "output": "{\"temperature\":25,\"condition\":\"cloudy\"}"}, - {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "It is 25 C and cloudy in Tokyo."}]} - ]) - ); - } - // -- Provider reference (warning, not error in Rust) -- #[test] fn convert_file_reference_produces_warning() { - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::file_reference( - "image/png", - json!({"openResponses": "file-ref-123"}), - )], - ..Default::default() + let prompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Reference { + reference: [("openResponses".into(), "file-ref-123".into())] + .into_iter() + .collect(), + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + })], + provider_options: None, }]; let (input, _, warnings) = convert_to_open_responses_input(&prompt); // The reference part is skipped, so user content is empty array. @@ -1312,16 +952,11 @@ mod do_generate_tests { let model = make_model(&server, "gemma-7b-it"); let prompt = vec![ - LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text("You are a helpful assistant.")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() + LanguageModelMessage::System { + content: "You are a helpful assistant.".into(), + provider_options: None, }, + LanguageModelMessage::user_text("Hello"), ]; model.do_generate(&default_options(prompt)).await.unwrap(); @@ -1337,21 +972,15 @@ mod do_generate_tests { let model = make_model(&server, "gemma-7b-it"); let prompt = vec![ - LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text("You are a helpful assistant.")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text("Always be concise.")], - ..Default::default() + LanguageModelMessage::System { + content: "You are a helpful assistant.".into(), + provider_options: None, }, - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() + LanguageModelMessage::System { + content: "Always be concise.".into(), + provider_options: None, }, + LanguageModelMessage::user_text("Hello"), ]; model.do_generate(&default_options(prompt)).await.unwrap(); @@ -1365,76 +994,6 @@ mod do_generate_tests { // -- Multi-turn tool conversation -- - #[tokio::test] - async fn multi_turn_tool_conversation_request_body() { - let server = MockServer::start().await; - mock_json(&server, lmstudio_basic_json()).await; - - let model = make_model(&server, "gemma-7b-it"); - let prompt = vec![ - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("What is the weather in Tokyo?")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::tool_call( - "call_weather_123", - "get_weather", - json!({"location": "Tokyo"}), - )], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result( - "call_weather_123", - json!({ - "type": "json", - "value": {"temperature": 22, "condition": "sunny", "humidity": 65} - }), - )], - ..Default::default() - }, - ]; - let options = CallOptions { - tools: Some(vec![Tool::from( - FunctionTool::new( - "get_weather", - json!({ - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"] - }), - ) - .with_description("Get the current weather for a location"), - )]), - ..CallOptions::new(prompt) - }; - model.do_generate(&options).await.unwrap(); - - let request = &server.received_requests().await.unwrap()[0]; - let body: Value = serde_json::from_slice(&request.body).unwrap(); - assert_eq!( - body, - json!({ - "model": "gemma-7b-it", - "input": [ - {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is the weather in Tokyo?"}]}, - {"type": "function_call", "call_id": "call_weather_123", "name": "get_weather", "arguments": "{\"location\":\"Tokyo\"}"}, - {"type": "function_call_output", "call_id": "call_weather_123", "output": "{\"temperature\":22,\"condition\":\"sunny\",\"humidity\":65}"} - ], - "tools": [{ - "type": "function", - "name": "get_weather", - "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, - "description": "Get the current weather for a location" - }] - }) - ); - } - // -- PDF input file -- #[tokio::test] @@ -1443,18 +1002,23 @@ mod do_generate_tests { mock_json(&server, openai_pdf_json()).await; let model = make_model(&server, "gpt-4.1-nano"); - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, + let prompt = vec![LanguageModelMessage::User { content: vec![ - ContentPart::text( - "What text does this PDF contain? Reply with just the text content, nothing else.", - ), - ContentPart::file_url( - "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf", - "application/pdf", - ), - ], - ..Default::default() +UserPart::Text(TextPart { +text: "What text does this PDF contain? Reply with just the text content, nothing else.".into(), +provider_options: None +}), +UserPart::File(FilePart { +data: FileData::Url { +url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into(), +original_url: None +}, +media_type: "application/pdf".into(), +filename: None, +provider_options: None +}) +], + provider_options: None, }]; model.do_generate(&default_options(prompt)).await.unwrap(); @@ -1482,16 +1046,23 @@ mod do_generate_tests { mock_json(&server, openai_pdf_json()).await; let model = make_model(&server, "gpt-4.1-nano"); - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, + let prompt = vec![LanguageModelMessage::User { content: vec![ - ContentPart::text("What text does this PDF contain?"), - ContentPart::file_url( - "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf", - "application/pdf", - ), - ], - ..Default::default() +UserPart::Text(TextPart { +text: "What text does this PDF contain?".into(), +provider_options: None +}), +UserPart::File(FilePart { +data: FileData::Url { +url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into(), +original_url: None +}, +media_type: "application/pdf".into(), +filename: None, +provider_options: None +}) +], + provider_options: None, }]; let result = model.do_generate(&default_options(prompt)).await.unwrap(); @@ -1511,16 +1082,23 @@ mod do_generate_tests { mock_json(&server, openai_pdf_json()).await; let model = make_model(&server, "gpt-4.1-nano"); - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, + let prompt = vec![LanguageModelMessage::User { content: vec![ - ContentPart::text("What text does this PDF contain?"), - ContentPart::file_url( - "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf", - "application/pdf", - ), - ], - ..Default::default() +UserPart::Text(TextPart { +text: "What text does this PDF contain?".into(), +provider_options: None +}), +UserPart::File(FilePart { +data: FileData::Url { +url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into(), +original_url: None +}, +media_type: "application/pdf".into(), +filename: None, +provider_options: None +}) +], + provider_options: None, }]; let result = model.do_generate(&default_options(prompt)).await.unwrap(); @@ -1788,16 +1366,23 @@ mod do_stream_tests { mock_sse(&server, body).await; let model = make_model(&server, "gpt-4.1-nano"); - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, + let prompt = vec![LanguageModelMessage::User { content: vec![ - ContentPart::text("What text does this PDF contain?"), - ContentPart::file_url( - "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf", - "application/pdf", - ), - ], - ..Default::default() +UserPart::Text(TextPart { +text: "What text does this PDF contain?".into(), +provider_options: None +}), +UserPart::File(FilePart { +data: FileData::Url { +url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into(), +original_url: None +}, +media_type: "application/pdf".into(), +filename: None, +provider_options: None +}) +], + provider_options: None, }]; let result = model.do_stream(&default_options(prompt)).await.unwrap(); let parts = collect_stream(result).await; diff --git a/aimux-providers/tests/openai_compatible_test.rs b/aimux-providers/tests/openai_compatible_test.rs index 645e1ad29..96830dd28 100644 --- a/aimux-providers/tests/openai_compatible_test.rs +++ b/aimux-providers/tests/openai_compatible_test.rs @@ -21,11 +21,9 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; @@ -60,11 +58,7 @@ use aimux_providers::{ /// The TS `TEST_PROMPT`: a single user text message "Hello". fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } /// `CallOptions` with only `prompt` set (everything else default/None). diff --git a/aimux-providers/tests/openai_provider_test.rs b/aimux-providers/tests/openai_provider_test.rs deleted file mode 100644 index 4dcfe790b..000000000 --- a/aimux-providers/tests/openai_provider_test.rs +++ /dev/null @@ -1,282 +0,0 @@ -//! OpenAI provider configuration and forward-compatible defaults tests. -//! -//! Translates the chat-completions-relevant parts of: -//! - `packages/openai/src/openai-provider.test.ts` — baseURL config, chat routing -//! - `packages/openai/src/openai-forward-compatible-defaults.test.ts` — reasoning-safe defaults -//! -//! Responses API / embedding / image parts are excluded (covered by B1/C1/C2). - -use serde_json::{Value, json}; -use serial_test::serial; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::content::ContentPart; -use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; -use aimux_core::options::CallOptions; -use aimux_core::provider::Provider; - -use aimux_providers::{OpenAIConfig, OpenAIProvider}; - -// ── helpers ─────────────────────────────────────────────────────────────────── - -fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] -} - -fn default_options(prompt: LanguageModelPrompt) -> CallOptions { - CallOptions::new(prompt) -} - -fn text_completion_body() -> Value { - json!({ - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1711115037, - "model": "test-model", - "choices": [{ - "index": 0, - "message": { "role": "assistant", "content": "ok" }, - "finish_reason": "stop" - }], - "usage": { "prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2 } - }) -} - -// ════════════════════════════════════════════════════════════════════════════ -// Provider configuration (openai-provider.test.ts) -// ════════════════════════════════════════════════════════════════════════════ - -mod provider_config { - use super::*; - - /// TS: `createOpenAI()` provider name is "openai". - #[test] - fn provider_name_is_openai() { - let config = OpenAIConfig::new("test-key"); - let provider = OpenAIProvider::new(config); - assert_eq!(provider.name(), "openai"); - } - - /// TS: default base URL is `https://api.openai.com/v1`. - #[test] - fn default_base_url() { - let config = OpenAIConfig::new("test-key"); - assert_eq!(config.base_url, "https://api.openai.com/v1"); - } - - /// TS: custom base URL overrides the default. - #[test] - fn custom_base_url_override() { - let config = - OpenAIConfig::new("test-key").with_base_url("https://proxy.openai.example/v1/"); - // `with_base_url` strips trailing slash. - assert_eq!(config.base_url, "https://proxy.openai.example/v1"); - } - - /// TS: `from_env` loads `OPENAI_API_KEY`. - #[serial] - #[test] - fn from_env_loads_openai_api_key() { - let saved = std::env::var("OPENAI_API_KEY").ok(); - unsafe { std::env::set_var("OPENAI_API_KEY", "env-test-key") }; - - let config = OpenAIConfig::from_env(); - assert!(config.is_ok(), "from_env should succeed with env var set"); - - unsafe { - match saved { - Some(v) => std::env::set_var("OPENAI_API_KEY", v), - None => std::env::remove_var("OPENAI_API_KEY"), - } - } - } - - /// TS: `from_env` fails without the env var. - #[serial] - #[test] - fn from_env_fails_without_env_var() { - let saved = std::env::var("OPENAI_API_KEY").ok(); - unsafe { std::env::remove_var("OPENAI_API_KEY") }; - - let config = OpenAIConfig::from_env(); - assert!(config.is_err(), "from_env should fail without env var"); - - unsafe { - if let Some(v) = saved { - std::env::set_var("OPENAI_API_KEY", v); - } - } - } - - /// TS: chat completions API routes to `/chat/completions`. - #[tokio::test] - async fn chat_routes_to_chat_completions() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/chat/completions")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_completion_body())) - .mount(&server) - .await; - - let config = OpenAIConfig::new("test-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.model("gpt-4o-mini"); - - model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - } - - /// TS: custom API key is sent in Authorization header. - #[tokio::test] - async fn custom_api_key_in_auth_header() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/chat/completions")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_completion_body())) - .mount(&server) - .await; - - let config = OpenAIConfig::new("my-custom-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.model("gpt-4o"); - - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - assert_eq!( - requests[0] - .headers - .get("authorization") - .and_then(|v| v.to_str().ok()), - Some("Bearer my-custom-key"), - ); - } - - /// TS: `languageModel` via Provider trait creates a working model. - #[tokio::test] - async fn language_model_via_trait() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/chat/completions")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_completion_body())) - .mount(&server) - .await; - - let config = OpenAIConfig::new("test-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider - .language_model("gpt-4o") - .expect("language_model should succeed"); - - model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - } - - /// TS: org ID is sent in the OpenAI-Organization header. - #[tokio::test] - async fn org_id_in_header() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/chat/completions")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_completion_body())) - .mount(&server) - .await; - - let config = OpenAIConfig::new("test-key") - .with_base_url(server.uri()) - .with_org_id("org-123"); - let provider = OpenAIProvider::new(config); - let model = provider.model("gpt-4o"); - - let _ = model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - assert_eq!( - requests[0] - .headers - .get("openai-organization") - .and_then(|v| v.to_str().ok()), - Some("org-123"), - ); - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// Forward-compatible defaults (openai-forward-compatible-defaults.test.ts) -// ════════════════════════════════════════════════════════════════════════════ - -mod forward_compatible_defaults { - use super::*; - - /// TS: gpt-99 (reasoning model) should use reasoning-safe Chat Completions - /// defaults: system→developer, max_completion_tokens instead of max_tokens, - /// temperature/top_p/penalties/logit_bias/logprobs stripped. - /// - /// We verify by checking the request body sent to the mock server. - #[tokio::test] - async fn reasoning_model_uses_safe_defaults() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/chat/completions")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_completion_body())) - .mount(&server) - .await; - - let config = OpenAIConfig::new("test-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.model("o1"); - - let prompt = vec![ - LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text("Follow the instructions.")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Say ok.")], - ..Default::default() - }, - ]; - let mut options = CallOptions::new(prompt); - options.max_output_tokens = Some(64); - options.temperature = Some(0.2); - options.top_p = Some(0.8); - - let _ = model.do_generate(&options).await.expect("should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - - // Reasoning models use `max_completion_tokens` not `max_tokens`. - assert!( - body.get("max_completion_tokens").is_some(), - "should use max_completion_tokens for reasoning models" - ); - // System message should be converted to developer for reasoning models. - assert_eq!( - body["messages"][0]["role"], "developer", - "system should become developer for reasoning models" - ); - } -} diff --git a/aimux-providers/tests/openrouter_test.rs b/aimux-providers/tests/openrouter_test.rs index 703e63248..cca5ff453 100644 --- a/aimux-providers/tests/openrouter_test.rs +++ b/aimux-providers/tests/openrouter_test.rs @@ -24,11 +24,9 @@ use serial_test::serial; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::{CallOptions, Tool}; use aimux_core::provider::Provider; use aimux_core::result::GenerateContent; @@ -45,11 +43,7 @@ mod common; // ── shared helpers ─────────────────────────────────────────────────────────── fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn default_options(prompt: LanguageModelPrompt) -> CallOptions { diff --git a/aimux-providers/tests/perplexity_test.rs b/aimux-providers/tests/perplexity_test.rs index 16805a692..e59e4ea6d 100644 --- a/aimux-providers/tests/perplexity_test.rs +++ b/aimux-providers/tests/perplexity_test.rs @@ -26,11 +26,9 @@ use serial_test::serial; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; @@ -41,11 +39,7 @@ use aimux_providers::{ProviderOptions, provider, provider_from_env}; // ── shared helpers ─────────────────────────────────────────────────────────── fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn default_options(prompt: LanguageModelPrompt) -> CallOptions { diff --git a/aimux-providers/tests/provider_error_test.rs b/aimux-providers/tests/provider_error_test.rs index a4a2857cc..fad95c2d0 100644 --- a/aimux-providers/tests/provider_error_test.rs +++ b/aimux-providers/tests/provider_error_test.rs @@ -23,11 +23,9 @@ //! yields `AiMuxError::ApiCall` with the observed code in `status_code`; //! 408/409/429/5xx are stored retryable (`is_retryable`), other 4xx are not. -use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::stream_part::StreamPart; use aimux_provider_utils::RetryConfig; @@ -41,11 +39,7 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; /// The TS `TEST_PROMPT`: a single user text message "Hello". fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } /// Build `CallOptions` with everything unset except `prompt`. diff --git a/aimux-providers/tests/reasoning_map_test.rs b/aimux-providers/tests/reasoning_map_test.rs index 68ffcb0c7..0f9bdc6d1 100644 --- a/aimux-providers/tests/reasoning_map_test.rs +++ b/aimux-providers/tests/reasoning_map_test.rs @@ -14,9 +14,7 @@ //! - **reasoning_effort 直传**：7 档无归一化（none/minimal/low/medium/high/xhigh 原样 //! 透传；provider-default 不发字段）。 -use aimux_core::content::ContentPart; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::types::{ReasoningEffort, Warning}; use aimux_providers::openai::OpenAICompatProfile; @@ -25,11 +23,7 @@ use aimux_providers::provider_registry_entry; use serde_json::json; fn user_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn opts_with_max_tokens(n: u32) -> CallOptions { diff --git a/aimux-providers/tests/stream_error_semantics_test.rs b/aimux-providers/tests/stream_error_semantics_test.rs deleted file mode 100644 index 734ef9df2..000000000 --- a/aimux-providers/tests/stream_error_semantics_test.rs +++ /dev/null @@ -1,139 +0,0 @@ -//! Stream-error classification and termination semantics (review follow-ups). -//! -//! Locks three behaviors: -//! - no fabricated HTTP 500 for non-numeric provider error codes (the "M3 -//! bug"): such errors carry no status and are not retryable; -//! - payload `retry_after_ms` hints surface through `retry_after_hint()`; -//! - the shared Responses reducer ends its stream right after a terminal -//! error event instead of waiting on a source that stays open. - -use std::collections::HashMap; - -use futures::StreamExt; -use serde_json::json; - -use aimux_core::AiMuxError; -use aimux_core::stream_part::StreamPart; -use aimux_core::types::FinishReasonUnified; -use aimux_provider_utils::stream_error_api_call; -use aimux_providers::openai::responses::responses_convert::build_responses_event_stream; - -fn as_api_call(error: &AiMuxError) -> &aimux_core::ApiCallError { - match error { - AiMuxError::ApiCall(detail) => detail, - other => panic!("expected ApiCall, got {other:?}"), - } -} - -#[test] -fn string_error_code_is_not_a_retryable_500() { - let payload = json!({"message": "Incorrect API key", "code": "invalid_api_key"}); - let error = stream_error_api_call( - "Incorrect API key", - Some("invalid_api_key".into()), - None, - &payload, - "https://example.test", - json!({}), - HashMap::new(), - ); - let detail = as_api_call(&error); - assert_eq!(detail.status_code, None); - assert!(!error.is_retryable()); -} - -#[test] -fn numeric_status_keeps_status_based_retryability() { - let payload = json!({"message": "rate limited", "code": 429}); - let error = stream_error_api_call( - "rate limited", - Some("429".into()), - Some(429), - &payload, - "https://example.test", - json!({}), - HashMap::new(), - ); - assert_eq!(as_api_call(&error).status_code, Some(429)); - assert!(error.is_retryable()); -} - -#[test] -fn payload_retry_after_ms_surfaces_as_hint() { - let payload = json!({"message": "rate limited", "code": 429, "retry_after_ms": 15000}); - let error = stream_error_api_call( - "rate limited", - None, - Some(429), - &payload, - "https://example.test", - json!({}), - HashMap::new(), - ); - assert_eq!(error.retry_after_hint(), Some(15000)); -} - -#[test] -fn real_retry_after_header_wins_over_payload() { - let payload = json!({"message": "rate limited", "retry_after_ms": 15000}); - let headers = HashMap::from([("retry-after-ms".to_string(), "7".to_string())]); - let error = stream_error_api_call( - "rate limited", - None, - Some(429), - &payload, - "https://example.test", - json!({}), - HashMap::new(), - ); - assert_eq!(error.retry_after_hint(), Some(15000)); - let error = stream_error_api_call( - "rate limited", - None, - Some(429), - &payload, - "https://example.test", - json!({}), - headers, - ); - assert_eq!(error.retry_after_hint(), Some(7)); -} - -#[tokio::test] -async fn responses_stream_ends_after_terminal_error_event() { - // First event is benign so the peek passes; the mid-stream error event is - // followed by a source that never closes (heartbeat-style server). - let first = json!({"type": "response.created", "response": {"id": "resp_1"}}); - let error_event = - json!({"type": "error", "error": {"message": "boom", "type": "server_error"}}); - let source = futures::stream::iter(vec![Ok(error_event)]).chain(futures::stream::pending()); - - let stream = build_responses_event_stream( - Some(Ok(first)), - source, - "openai".to_string(), - Vec::new(), - false, - "https://example.test".to_string(), - json!({}), - HashMap::new(), - ) - .expect("stream setup"); - - let parts: Vec<_> = stream.collect().await; - assert!( - parts - .iter() - .any(|part| matches!(part, Ok(StreamPart::Error { .. }))), - "error part yielded" - ); - let last = parts.last().expect("stream not empty"); - assert!( - matches!( - last, - Ok(StreamPart::Finish { finish_reason, .. }) - if finish_reason.unified == FinishReasonUnified::Error - ), - "stream ends with an error Finish, got {last:?}" - ); -} diff --git a/aimux-providers/tests/thin_wrapper_config_test.rs b/aimux-providers/tests/thin_wrapper_config_test.rs index 7e5624f2f..5eb2a5c71 100644 --- a/aimux-providers/tests/thin_wrapper_config_test.rs +++ b/aimux-providers/tests/thin_wrapper_config_test.rs @@ -21,10 +21,8 @@ use serde_json::{Value, json}; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::provider::Provider; @@ -35,11 +33,7 @@ use aimux_providers::{ // ── shared helpers ─────────────────────────────────────────────────────────── fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] + vec![LanguageModelMessage::user_text("Hello")] } fn default_options(prompt: LanguageModelPrompt) -> CallOptions { diff --git a/aimux-providers/tests/tool_result_and_reasoning_fix_test.rs b/aimux-providers/tests/tool_result_and_reasoning_fix_test.rs deleted file mode 100644 index f9807c573..000000000 --- a/aimux-providers/tests/tool_result_and_reasoning_fix_test.rs +++ /dev/null @@ -1,315 +0,0 @@ -// Panic convert wrappers are #[deprecated]; these tests still use them. -#![allow(deprecated)] -//! Regression tests for two issues reported against aimux 0.1.1 when driving -//! OpenAI-compatible thinking models (e.g. DeepSeek `deepseek-v4-flash`) in -//! multi-turn tool-call conversations: -//! -//! 1. **`tool` role with `ContentPart[]` was rejected.** `ModelPrompt` -//! deserialization of a `tool_result` part built with the legacy `output` -//! field (the shape emitted by the Vercel AI SDK and the 0.1.0 TypeScript -//! bindings) failed with "data did not match any variant of untagged enum -//! ModelPrompt". `ContentPart::ToolResult.result` now accepts `output` as a -//! deserialization alias (`#[serde(alias = "output")]`). -//! -//! 2. **`reasoning` ContentPart was dropped on the request side.** Thinking -//! models require prior assistant `reasoning_content` to be replayed on -//! later turns, including tool-call turns. The OpenAI message converter -//! now lifts `ContentPart::Reasoning` parts to a top-level -//! `reasoning_content` string on assistant messages (mirroring the Vercel -//! AI SDK `openai-compatible` assistant conversion). -//! -//! These are pure-function tests against -//! `convert_prompt_to_openai_messages` (no network). - -use aimux_core::content::ContentPart; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::{ModelMessage, ModelPrompt, Role}; -use aimux_providers::openai::convert::convert_prompt_to_openai_messages; -use serde_json::{Value, json}; - -// ── Issue 1: tool role ContentPart[] / `output` alias ─────────────────────── - -/// A `tool` message whose `tool_result` part uses the legacy `output` field -/// (Vercel AI SDK / 0.1.0 TS bindings shape) round-trips through `ModelPrompt`. -#[test] -fn tool_message_with_tool_result_output_field_deserializes() { - let json_str = r#"[ - {"role":"tool","content":[{"type":"tool_result","tool_call_id":"tc1","output":"ok"}]} - ]"#; - let prompt: ModelPrompt = - serde_json::from_str(json_str).expect("output-field tool_result must deserialize"); - let msgs = match prompt { - ModelPrompt::Messages(m) => m, - other => panic!("expected Messages, got {other:?}"), - }; - assert_eq!(msgs.len(), 1); - assert_eq!(msgs[0].role, Role::Tool); - // The alias only affects deserialization; the in-memory field is `result`. - match &msgs[0].content { - aimux_core::message::MessageContent::Parts(parts) => { - assert_eq!(parts.len(), 1); - match &parts[0] { - ContentPart::ToolResult { - tool_call_id, - result, - .. - } => { - assert_eq!(tool_call_id, "tc1"); - assert_eq!(result, &json!("ok")); - } - other => panic!("expected ToolResult, got {other:?}"), - } - } - other => panic!("expected Parts, got {other:?}"), - } -} - -/// The current `result` field name still deserializes (no regression). -#[test] -fn tool_message_with_tool_result_result_field_deserializes() { - let json_str = r#"[ - {"role":"tool","content":[{"type":"tool_result","tool_call_id":"tc1","result":"ok"}]} - ]"#; - let prompt: ModelPrompt = - serde_json::from_str(json_str).expect("result-field tool_result must deserialize"); - assert!(matches!(prompt, ModelPrompt::Messages(_))); -} - -/// A full tool round-trip (assistant tool_call → tool result → …) built from -/// raw JSON reproducing the user's failing case now converts to OpenAI -/// messages with `tool_call_id` on the tool message. -#[test] -fn full_tool_round_trip_from_json_converts_with_tool_call_id() { - let json_str = r#"[ - {"role":"user","content":"write a file"}, - {"role":"assistant","content":[ - {"type":"tool_call","tool_call_id":"tc1","tool_name":"write_file","input":{"path":"/tmp/test.txt"}} - ]}, - {"role":"tool","content":[ - {"type":"tool_result","tool_call_id":"tc1","output":"Successfully wrote to /tmp/test.txt"} - ]} - ]"#; - let prompt: ModelPrompt = serde_json::from_str(json_str).expect("must deserialize"); - let msgs = match prompt { - ModelPrompt::Messages(m) => m, - _ => unreachable!(), - }; - // Convert to provider-facing prompt then to OpenAI messages. - let provider_prompt: LanguageModelPrompt = msgs - .iter() - .map(|m| LanguageModelPromptMessage { - role: m.role, - content: match &m.content { - aimux_core::message::MessageContent::Text(t) => vec![ContentPart::text(t)], - aimux_core::message::MessageContent::Parts(p) => p.clone(), - }, - provider_options: None, - }) - .collect(); - let out = convert_prompt_to_openai_messages(&provider_prompt); - assert_eq!(out.len(), 3); - // The tool message must carry tool_call_id (the core of issue 1). - assert_eq!(out[2]["role"], json!("tool")); - assert_eq!(out[2]["tool_call_id"], json!("tc1")); - assert_eq!( - out[2]["content"], - json!("Successfully wrote to /tmp/test.txt") - ); -} - -// ── Issue 2: reasoning_content replay on the request side ─────────────────── - -/// An assistant message with reasoning + tool_call lifts the reasoning to a -/// top-level `reasoning_content` field (the DeepSeek V4 thinking-mode -/// requirement). This is the exact scenario from the user's table: assistant -/// `[{reasoning},{tool_call}]` + tool `string`. -#[test] -fn assistant_reasoning_with_tool_call_emits_reasoning_content() { - let prompt: LanguageModelPrompt = vec![ - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("inspect the repo")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ - ContentPart::reasoning("I need to inspect files before answering."), - ContentPart::tool_call( - "call_1".to_string(), - "read_file".to_string(), - json!({ "path": "README.md" }), - ), - ], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Tool, - content: vec![ContentPart::tool_result( - "call_1".to_string(), - json!("contents"), - )], - ..Default::default() - }, - ]; - let out = convert_prompt_to_openai_messages(&prompt); - assert_eq!(out.len(), 3); - - // Assistant message: reasoning_content present, content null, tool_calls set. - let assistant = &out[1]; - assert_eq!(assistant["role"], json!("assistant")); - assert_eq!( - assistant["reasoning_content"], - json!("I need to inspect files before answering.") - ); - assert_eq!(assistant["content"], Value::Null); - assert_eq!( - assistant["tool_calls"], - json!([{ - "type": "function", - "id": "call_1", - "function": { - "name": "read_file", - "arguments": "{\"path\":\"README.md\"}" - } - }]) - ); -} - -/// An assistant message with reasoning + text (no tool calls) lifts reasoning -/// to `reasoning_content` and keeps the text as `content`. -#[test] -fn assistant_reasoning_with_text_emits_reasoning_content() { - let prompt: LanguageModelPrompt = vec![ - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("What is 2+2?")], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ - ContentPart::reasoning("2 plus 2 equals 4."), - ContentPart::text("4"), - ], - ..Default::default() - }, - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("thanks")], - ..Default::default() - }, - ]; - let out = convert_prompt_to_openai_messages(&prompt); - let assistant = &out[1]; - assert_eq!(assistant["role"], json!("assistant")); - assert_eq!(assistant["content"], json!("4")); - assert_eq!(assistant["reasoning_content"], json!("2 plus 2 equals 4.")); - // No tool_calls key should be emitted. - assert!( - assistant.get("tool_calls").is_none(), - "tool_calls must not be emitted when there are none" - ); -} - -/// An assistant message with no reasoning part emits no `reasoning_content` -/// field (no spurious empty string — matches the Vercel `openai-compatible` -/// `reasoning.length > 0` guard). -#[test] -fn assistant_without_reasoning_omits_reasoning_content() { - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ContentPart::tool_call( - "call_1".to_string(), - "write_file".to_string(), - json!({ "path": "/tmp/x" }), - )], - ..Default::default() - }]; - let out = convert_prompt_to_openai_messages(&prompt); - let assistant = &out[0]; - assert!( - assistant.get("reasoning_content").is_none(), - "reasoning_content must be omitted when no reasoning part is present" - ); -} - -/// Multiple reasoning parts are concatenated (mirrors the SDK join behaviour). -#[test] -fn multiple_reasoning_parts_are_concatenated() { - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::Assistant, - content: vec![ - ContentPart::reasoning("part one. "), - ContentPart::reasoning("part two."), - ContentPart::text("answer"), - ], - ..Default::default() - }]; - let out = convert_prompt_to_openai_messages(&prompt); - let assistant = &out[0]; - assert_eq!(assistant["content"], json!("answer")); - assert_eq!(assistant["reasoning_content"], json!("part one. part two.")); -} - -/// A user message never carries reasoning, so `reasoning_content` is never -/// emitted for non-assistant roles even if a reasoning part somehow appears. -/// The reasoning part is still dropped from the content shape, and the -/// remaining text part collapses to a string (matching the plain-text path). -#[test] -fn non_assistant_role_never_emits_reasoning_content() { - let prompt: LanguageModelPrompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ - ContentPart::reasoning("should be ignored"), - ContentPart::text("hello"), - ], - ..Default::default() - }]; - let out = convert_prompt_to_openai_messages(&prompt); - let user = &out[0]; - assert_eq!(user["role"], json!("user")); - assert!( - user.get("reasoning_content").is_none(), - "user messages must not carry reasoning_content" - ); - // The text part collapses to a plain string; reasoning is dropped entirely. - assert_eq!(user["content"], json!("hello")); -} - -/// The user-facing `ModelMessage` + `ModelPrompt` path (not just the -/// provider-facing prompt) also surfaces reasoning_content, exercising the -/// `convert_to_language_model_prompt` → `convert_prompt_to_openai_messages` -/// pipeline end to end. -#[test] -fn model_prompt_path_emits_reasoning_content() { - let prompt = ModelPrompt::Messages(vec![ - ModelMessage::user("hi"), - ModelMessage { - role: Role::Assistant, - content: aimux_core::message::MessageContent::Parts(vec![ - ContentPart::reasoning("thinking about hi"), - ContentPart::text("hello!"), - ]), - }, - ]); - let msgs = match prompt { - ModelPrompt::Messages(m) => m, - _ => unreachable!(), - }; - let provider_prompt: LanguageModelPrompt = msgs - .iter() - .map(|m| LanguageModelPromptMessage { - role: m.role, - content: match &m.content { - aimux_core::message::MessageContent::Text(t) => vec![ContentPart::text(t)], - aimux_core::message::MessageContent::Parts(p) => p.clone(), - }, - provider_options: None, - }) - .collect(); - let out = convert_prompt_to_openai_messages(&provider_prompt); - let assistant = &out[1]; - assert_eq!(assistant["reasoning_content"], json!("thinking about hi")); - assert_eq!(assistant["content"], json!("hello!")); -} diff --git a/aimux-providers/tests/vertex_anthropic_test.rs b/aimux-providers/tests/vertex_anthropic_test.rs deleted file mode 100644 index ca23852b4..000000000 --- a/aimux-providers/tests/vertex_anthropic_test.rs +++ /dev/null @@ -1,400 +0,0 @@ -//! Wiremock tests for Anthropic partner models on Vertex AI (`rawPredict`). -//! -//! Claude models on Vertex AI are reached via the -//! `publishers/anthropic/models/{model}:rawPredict` (and `:streamRawPredict`) -//! endpoints. The request body is a standard Anthropic Messages request -//! wrapped in an `anthropic_version` envelope (model identity in the URL, not -//! the body); the response is a verbatim Anthropic Messages response. These -//! tests mock the Vertex AI endpoints and never touch the public network or -//! real credentials. - -use aimux_core::tool::RawToolCall; -use futures::StreamExt; -use serde_json::{Value, json}; -use wiremock::matchers::{header, method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::content::ContentPart; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; -use aimux_core::options::CallOptions; -use aimux_core::result::{GenerateContent, StreamResult}; -use aimux_core::stream_part::StreamPart; -use aimux_core::types::FinishReasonUnified; - -use aimux_providers::vertex::{ - VertexAnthropicModel, VertexAuth, VertexProvider, VertexProviderConfig, -}; - -// ── Constants ───────────────────────────────────────────────────────────────── - -const MODEL_ID: &str = "claude-sonnet-4-20250514"; -/// Endpoint path once the `/publishers/google` suffix has been swapped for -/// `/publishers/anthropic`. -const RAW_PREDICT_PATH: &str = "/publishers/anthropic/models/claude-sonnet-4-20250514:rawPredict"; -const STREAM_RAW_PREDICT_PATH: &str = - "/publishers/anthropic/models/claude-sonnet-4-20250514:streamRawPredict"; - -// ── Helpers ─────────────────────────────────────────────────────────────────── - -fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] -} - -fn default_options(prompt: LanguageModelPrompt) -> CallOptions { - CallOptions::new(prompt) -} - -/// Build a `VertexProvider` whose base URL ends in `/publishers/google` and -/// points at the mock server, then return the Anthropic partner model. This -/// exercises the `/publishers/google` → `/publishers/anthropic` base-URL -/// rewrite performed by [`VertexProvider::anthropic_model`]. -fn make_model(server: &MockServer) -> VertexAnthropicModel { - let config = VertexProviderConfig { - base_url: format!("{}/publishers/google", server.uri()), - project: Some("test-project".to_string()), - location: Some("us-central1".to_string()), - auth: VertexAuth::BearerToken("test-token".to_string()), - api_key_source: None, - retry_config: aimux_provider_utils::RetryConfig::default(), - }; - VertexProvider::new(config) - .anthropic_model(MODEL_ID) - .expect("anthropic_model should succeed") -} - -async fn mock_raw_predict_json(server: &MockServer, status: u16, body: Value) { - Mock::given(method("POST")) - .and(path(RAW_PREDICT_PATH)) - .respond_with(ResponseTemplate::new(status).set_body_json(body)) - .mount(server) - .await; -} - -async fn mock_stream_raw_predict_sse(server: &MockServer, sse_body: &str) { - Mock::given(method("POST")) - .and(path(STREAM_RAW_PREDICT_PATH)) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body.to_string()), - ) - .mount(server) - .await; -} - -fn sse(data: &Value) -> String { - format!("data: {data}\n\n") -} - -fn sse_stream(events: &[Value]) -> String { - events.iter().map(sse).collect() -} - -fn as_text(item: &GenerateContent) -> &str { - match item { - GenerateContent::Text { text, .. } => text, - _ => panic!("expected Text content, got {item:?}"), - } -} - -async fn collect_stream(result: StreamResult) -> Vec { - let mut parts = Vec::new(); - let mut stream = result.stream; - while let Some(part) = stream.next().await { - match part { - Ok(p) => parts.push(p), - Err(e) => panic!("stream error: {e:?}"), - } - } - parts -} - -fn text_response(text: &str) -> Value { - json!({ - "id": "msg_017TfcQ4AgGxKyBduUpqYPZn", - "type": "message", - "role": "assistant", - "content": [{ "type": "text", "text": text }], - "model": "claude-sonnet-4-20250514", - "stop_reason": "end_turn", - "stop_sequence": null, - "usage": { "input_tokens": 4, "output_tokens": 30 } - }) -} - -// ═════════════════════════════════════════════════════════════════════════════ -// doGenerate tests -// ═════════════════════════════════════════════════════════════════════════════ - -#[tokio::test] -async fn vertex_anthropic_generate_keeps_direct_caller_metadata() { - let server = MockServer::start().await; - mock_raw_predict_json( - &server, - 200, - json!({ - "id": "msg_direct", - "type": "message", - "role": "assistant", - "content": [{ - "type": "tool_use", - "id": "toolu_direct", - "name": "get_weather", - "input": { "city": "Tokyo" }, - "caller": { "type": "direct" }, - }], - "model": MODEL_ID, - "stop_reason": "tool_use", - "usage": { "input_tokens": 100, "output_tokens": 50 }, - }), - ) - .await; - - let result = make_model(&server) - .do_generate(&default_options(test_prompt())) - .await - .unwrap(); - let metadata = result.content.iter().find_map(|part| match part { - GenerateContent::ToolCall(RawToolCall { - provider_metadata, .. - }) => provider_metadata.as_ref(), - _ => None, - }); - assert_eq!( - metadata.expect("caller metadata")["anthropic"]["caller"], - json!({ "type": "direct" }) - ); -} - -/// Test: the request hits the `/publishers/anthropic/models/{model}:rawPredict` -/// path (i.e. the `/publishers/google` suffix is replaced) and carries the -/// Bearer token Authorization header. -#[tokio::test] -async fn vertex_anthropic_url_and_auth() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path(RAW_PREDICT_PATH)) - .and(header("Authorization", "Bearer test-token")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_response("OK"))) - .mount(&server) - .await; - - let model = make_model(&server); - let result = model - .do_generate(&default_options(test_prompt())) - .await - .expect("do_generate should succeed"); - - assert_eq!(as_text(&result.content[0]), "OK"); - - // The single recorded request hit the Anthropic publisher path. - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - assert_eq!(requests[0].url.path(), RAW_PREDICT_PATH); - assert_eq!( - requests[0].headers.get("authorization").unwrap(), - "Bearer test-token" - ); -} - -// ═════════════════════════════════════════════════════════════════════════════ -// doStream tests -// ═════════════════════════════════════════════════════════════════════════════ - -/// Test: streaming text via the Anthropic SSE events returned by -/// streamRawPredict. -#[tokio::test] -async fn vertex_anthropic_stream_text() { - let server = MockServer::start().await; - - let sse_body = sse_stream(&[ - json!({ - "type": "message_start", - "message": { - "id": "msg_001", - "model": "claude-sonnet-4-20250514", - "usage": { "input_tokens": 3 } - } - }), - json!({ - "type": "content_block_start", - "index": 0, - "content_block": { "type": "text", "text": "" } - }), - json!({ - "type": "content_block_delta", - "index": 0, - "delta": { "type": "text_delta", "text": "Hello" } - }), - json!({ - "type": "content_block_delta", - "index": 0, - "delta": { "type": "text_delta", "text": " streamed!" } - }), - json!({ - "type": "content_block_stop", - "index": 0 - }), - json!({ - "type": "message_delta", - "delta": { "stop_reason": "end_turn" }, - "usage": { "output_tokens": 5 } - }), - json!({ "type": "message_stop" }), - ]); - - mock_stream_raw_predict_sse(&server, &sse_body).await; - - let model = make_model(&server); - let result = model - .do_stream(&default_options(test_prompt())) - .await - .expect("do_stream should succeed"); - - let parts = collect_stream(result).await; - - let text_deltas: Vec = parts - .iter() - .filter_map(|p| match p { - StreamPart::TextDelta { delta, .. } => Some(delta.clone()), - _ => None, - }) - .collect(); - assert_eq!(text_deltas, vec!["Hello", " streamed!"]); - - let finish = parts.iter().find_map(|p| match p { - StreamPart::Finish { finish_reason, .. } => Some(finish_reason.clone()), - _ => None, - }); - let finish = finish.expect("should have a Finish part"); - assert_eq!(finish.unified, FinishReasonUnified::Stop); - - // The stream request hit the streamRawPredict path. - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - assert_eq!(requests[0].url.path(), STREAM_RAW_PREDICT_PATH); - let wire: Value = serde_json::from_slice(&requests[0].body).expect("valid json body"); - assert_eq!(wire["anthropic_version"], "vertex-2023-10-16"); - assert_eq!(wire["stream"], true); -} - -// ═════════════════════════════════════════════════════════════════════════════ - -#[tokio::test] -async fn vertex_anthropic_stream_keeps_programmatic_caller_metadata() { - let server = MockServer::start().await; - let sse_body = sse_stream(&[ - json!({ - "type": "message_start", - "message": { - "id": "msg_programmatic", - "model": MODEL_ID, - "usage": { "input_tokens": 100 }, - }, - }), - json!({ - "type": "content_block_start", - "index": 0, - "content_block": { - "type": "tool_use", - "id": "toolu_programmatic", - "name": "query_database", - "input": { "sql": "SELECT 1" }, - "caller": { - "type": "code_execution_20250825", - "tool_id": "srvtoolu_code", - }, - }, - }), - json!({ "type": "content_block_stop", "index": 0 }), - json!({ - "type": "message_delta", - "delta": { "stop_reason": "tool_use" }, - "usage": { "output_tokens": 50 }, - }), - json!({ "type": "message_stop" }), - ]); - mock_stream_raw_predict_sse(&server, &sse_body).await; - let model = make_model(&server); - - let parts = collect_stream( - model - .do_stream(&default_options(test_prompt())) - .await - .unwrap(), - ) - .await; - let metadata = parts.iter().find_map(|part| match part { - StreamPart::ToolCall(RawToolCall { - provider_metadata, .. - }) => provider_metadata.as_ref(), - _ => None, - }); - assert_eq!( - metadata.expect("caller metadata")["anthropic"]["caller"], - json!({ - "type": "code_execution_20250825", - "toolId": "srvtoolu_code", - }) - ); -} - -// Error handling -// ═════════════════════════════════════════════════════════════════════════════ - -/// Test: a 401 response maps to `AiMuxError::ApiCall` (401 in `status_code`). -#[tokio::test] -async fn vertex_anthropic_generate_auth_error() { - let server = MockServer::start().await; - mock_raw_predict_json( - &server, - 401, - json!({ - "error": { - "code": 401, - "message": "Request had invalid authentication credentials.", - "status": "UNAUTHENTICATED" - } - }), - ) - .await; - - let model = make_model(&server); - let result = model.do_generate(&default_options(test_prompt())).await; - assert!( - matches!(result, Err(ref e) if e.status_code() == Some(401)), - "expected Auth, got {result:?}" - ); -} - -/// Test: a generic HTTP error (500) maps to `AiMuxError::ApiCall`. -#[tokio::test] -async fn vertex_anthropic_generate_provider_error() { - let server = MockServer::start().await; - mock_raw_predict_json( - &server, - 500, - json!({ - "error": { - "code": 500, - "message": "Internal error.", - "status": "INTERNAL" - } - }), - ) - .await; - - let model = make_model(&server); - let result = model.do_generate(&default_options(test_prompt())).await; - assert!( - matches!(result, Err(AiMuxError::ApiCall(_))), - "expected Provider, got {result:?}" - ); -} diff --git a/aimux-providers/tests/vertex_remaining_test.rs b/aimux-providers/tests/vertex_remaining_test.rs deleted file mode 100644 index 18eaf7874..000000000 --- a/aimux-providers/tests/vertex_remaining_test.rs +++ /dev/null @@ -1,504 +0,0 @@ -//! Remaining Google Vertex AI provider tests — ported from the TS SDK suite. -//! -//! Mirrors `reference/ai/packages/google-vertex/src/google-vertex-provider.test.ts` -//! (provider configuration: auth headers, base URL, env-var resolution, project -//! handling, Express-mode API key, tuned-model restrictions). -//! -//! The TS tests mock `createAuthTokenGenerator` / `createGoogleVertex` and assert -//! on the options passed to the base provider. The Rust provider has no lazy -//! token generator — a bearer token (or API key) is supplied directly to -//! `VertexProviderConfig` — so each TS scenario is translated to the equivalent -//! observable Rust behaviour: which auth variant is selected, which headers are -//! sent, and how the project/location shape the base URL. -//! -//! Model-level generate/stream behaviour is already covered by -//! `vertex_model_test.rs`; this file focuses on provider configuration only. - -use serde_json::json; -use serial_test::serial; -use wiremock::matchers::{header, method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::content::ContentPart; -use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; -use aimux_core::message::Role; -use aimux_core::options::CallOptions; -use aimux_core::provider::Provider; -use aimux_providers::vertex::{VertexAuth, VertexProvider, VertexProviderConfig}; - -// ── Shared helpers ─────────────────────────────────────────────────────────── - -/// All environment variables consulted by `VertexProviderConfig::from_env`. -const ENV_VARS: &[&str] = &[ - "GOOGLE_VERTEX_API_KEY", - "GOOGLE_VERTEX_ACCESS_TOKEN", - "GOOGLE_VERTEX_PROJECT", - "GOOGLE_VERTEX_LOCATION", -]; - -/// Remove every Vertex env var (test isolation for `#[serial]` env tests). -fn clear_vertex_env() { - for var in ENV_VARS { - unsafe { - std::env::remove_var(var); - } - } -} - -fn test_prompt() -> LanguageModelPrompt { - vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("Hello")], - ..Default::default() - }] -} - -fn default_options(prompt: LanguageModelPrompt) -> CallOptions { - CallOptions::new(prompt) -} - -/// Mount a minimal 200 generateContent mock on the standard model path. -async fn mock_ok(server: &MockServer) { - Mock::given(method("POST")) - .and(path("/models/gemini-2.0-flash:generateContent")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "candidates": [{ - "content": { "parts": [{ "text": "hi" }], "role": "model" }, - "finishReason": "STOP", "index": 0 - }] - }))) - .mount(server) - .await; -} - -// ════════════════════════════════════════════════════════════════════════════ -// Provider identity (TS: google-vertex-provider.test.ts — base wiring) -// ════════════════════════════════════════════════════════════════════════════ - -#[tokio::test] -async fn provider_name_is_google_vertex() { - let config = VertexProviderConfig::new("test-token", "my-project", "us-central1"); - let provider = VertexProvider::new(config); - assert_eq!(provider.name(), "google.vertex"); -} - -#[tokio::test] -async fn language_model_trait_method_returns_boxed_model() { - let config = VertexProviderConfig::new("test-token", "my-project", "us-central1"); - let provider = VertexProvider::new(config); - let model = provider.language_model("gemini-2.0-flash").expect("model"); - assert_eq!(model.model_id(), "gemini-2.0-flash"); - assert_eq!(model.provider(), "google.vertex"); -} - -#[tokio::test] -async fn model_factory_returns_vertex_model() { - let config = VertexProviderConfig::new("test-token", "my-project", "us-central1"); - let provider = VertexProvider::new(config); - let model = provider.model("gemini-2.0-flash").expect("model"); - assert_eq!(model.model_id(), "gemini-2.0-flash"); - assert_eq!(model.provider(), "google.vertex"); -} - -/// TS: "creates the auth token generator once per provider instance" — a single -/// provider instance can mint multiple models that all share its config. -#[tokio::test] -async fn one_provider_instance_creates_multiple_models() { - let config = VertexProviderConfig::new("test-token", "my-project", "us-central1"); - let provider = VertexProvider::new(config); - - let m1 = provider.model("gemini-2.0-flash").expect("model 1"); - let m2 = provider.model("gemini-2.5-pro").expect("model 2"); - - assert_eq!(m1.model_id(), "gemini-2.0-flash"); - assert_eq!(m2.model_id(), "gemini-2.5-pro"); - assert_eq!(m1.provider(), "google.vertex"); - assert_eq!(m2.provider(), "google.vertex"); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Auth headers (TS: "default headers function should return auth token", -// "should use custom headers in addition to auth token", -// "should pass options through to base provider when apiKey -// is provided") -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: bearer-token auth resolves to an `Authorization: Bearer {token}` header. -#[tokio::test] -async fn bearer_token_sent_via_authorization_header() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/models/gemini-2.0-flash:generateContent")) - .and(header("authorization", "Bearer my-bearer-token")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "candidates": [{ - "content": { "parts": [{ "text": "hi" }], "role": "model" }, - "finishReason": "STOP", "index": 0 - }] - }))) - .mount(&server) - .await; - - let config = VertexProviderConfig::new("my-bearer-token", "my-project", "us-central1") - .with_base_url(server.uri()); - let provider = VertexProvider::new(config); - let model = provider.model("gemini-2.0-flash").expect("model"); - - model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed — mock requires the Authorization: Bearer header"); -} - -/// TS: custom headers are sent alongside the auth token. -#[tokio::test] -async fn custom_headers_sent_alongside_bearer_token() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/models/gemini-2.0-flash:generateContent")) - .and(header("authorization", "Bearer my-bearer-token")) - .and(header("custom-header", "custom-value")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "candidates": [{ - "content": { "parts": [{ "text": "hi" }], "role": "model" }, - "finishReason": "STOP", "index": 0 - }] - }))) - .mount(&server) - .await; - - let config = VertexProviderConfig::new("my-bearer-token", "my-project", "us-central1") - .with_base_url(server.uri()); - let provider = VertexProvider::new(config); - let model = provider.model("gemini-2.0-flash").expect("model"); - - let mut opts = default_options(test_prompt()); - let mut headers = std::collections::HashMap::new(); - headers.insert("Custom-Header".to_string(), "custom-value".to_string()); - opts.headers = Some(headers); - - model - .do_generate(&opts) - .await - .expect("should succeed — mock requires both headers"); -} - -/// TS: "should pass options through to base provider when apiKey is provided" — -/// Express-mode API key uses `x-goog-api-key` (and does NOT send an -/// `Authorization: Bearer` header, i.e. the token generator is not invoked). -#[tokio::test] -async fn api_key_uses_x_goog_api_key_header() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/models/gemini-2.0-flash:generateContent")) - .and(header("x-goog-api-key", "express-api-key")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "candidates": [{ - "content": { "parts": [{ "text": "hi" }], "role": "model" }, - "finishReason": "STOP", "index": 0 - }] - }))) - .mount(&server) - .await; - - let config = VertexProviderConfig::with_api_key("express-api-key").with_base_url(server.uri()); - let provider = VertexProvider::new(config); - let model = provider.model("gemini-2.0-flash").expect("model"); - - model - .do_generate(&default_options(test_prompt())) - .await - .expect("should succeed — mock requires the x-goog-api-key header"); -} - -/// Express-mode config selects the `ApiKey` auth variant (no bearer token). -#[tokio::test] -async fn with_api_key_selects_apikey_auth_variant() { - let config = VertexProviderConfig::with_api_key("express-key"); - match &config.auth { - VertexAuth::ApiKey(k) => assert_eq!(k, "express-key"), - other => panic!("expected ApiKey auth, got {other:?}"), - } -} - -/// Bearer-token config selects the `BearerToken` auth variant. -#[tokio::test] -async fn new_selects_bearer_token_auth_variant() { - let config = VertexProviderConfig::new("the-token", "proj", "us-central1"); - match &config.auth { - VertexAuth::BearerToken(t) => assert_eq!(t, "the-token"), - other => panic!("expected BearerToken auth, got {other:?}"), - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// Base URL construction (TS: project / location handling) -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: the project is threaded into the base URL (projectId → URL path). -#[tokio::test] -async fn new_embeds_project_in_base_url() { - let config = VertexProviderConfig::new("token", "my-gcp-project", "us-central1"); - assert!( - config.base_url.contains("/projects/my-gcp-project/"), - "base_url should embed the project: {}", - config.base_url - ); -} - -/// TS: the location is threaded into the base URL. -#[tokio::test] -async fn new_embeds_location_in_base_url() { - let config = VertexProviderConfig::new("token", "proj", "us-central1"); - assert!( - config.base_url.contains("/locations/us-central1/"), - "base_url should embed the location: {}", - config.base_url - ); -} - -/// Location `us-central1` → `{location}-aiplatform.googleapis.com` host. -#[tokio::test] -async fn base_url_for_regional_location() { - let config = VertexProviderConfig::new("token", "proj", "us-central1"); - assert_eq!( - config.base_url, - "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/proj/locations/us-central1/publishers/google" - ); -} - -/// Location `global` → bare `aiplatform.googleapis.com` host. -#[tokio::test] -async fn base_url_for_global_location() { - let config = VertexProviderConfig::new("token", "proj", "global"); - assert_eq!( - config.base_url, - "https://aiplatform.googleapis.com/v1beta1/projects/proj/locations/global/publishers/google" - ); -} - -/// Location `eu`/`us` → multi-region `aiplatform.{region}.rep.googleapis.com`. -#[tokio::test] -async fn base_url_for_eu_multi_region() { - let config = VertexProviderConfig::new("token", "proj", "eu"); - assert_eq!( - config.base_url, - "https://aiplatform.eu.rep.googleapis.com/v1beta1/projects/proj/locations/eu/publishers/google" - ); -} - -#[tokio::test] -async fn base_url_for_us_multi_region() { - let config = VertexProviderConfig::new("token", "proj", "us"); - assert_eq!( - config.base_url, - "https://aiplatform.us.rep.googleapis.com/v1beta1/projects/proj/locations/us/publishers/google" - ); -} - -/// Express-mode (API key) uses the `/v1/publishers/google` endpoint (no -/// project/location scoping). -#[tokio::test] -async fn express_mode_base_url() { - let config = VertexProviderConfig::with_api_key("key"); - assert_eq!( - config.base_url, - "https://aiplatform.googleapis.com/v1/publishers/google" - ); -} - -/// `with_base_url` overrides the constructed URL and strips a trailing slash. -#[tokio::test] -async fn with_base_url_strips_trailing_slash() { - let config = VertexProviderConfig::new("token", "proj", "us-central1") - .with_base_url("https://example.com/v1beta1/"); - assert_eq!(config.base_url, "https://example.com/v1beta1"); -} - -/// `with_base_url` is honoured end-to-end (request hits the override host). -#[tokio::test] -async fn with_base_url_is_used_for_requests() { - let server = MockServer::start().await; - mock_ok(&server).await; - - let config = - VertexProviderConfig::new("token", "proj", "us-central1").with_base_url(server.uri()); - let provider = VertexProvider::new(config); - let model = provider.model("gemini-2.0-flash").expect("model"); - - model - .do_generate(&default_options(test_prompt())) - .await - .expect("request should hit the mock server"); -} - -// ════════════════════════════════════════════════════════════════════════════ -// Tuned-model restriction (TS: Express mode cannot address tuned endpoints) -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: a tuned model (`endpoints/…`) cannot be used with Express-mode API key -/// auth — the provider rejects it up front. -#[tokio::test] -async fn tuned_model_rejected_with_api_key_auth() { - let config = VertexProviderConfig::with_api_key("express-key"); - let provider = VertexProvider::new(config); - let result = provider.model("endpoints/1234567890"); - assert!( - result.is_err(), - "tuned models should be rejected under Express-mode API key auth" - ); -} - -/// Tuned models are allowed with standard (bearer-token) auth. -#[tokio::test] -async fn tuned_model_allowed_with_bearer_token_auth() { - let config = VertexProviderConfig::new("token", "proj", "us-central1"); - let provider = VertexProvider::new(config); - let model = provider - .model("endpoints/1234567890") - .expect("tuned models should be allowed with bearer-token auth"); - assert_eq!(model.model_id(), "endpoints/1234567890"); -} - -// ════════════════════════════════════════════════════════════════════════════ -// from_env (TS: env-var driven provider creation) -// ════════════════════════════════════════════════════════════════════════════ - -/// `GOOGLE_VERTEX_API_KEY` selects Express mode (API key auth). -#[tokio::test] -#[serial] -async fn from_env_uses_api_key_when_set() { - clear_vertex_env(); - unsafe { - std::env::set_var("GOOGLE_VERTEX_API_KEY", "env-api-key"); - } - let config = VertexProviderConfig::from_env().expect("from_env"); - match config.auth { - VertexAuth::ApiKey(k) => assert_eq!(k, "env-api-key"), - other => panic!("expected ApiKey auth, got {other:?}"), - } - clear_vertex_env(); -} - -/// `GOOGLE_VERTEX_API_KEY` takes precedence over the access-token vars. -#[tokio::test] -#[serial] -async fn from_env_prefers_api_key_over_access_token() { - clear_vertex_env(); - unsafe { - std::env::set_var("GOOGLE_VERTEX_API_KEY", "env-api-key"); - std::env::set_var("GOOGLE_VERTEX_ACCESS_TOKEN", "env-token"); - std::env::set_var("GOOGLE_VERTEX_PROJECT", "env-project"); - } - let config = VertexProviderConfig::from_env().expect("from_env"); - // API key wins → Express mode base URL (no project scoping). - assert!( - config.base_url.contains("/v1/publishers/google"), - "API key should select Express mode: {}", - config.base_url - ); - match config.auth { - VertexAuth::ApiKey(k) => assert_eq!(k, "env-api-key"), - other => panic!("expected ApiKey auth, got {other:?}"), - } - clear_vertex_env(); -} - -/// Access token + project + location → bearer-token auth with scoped base URL. -#[tokio::test] -#[serial] -async fn from_env_uses_access_token_project_location() { - clear_vertex_env(); - unsafe { - std::env::set_var("GOOGLE_VERTEX_ACCESS_TOKEN", "env-token"); - std::env::set_var("GOOGLE_VERTEX_PROJECT", "env-project"); - std::env::set_var("GOOGLE_VERTEX_LOCATION", "europe-west1"); - } - let config = VertexProviderConfig::from_env().expect("from_env"); - match config.auth { - VertexAuth::BearerToken(t) => assert_eq!(t, "env-token"), - other => panic!("expected BearerToken auth, got {other:?}"), - } - assert!(config.base_url.contains("/projects/env-project/")); - assert!(config.base_url.contains("/locations/europe-west1/")); - clear_vertex_env(); -} - -/// `GOOGLE_VERTEX_LOCATION` defaults to `us-central1` when unset. -#[tokio::test] -#[serial] -async fn from_env_defaults_location_to_us_central1() { - clear_vertex_env(); - unsafe { - std::env::set_var("GOOGLE_VERTEX_ACCESS_TOKEN", "env-token"); - std::env::set_var("GOOGLE_VERTEX_PROJECT", "env-project"); - } - let config = VertexProviderConfig::from_env().expect("from_env"); - assert!( - config.base_url.contains("/locations/us-central1/"), - "location should default to us-central1: {}", - config.base_url - ); - clear_vertex_env(); -} - -/// Missing every var → error. -#[tokio::test] -#[serial] -async fn from_env_errors_when_all_missing() { - clear_vertex_env(); - let result = VertexProviderConfig::from_env(); - assert!(result.is_err(), "from_env should error with no credentials"); - clear_vertex_env(); -} - -/// Access token present but project missing → error. -#[tokio::test] -#[serial] -async fn from_env_errors_when_project_missing() { - clear_vertex_env(); - unsafe { - std::env::set_var("GOOGLE_VERTEX_ACCESS_TOKEN", "env-token"); - } - let result = VertexProviderConfig::from_env(); - assert!( - result.is_err(), - "from_env should error when GOOGLE_VERTEX_PROJECT is missing" - ); - clear_vertex_env(); -} - -/// Project present but access token missing → error. -#[tokio::test] -#[serial] -async fn from_env_errors_when_access_token_missing() { - clear_vertex_env(); - unsafe { - std::env::set_var("GOOGLE_VERTEX_PROJECT", "env-project"); - } - let result = VertexProviderConfig::from_env(); - assert!( - result.is_err(), - "from_env should error when GOOGLE_VERTEX_ACCESS_TOKEN is missing" - ); - clear_vertex_env(); -} - -/// An empty `GOOGLE_VERTEX_API_KEY` falls through to access-token auth. -#[tokio::test] -#[serial] -async fn from_env_empty_api_key_falls_through_to_access_token() { - clear_vertex_env(); - unsafe { - std::env::set_var("GOOGLE_VERTEX_API_KEY", " "); - std::env::set_var("GOOGLE_VERTEX_ACCESS_TOKEN", "env-token"); - std::env::set_var("GOOGLE_VERTEX_PROJECT", "env-project"); - } - let config = VertexProviderConfig::from_env().expect("from_env"); - match config.auth { - VertexAuth::BearerToken(t) => assert_eq!(t, "env-token"), - other => panic!("expected BearerToken auth, got {other:?}"), - } - clear_vertex_env(); -} diff --git a/bindings/node/src/types.ts b/bindings/node/src/types.ts index 58bc4f504..60b24edcf 100644 --- a/bindings/node/src/types.ts +++ b/bindings/node/src/types.ts @@ -5,6 +5,7 @@ export type * from './types/AiMuxError' export type * from './types/ApiCallError' export type * from './types/AspectRatio' +export type * from './types/AssistantPart' export type * from './types/AudioChunk' export type * from './types/AudioData' export type * from './types/AudioInput' @@ -24,6 +25,7 @@ export type * from './types/ChatCompletionToolCall' export type * from './types/ChatCompletionUsage' export type * from './types/CompletionTokensDetails' export type * from './types/ContentPart' +export type * from './types/CustomPart' export type * from './types/DecisionAnswer' export type * from './types/DecisionBooleanCriteria' export type * from './types/DecisionCallOptions' @@ -41,6 +43,7 @@ export type * from './types/EmbeddingResult' export type * from './types/EmbeddingUsage' export type * from './types/FileBytes' export type * from './types/FileData' +export type * from './types/FilePart' export type * from './types/Fingerprint' export type * from './types/FinishReason' export type * from './types/FinishReasonUnified' @@ -66,7 +69,8 @@ export type * from './types/ImageUsage' export type * from './types/InputAudioFormat' export type * from './types/InputRecord' export type * from './types/InputTokenUsage' -export type * from './types/LanguageModelPromptMessage' +export type * from './types/LanguageModelMessage' +export type * from './types/LanguageModelReasoningPart' export type * from './types/MessageContent' export type * from './types/Modality' export type * from './types/ModelCapabilities' @@ -91,6 +95,7 @@ export type * from './types/RawToolApprovalRequest' export type * from './types/RawToolCall' export type * from './types/ReasoningEffort' export type * from './types/ReasoningFileOutput' +export type * from './types/ReasoningFilePart' export type * from './types/ReasoningMode' export type * from './types/ReasoningOutput' export type * from './types/ReasoningPart' @@ -132,14 +137,21 @@ export type * from './types/StreamPart' export type * from './types/StreamResponse' export type * from './types/StreamTextResultAggregated' export type * from './types/TextContent' +export type * from './types/TextPart' export type * from './types/TimeoutConfiguration' export type * from './types/TimingRecord' export type * from './types/Tool' export type * from './types/ToolApprovalRequestOutput' +export type * from './types/ToolApprovalResponsePart' export type * from './types/ToolCall' +export type * from './types/ToolCallPart' export type * from './types/ToolCallRepairReply' export type * from './types/ToolChoice' +export type * from './types/ToolPart' export type * from './types/ToolResult' +export type * from './types/ToolResultContent' +export type * from './types/ToolResultOutput' +export type * from './types/ToolResultPart' export type * from './types/TraceFilter' export type * from './types/TraceRecord' export type * from './types/TraceStats' @@ -154,6 +166,7 @@ export type * from './types/UploadFileData' export type * from './types/UploadFileResult' export type * from './types/Usage' export type * from './types/UsageSnapshot' +export type * from './types/UserPart' export type * from './types/Verdict' export type * from './types/VerdictConfidence' export type * from './types/VerdictKind' diff --git a/bindings/node/src/types/AssistantPart.ts b/bindings/node/src/types/AssistantPart.ts new file mode 100644 index 000000000..850d5780e --- /dev/null +++ b/bindings/node/src/types/AssistantPart.ts @@ -0,0 +1,13 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { CustomPart } from "./CustomPart"; +import type { FilePart } from "./FilePart"; +import type { LanguageModelReasoningPart } from "./LanguageModelReasoningPart"; +import type { ReasoningFilePart } from "./ReasoningFilePart"; +import type { TextPart } from "./TextPart"; +import type { ToolCallPart } from "./ToolCallPart"; +import type { ToolResultPart } from "./ToolResultPart"; + +/** + * Parts allowed in an assistant message. + */ +export type AssistantPart = { "type": "text" } & TextPart | { "type": "file" } & FilePart | { "type": "reasoning" } & LanguageModelReasoningPart | { "type": "reasoning-file" } & ReasoningFilePart | { "type": "custom" } & CustomPart | { "type": "tool-call" } & ToolCallPart | { "type": "tool-result" } & ToolResultPart; diff --git a/bindings/node/src/types/CallOptions.ts b/bindings/node/src/types/CallOptions.ts index 3f8c21974..db98aa764 100644 --- a/bindings/node/src/types/CallOptions.ts +++ b/bindings/node/src/types/CallOptions.ts @@ -1,5 +1,5 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -import type { LanguageModelPromptMessage } from "./LanguageModelPromptMessage"; +import type { LanguageModelMessage } from "./LanguageModelMessage"; import type { ReasoningEffort } from "./ReasoningEffort"; import type { ResponseFormat } from "./ResponseFormat"; import type { TimeoutConfiguration } from "./TimeoutConfiguration"; @@ -18,7 +18,7 @@ export type CallOptions = { /** * The standardized prompt (message array). Required. */ -prompt: Array, +prompt: Array, /** * Maximum tokens to generate. */ diff --git a/bindings/node/src/types/ContentPart.ts b/bindings/node/src/types/ContentPart.ts index bd1ddab7e..fd965cc5e 100644 --- a/bindings/node/src/types/ContentPart.ts +++ b/bindings/node/src/types/ContentPart.ts @@ -7,32 +7,32 @@ import type { JsonValue } from "./serde_json/JsonValue"; * * Shared between `ModelMessage` (user-facing) and `LanguageModelPrompt` (provider-facing). */ -export type ContentPart = { "type": "text", text: string, +export type ContentPart = { "type": "text", text: string, /** * Provider-specific options for this part (e.g. `openai.promptCacheBreakpoint`). */ -provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "custom", kind: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "reasoning_file", data: GeneratedFileData, media_type: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "tool_approval_request", approval_id: string, tool_call_id: string, reason?: string | null, is_automatic?: boolean | null, signature?: string | null, input_schema_input?: JsonValue | null, } | { "type": "image", image: Array, media_type: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "file", data: Array, media_type: string, filename?: string | null, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "file_base64", data: string, media_type: string, filename?: string | null, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "file_url", url: string, media_type: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "file_reference", media_type: string, reference: JsonValue, filename?: string | null, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "reasoning", text: string, signature: string | null, +provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "custom", kind: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "reasoning_file", data: GeneratedFileData, media_type: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "tool_approval_request", approval_id: string, tool_call_id: string, reason?: string | null, is_automatic?: boolean | null, signature?: string | null, input_schema_input?: JsonValue | null, } | { "type": "image", image: Array, media_type: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "file", data: Array, media_type: string, filename?: string | null, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "file_base64", data: string, media_type: string, filename?: string | null, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "file_url", url: string, media_type: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "file_reference", media_type: string, reference: JsonValue, filename?: string | null, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "reasoning", text: string, signature: string | null, /** * Provider-specific options for this part (e.g. * `anthropic.signature`, `anthropic.redactedData`, * `anthropic.cacheControl`). */ -provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "tool_call", tool_call_id: string, tool_name: string, +provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "tool_call", tool_call_id: string, tool_name: string, /** * Arguments as a JSON value (usually an object). */ -input: JsonValue, +input: JsonValue, /** * Whether the tool call is executed by the provider rather than by * the client. This is part of the prompt contract because providers * need it to replay server tool calls on a later turn. */ -provider_executed?: boolean | null, +provider_executed?: boolean | null, /** * Provider-specific options for this part (e.g. * `anthropic.cacheControl`). */ -provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "tool_result", tool_call_id: string, +provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "tool_result", tool_call_id: string, /** * The tool's output (usually a JSON value or plain text). * @@ -41,25 +41,25 @@ provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } * so multi-part tool messages constructed either way round-trip * correctly. Serialization always emits `result`. */ -result: JsonValue, +result: JsonValue, /** * The name of the tool that produced this result (optional on the * user-input side; providers that need it can look it up from the * preceding tool call). */ -tool_name?: string | null, +tool_name?: string | null, /** * Whether the result is an error or error message. */ -is_error?: boolean | null, +is_error?: boolean | null, /** * Whether the result is preliminary (replaces prior, e.g. image previews). */ -preliminary?: boolean | null, +preliminary?: boolean | null, /** * Whether the tool is dynamic (defined at runtime, e.g. MCP tools). */ -dynamic?: boolean | null, +dynamic?: boolean | null, /** * Provider-specific options for this part (e.g. * `anthropic.cacheControl`). diff --git a/bindings/node/src/types/CustomPart.ts b/bindings/node/src/types/CustomPart.ts new file mode 100644 index 000000000..39f66564d --- /dev/null +++ b/bindings/node/src/types/CustomPart.ts @@ -0,0 +1,7 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * Provider-specific content identified by its provider and kind. + */ +export type CustomPart = { kind: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/FilePart.ts b/bindings/node/src/types/FilePart.ts new file mode 100644 index 000000000..e412d1c0a --- /dev/null +++ b/bindings/node/src/types/FilePart.ts @@ -0,0 +1,8 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { FileData } from "./FileData"; +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * File content part (the only file part; `data` is a tagged `FileData`). + */ +export type FilePart = { data: FileData, media_type: string, filename?: string | null, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/InputRecord.ts b/bindings/node/src/types/InputRecord.ts index c6e44df47..12ea93e86 100644 --- a/bindings/node/src/types/InputRecord.ts +++ b/bindings/node/src/types/InputRecord.ts @@ -1,5 +1,5 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -import type { LanguageModelPromptMessage } from "./LanguageModelPromptMessage"; +import type { LanguageModelMessage } from "./LanguageModelMessage"; import type { RecordingOperation } from "./RecordingOperation"; import type { JsonValue } from "./serde_json/JsonValue"; @@ -10,7 +10,7 @@ export type InputRecord = { operation?: RecordingOperation, /** * 完整 prompt(消息数组,含 ContentPart::Image 等多模态)。 */ -prompt: Array, +prompt: Array, /** * 序列化的 CallOptions(abort_signal/call_id 已 serde skip); * headers/provider_options/body_overrides 已递归脱敏。 diff --git a/bindings/node/src/types/LanguageModelMessage.ts b/bindings/node/src/types/LanguageModelMessage.ts new file mode 100644 index 000000000..a7052a024 --- /dev/null +++ b/bindings/node/src/types/LanguageModelMessage.ts @@ -0,0 +1,13 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { AssistantPart } from "./AssistantPart"; +import type { ToolPart } from "./ToolPart"; +import type { UserPart } from "./UserPart"; +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * A single provider-facing message, a union by role. + * + * `provider_options` is message-level (e.g. `anthropic.cacheControl`), + * mirroring `LanguageModelV4Message.providerOptions`. + */ +export type LanguageModelMessage = { "role": "system", content: string, provider_options: { [key in string]: { [key in string]: JsonValue } } | null, } | { "role": "user", content: Array, provider_options: { [key in string]: { [key in string]: JsonValue } } | null, } | { "role": "assistant", content: Array, provider_options: { [key in string]: { [key in string]: JsonValue } } | null, } | { "role": "tool", content: Array, provider_options: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/LanguageModelPromptMessage.ts b/bindings/node/src/types/LanguageModelPromptMessage.ts deleted file mode 100644 index 459e4fd5b..000000000 --- a/bindings/node/src/types/LanguageModelPromptMessage.ts +++ /dev/null @@ -1,19 +0,0 @@ -// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -import type { ContentPart } from "./ContentPart"; -import type { Role } from "./Role"; -import type { JsonValue } from "./serde_json/JsonValue"; - -/** - * A single provider-facing message. - * - * Unlike `ModelMessage` (which can have string content), this always has - * `content: Vec` — strings are normalized to `ContentPart::Text`. - */ -export type LanguageModelPromptMessage = { role: Role, content: Array, -/** - * Message-level provider-specific options (e.g. - * `anthropic.cacheControl`), applied to the last part of the message when - * converting to a provider prompt. Mirrors the Vercel AI SDK - * `LanguageModelV4Message.providerOptions`. - */ -provider_options: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/LanguageModelReasoningPart.ts b/bindings/node/src/types/LanguageModelReasoningPart.ts new file mode 100644 index 000000000..4fb52fa16 --- /dev/null +++ b/bindings/node/src/types/LanguageModelReasoningPart.ts @@ -0,0 +1,7 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * Reasoning / thinking content part. + */ +export type LanguageModelReasoningPart = { text: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/ReasoningFilePart.ts b/bindings/node/src/types/ReasoningFilePart.ts new file mode 100644 index 000000000..b8c1e4aab --- /dev/null +++ b/bindings/node/src/types/ReasoningFilePart.ts @@ -0,0 +1,8 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { GeneratedFileData } from "./GeneratedFileData"; +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * A file generated as part of reasoning. + */ +export type ReasoningFilePart = { data: GeneratedFileData, media_type: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/StreamPart.ts b/bindings/node/src/types/StreamPart.ts index adab7631d..a7e850989 100644 --- a/bindings/node/src/types/StreamPart.ts +++ b/bindings/node/src/types/StreamPart.ts @@ -19,27 +19,27 @@ import type { JsonValue } from "./serde_json/JsonValue"; * [`StreamPart`] (raw argument text); `stream_text` emits [`TextStreamPart`] * (parsed and validated [`ToolCall`]). */ -export type StreamPart = { "TextStart": { id: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "TextDelta": { id: string, delta: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "TextEnd": { id: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "StreamStart": { warnings: Array, } } | { "Finish": { finish_reason: FinishReason, usage: Usage, provider_metadata: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "Error": { error: AiMuxError, } } | { "ToolInputStart": { id: string, tool_name: string, +export type StreamPart = { "TextStart": { id: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "TextDelta": { id: string, delta: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "TextEnd": { id: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "StreamStart": { warnings: Array, } } | { "Finish": { finish_reason: FinishReason, usage: Usage, provider_metadata: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "Error": { error: AiMuxError, } } | { "ToolInputStart": { id: string, tool_name: string, /** * Whether the tool call will be executed by the provider. */ -provider_executed?: boolean | null, +provider_executed?: boolean | null, /** * Whether the tool is dynamic (defined at runtime, e.g. MCP tools). */ -dynamic?: boolean | null, +dynamic?: boolean | null, /** * Optional title for the tool call. */ -title?: string | null, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ToolInputDelta": { id: string, delta: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ToolInputEnd": { id: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ToolCall": ToolCall } | { "ToolResult": ToolResult } | { "File": GeneratedFile } | { "ReasoningFile": ReasoningFileOutput } | { "Custom": { kind: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ToolApprovalRequest": ToolApprovalRequestOutput } | { "ReasoningStart": { id: string, +title?: string | null, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ToolInputDelta": { id: string, delta: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ToolInputEnd": { id: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ToolCall": ToolCall } | { "ToolResult": ToolResult } | { "File": GeneratedFile } | { "ReasoningFile": ReasoningFileOutput } | { "Custom": { kind: string, provider_metadata?: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ToolApprovalRequest": ToolApprovalRequestOutput } | { "ReasoningStart": { id: string, /** * Provider-specific metadata (e.g. xAI `itemId`). */ -provider_metadata: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ReasoningDelta": { id: string, delta: string, +provider_metadata: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ReasoningDelta": { id: string, delta: string, /** * Provider-specific metadata (e.g. xAI `itemId`). */ -provider_metadata: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ReasoningEnd": { id: string, +provider_metadata: { [key in string]: { [key in string]: JsonValue } } | null, } } | { "ReasoningEnd": { id: string, /** * Provider-specific metadata (e.g. xAI `itemId`, `reasoningEncryptedContent`). */ diff --git a/bindings/node/src/types/TextPart.ts b/bindings/node/src/types/TextPart.ts new file mode 100644 index 000000000..2e01804cf --- /dev/null +++ b/bindings/node/src/types/TextPart.ts @@ -0,0 +1,7 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * Text content part. + */ +export type TextPart = { text: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/ToolApprovalResponsePart.ts b/bindings/node/src/types/ToolApprovalResponsePart.ts new file mode 100644 index 000000000..12163b325 --- /dev/null +++ b/bindings/node/src/types/ToolApprovalResponsePart.ts @@ -0,0 +1,7 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * The user's decision for a provider-executed tool approval request. + */ +export type ToolApprovalResponsePart = { approval_id: string, approved: boolean, reason?: string | null, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/ToolCallPart.ts b/bindings/node/src/types/ToolCallPart.ts new file mode 100644 index 000000000..bac25885b --- /dev/null +++ b/bindings/node/src/types/ToolCallPart.ts @@ -0,0 +1,7 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * Tool call content part. + */ +export type ToolCallPart = { tool_call_id: string, tool_name: string, input: JsonValue, provider_executed?: boolean | null, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/ToolPart.ts b/bindings/node/src/types/ToolPart.ts new file mode 100644 index 000000000..df0fe1c1b --- /dev/null +++ b/bindings/node/src/types/ToolPart.ts @@ -0,0 +1,8 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { ToolApprovalResponsePart } from "./ToolApprovalResponsePart"; +import type { ToolResultPart } from "./ToolResultPart"; + +/** + * Parts allowed in a tool message. + */ +export type ToolPart = { "type": "tool-result" } & ToolResultPart | { "type": "tool-approval-response" } & ToolApprovalResponsePart; diff --git a/bindings/node/src/types/ToolResultContent.ts b/bindings/node/src/types/ToolResultContent.ts new file mode 100644 index 000000000..3511c045d --- /dev/null +++ b/bindings/node/src/types/ToolResultContent.ts @@ -0,0 +1,9 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { FilePart } from "./FilePart"; +import type { TextPart } from "./TextPart"; +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * Content within a tool result. + */ +export type ToolResultContent = { "type": "text" } & TextPart | { "type": "file" } & FilePart | { "type": "custom", provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/ToolResultOutput.ts b/bindings/node/src/types/ToolResultOutput.ts new file mode 100644 index 000000000..2855f2918 --- /dev/null +++ b/bindings/node/src/types/ToolResultOutput.ts @@ -0,0 +1,8 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { ToolResultContent } from "./ToolResultContent"; +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * The provider-facing result of a tool call. + */ +export type ToolResultOutput = { "type": "text", value: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "json", value: JsonValue, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "execution-denied", reason?: string | null, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "error-text", value: string, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "error-json", value: JsonValue, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, } | { "type": "content", value: Array, }; diff --git a/bindings/node/src/types/ToolResultPart.ts b/bindings/node/src/types/ToolResultPart.ts new file mode 100644 index 000000000..6a479932d --- /dev/null +++ b/bindings/node/src/types/ToolResultPart.ts @@ -0,0 +1,8 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { ToolResultOutput } from "./ToolResultOutput"; +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * Tool result content part. + */ +export type ToolResultPart = { tool_call_id: string, tool_name: string, output: ToolResultOutput, provider_options?: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/UserPart.ts b/bindings/node/src/types/UserPart.ts new file mode 100644 index 000000000..beecc56e5 --- /dev/null +++ b/bindings/node/src/types/UserPart.ts @@ -0,0 +1,8 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { FilePart } from "./FilePart"; +import type { TextPart } from "./TextPart"; + +/** + * Parts allowed in a user message. + */ +export type UserPart = { "type": "text" } & TextPart | { "type": "file" } & FilePart; diff --git a/tools/aimux-replay/src/main.rs b/tools/aimux-replay/src/main.rs index da111e534..dd095e062 100644 --- a/tools/aimux-replay/src/main.rs +++ b/tools/aimux-replay/src/main.rs @@ -18,6 +18,7 @@ use anyhow::{Context, Result, bail}; use clap::Parser; use tokio::runtime::Runtime; +use aimux_core::language_model_message::{AssistantPart, LanguageModelMessage, UserPart}; use aimux_core::recording::Recording; use aimux_core::replay::{ReplayOverrides, replay_with_model}; use aimux_providers::rebuild_provider; @@ -129,12 +130,18 @@ fn run_dry(rec: &Recording) -> Result<()> { } /// 取消息首个文本内容(截断展示用)。 -fn prompt_text( - m: &aimux_core::language_model_message::LanguageModelPromptMessage, -) -> Option { - match m.content.first() { - Some(aimux_core::content::ContentPart::Text { text, .. }) => Some(text.clone()), - _ => None, +fn prompt_text(m: &LanguageModelMessage) -> Option { + match m { + LanguageModelMessage::System { content, .. } => Some(content.clone()), + LanguageModelMessage::User { content, .. } => match content.first() { + Some(UserPart::Text(part)) => Some(part.text.clone()), + _ => None, + }, + LanguageModelMessage::Assistant { content, .. } => match content.first() { + Some(AssistantPart::Text(part)) => Some(part.text.clone()), + _ => None, + }, + LanguageModelMessage::Tool { .. } => None, } } @@ -185,13 +192,7 @@ fn run_replay( }; let overrides = prompt.map(|text| ReplayOverrides { - prompt: Some(vec![ - aimux_core::language_model_message::LanguageModelPromptMessage { - role: aimux_core::message::Role::User, - content: vec![aimux_core::content::ContentPart::text(text)], - provider_options: None, - }, - ]), + prompt: Some(vec![LanguageModelMessage::user_text(text)]), ..Default::default() }); diff --git a/tools/aimux-web/examples/gen_demo_fixtures.rs b/tools/aimux-web/examples/gen_demo_fixtures.rs index 8de6ccff3..071b57af9 100644 --- a/tools/aimux-web/examples/gen_demo_fixtures.rs +++ b/tools/aimux-web/examples/gen_demo_fixtures.rs @@ -9,9 +9,7 @@ //! //! The rest only enrich the Traces / Sessions listing. -use aimux_core::content::ContentPart; -use aimux_core::language_model_message::LanguageModelPromptMessage; -use aimux_core::message::Role; +use aimux_core::language_model_message::{AssistantPart, LanguageModelMessage, TextPart}; use aimux_core::options::CallOptions; use aimux_core::recording::{ HttpExchange, HttpRecord, InputRecord, OutcomeRecord, OutcomeStatus, ProviderRecord, Recording, @@ -21,21 +19,13 @@ use aimux_core::recording::{ const AGENT_SYSTEM: &str = "You are a helpful assistant. Always use the calculator tool for arithmetic, then answer."; -fn msg(role: Role, text: &str) -> LanguageModelPromptMessage { - LanguageModelPromptMessage { - role, - content: vec![ContentPart::text(text)], - provider_options: None, - } -} - #[allow(clippy::too_many_arguments)] fn recording( call_id: &str, recorded_at: &str, provider: &str, model_id: &str, - prompt: Vec, + prompt: Vec, session: Option<&str>, step: Option, status: OutcomeStatus, @@ -140,7 +130,7 @@ data: [DONE]"#; "2026-08-14T09:31:10.000Z", "openai", "gpt-4o", - vec![msg(Role::User, "what is 1 + 1?")], + vec![LanguageModelMessage::user_text("what is 1 + 1?")], Some("sess-playground"), Some(0), OutcomeStatus::Success, @@ -159,8 +149,11 @@ data: [DONE]"#; "openai", "gpt-4o", vec![ - msg(Role::System, AGENT_SYSTEM), - msg(Role::User, "what is 17 * 19?"), + LanguageModelMessage::System { + content: AGENT_SYSTEM.into(), + provider_options: None, + }, + LanguageModelMessage::user_text("what is 17 * 19?"), ], Some("sess-agent"), Some(0), @@ -180,8 +173,11 @@ data: [DONE]"#; "deepseek", "deepseek-chat", vec![ - msg(Role::System, "You are a careful technical analyst."), - msg(Role::User, "Explain how prefix caching works."), + LanguageModelMessage::System { + content: "You are a careful technical analyst.".into(), + provider_options: None, + }, + LanguageModelMessage::user_text("Explain how prefix caching works."), ], Some("sess-deepseek"), Some(0), @@ -200,10 +196,19 @@ data: [DONE]"#; "deepseek", "deepseek-chat", vec![ - msg(Role::System, "You are a careful technical analyst."), - msg(Role::User, "Explain how prefix caching works."), - msg(Role::Assistant, "Prefix caching reuses the KV cache..."), - msg(Role::User, "What is the hit rate formula?"), + LanguageModelMessage::System { + content: "You are a careful technical analyst.".into(), + provider_options: None, + }, + LanguageModelMessage::user_text("Explain how prefix caching works."), + LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: "Prefix caching reuses the KV cache...".into(), + provider_options: None, + })], + provider_options: None, + }, + LanguageModelMessage::user_text("What is the hit rate formula?"), ], Some("sess-deepseek"), Some(1), @@ -222,7 +227,7 @@ data: [DONE]"#; "2026-08-14T09:25:15.000Z", "groq", "llama-3.3-70b-versatile", - vec![msg(Role::User, "hi")], + vec![LanguageModelMessage::user_text("hi")], Some("sess-err"), Some(0), OutcomeStatus::Error, @@ -237,7 +242,7 @@ data: [DONE]"#; "2026-08-14T09:10:05.000Z", "anthropic", "claude-3-5-sonnet-latest", - vec![msg(Role::User, "Write a haiku about Rust.")], + vec![LanguageModelMessage::user_text("Write a haiku about Rust.")], None, None, OutcomeStatus::Success, diff --git a/tools/aimux-web/examples/gen_fixture.rs b/tools/aimux-web/examples/gen_fixture.rs index d9cdb1cb6..edf52366e 100644 --- a/tools/aimux-web/examples/gen_fixture.rs +++ b/tools/aimux-web/examples/gen_fixture.rs @@ -3,9 +3,7 @@ //! //! Usage: `cargo run -p aimux-web --example gen_fixture [out.jsonl]` -use aimux_core::content::ContentPart; -use aimux_core::language_model_message::LanguageModelPromptMessage; -use aimux_core::message::Role; +use aimux_core::language_model_message::LanguageModelMessage; use aimux_core::options::CallOptions; use aimux_core::recording::{ HttpExchange, HttpRecord, InputRecord, OutcomeRecord, OutcomeStatus, ProviderRecord, Recording, @@ -18,18 +16,11 @@ fn main() { .unwrap_or_else(|| "fixture.jsonl".to_string()); let prompt = vec![ - LanguageModelPromptMessage { - role: Role::System, - content: vec![ContentPart::text( - "You are a helpful assistant. Always use the calculator tool.", - )], - provider_options: None, - }, - LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("what is 1 + 1?")], + LanguageModelMessage::System { + content: "You are a helpful assistant. Always use the calculator tool.".into(), provider_options: None, }, + LanguageModelMessage::user_text("what is 1 + 1?"), ]; let options = CallOptions { diff --git a/tools/aimux-web/src/api/replay.rs b/tools/aimux-web/src/api/replay.rs index 776685f07..beb2a9d10 100644 --- a/tools/aimux-web/src/api/replay.rs +++ b/tools/aimux-web/src/api/replay.rs @@ -100,7 +100,7 @@ async fn run_inner(state: AppState, req: ReplayRequest) -> Result Result>(); - out.push(ModelMessage { - role, - content: MessageContent::Parts(parts), - }); + let content = if role == Role::System { + let [WireContentPart::Text { text }] = m.content.as_slice() else { + return Err(AiMuxError::InvalidPrompt( + "system content must be plain text".into(), + )); + }; + MessageContent::Text(text.clone()) + } else { + MessageContent::Parts(m.content.iter().map(to_content_part).collect()) + }; + out.push(ModelMessage { role, content }); } Ok(ModelPrompt::Messages(out)) }