From 0518502482899346740154ab8564a252ef0a18e5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 08:49:11 +0000 Subject: [PATCH 1/2] refactor(xai): move the xAI files to the upstream layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One Rust file per @ai-sdk/xai source file (RFC-0046 §4.2): mod.rs keeps only the module list (index.ts), with the provider in provider.rs and the error and response handlers in error.rs; convert.rs becomes supports_reasoning_effort.rs; speech, transcription and video become *_model.rs. responses/mod.rs becomes responses_language_model.rs (with the request body and tool-name helpers), types.rs becomes responses_api.rs, and responses/convert.rs splits into convert_to_responses_input.rs, convert_responses_usage.rs, map_responses_finish_reason.rs, responses_prepare_tools.rs and responses_language_model_options.rs. The four tests/xai_*_test.rs binaries join their files in src/xai/tests/. Move only: the test count (338) and the test names are unchanged; function bodies change only in module paths. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0116HPMoqaLfFr1TyqDq1hrW --- aimux-providers/layout.toml | 36 +- .../src/openai/responses/responses_convert.rs | 2 +- aimux-providers/src/xai/convert.rs | 90 - aimux-providers/src/xai/error.rs | 222 +++ aimux-providers/src/xai/mod.rs | 459 +---- aimux-providers/src/xai/provider.rs | 224 +++ aimux-providers/src/xai/responses/convert.rs | 1013 ----------- .../xai/responses/convert_responses_usage.rs | 47 + .../responses/convert_to_responses_input.rs | 350 ++++ .../responses/map_responses_finish_reason.rs | 14 + aimux-providers/src/xai/responses/mod.rs | 1240 +------------ .../responses/{types.rs => responses_api.rs} | 0 .../xai/responses/responses_language_model.rs | 1536 +++++++++++++++++ .../responses_language_model_options.rs | 86 + .../xai/responses/responses_prepare_tools.rs | 257 +++ .../src/xai/{speech.rs => speech_model.rs} | 4 +- .../src/xai/supports_reasoning_effort.rs | 25 + aimux-providers/src/xai/tests/error.rs | 2 +- .../responses/convert_responses_usage.rs | 4 +- .../responses/convert_to_responses_input.rs | 2 +- .../responses/responses_language_model.rs | 281 ++- .../responses/responses_prepare_tools.rs | 2 +- aimux-providers/src/xai/tests/speech_model.rs | 424 ++++- .../xai/tests/supports_reasoning_effort.rs | 2 +- .../src/xai/tests/transcription_model.rs | 816 ++++++++- aimux-providers/src/xai/tests/video_model.rs | 834 ++++++++- ...ranscription.rs => transcription_model.rs} | 2 +- .../src/xai/{video.rs => video_model.rs} | 4 +- aimux-providers/tests/xai_speech_test.rs | 435 ----- .../tests/xai_transcription_test.rs | 828 --------- .../tests/xai_upstream_e2e_test.rs | 286 --- aimux-providers/tests/xai_video_test.rs | 868 ---------- 32 files changed, 5162 insertions(+), 5233 deletions(-) delete mode 100644 aimux-providers/src/xai/convert.rs create mode 100644 aimux-providers/src/xai/error.rs create mode 100644 aimux-providers/src/xai/provider.rs delete mode 100644 aimux-providers/src/xai/responses/convert.rs create mode 100644 aimux-providers/src/xai/responses/convert_responses_usage.rs create mode 100644 aimux-providers/src/xai/responses/convert_to_responses_input.rs create mode 100644 aimux-providers/src/xai/responses/map_responses_finish_reason.rs rename aimux-providers/src/xai/responses/{types.rs => responses_api.rs} (100%) create mode 100644 aimux-providers/src/xai/responses/responses_language_model.rs create mode 100644 aimux-providers/src/xai/responses/responses_language_model_options.rs create mode 100644 aimux-providers/src/xai/responses/responses_prepare_tools.rs rename aimux-providers/src/xai/{speech.rs => speech_model.rs} (98%) create mode 100644 aimux-providers/src/xai/supports_reasoning_effort.rs rename aimux-providers/src/xai/{transcription.rs => transcription_model.rs} (99%) rename aimux-providers/src/xai/{video.rs => video_model.rs} (99%) delete mode 100644 aimux-providers/tests/xai_speech_test.rs delete mode 100644 aimux-providers/tests/xai_transcription_test.rs delete mode 100644 aimux-providers/tests/xai_upstream_e2e_test.rs delete mode 100644 aimux-providers/tests/xai_video_test.rs diff --git a/aimux-providers/layout.toml b/aimux-providers/layout.toml index 4cc6e0207..728a25209 100644 --- a/aimux-providers/layout.toml +++ b/aimux-providers/layout.toml @@ -22,7 +22,7 @@ issue = 275 [[xai.files]] upstream = "src/get-response-metadata.ts" -rust = "xai/responses/mod.rs" +rust = "xai/responses/responses_language_model.rs" status = "mapped" [xai.files.tests] @@ -61,56 +61,56 @@ reason = "Realtime is out of scope (docs/aisdk-architecture-alignment.md, 本期 [[xai.files]] upstream = "src/responses/convert-to-xai-responses-input.ts" -rust = "xai/responses/convert.rs" +rust = "xai/responses/convert_to_responses_input.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/responses/convert-xai-responses-usage.ts" -rust = "xai/responses/convert.rs" +rust = "xai/responses/convert_responses_usage.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/responses/map-xai-responses-finish-reason.ts" -rust = "xai/responses/convert.rs" +rust = "xai/responses/map_responses_finish_reason.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/responses/xai-responses-api.ts" -rust = "xai/responses/types.rs" +rust = "xai/responses/responses_api.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/responses/xai-responses-language-model-options.ts" -rust = "xai/responses/convert.rs" +rust = "xai/responses/responses_language_model_options.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/responses/xai-responses-language-model.ts" -rust = "xai/responses/mod.rs" +rust = "xai/responses/responses_language_model.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/responses/xai-responses-prepare-tools.ts" -rust = "xai/responses/convert.rs" +rust = "xai/responses/responses_prepare_tools.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/supports-reasoning-effort.ts" -rust = "xai/convert.rs" +rust = "xai/supports_reasoning_effort.rs" status = "mapped" [xai.files.tests] @@ -172,14 +172,14 @@ reason = "Batch is out of scope (docs/aisdk-architecture-alignment.md, 本期范 [[xai.files]] upstream = "src/xai-error.ts" -rust = "xai/mod.rs" +rust = "xai/error.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/xai-file-part-options.ts" -rust = "xai/responses/convert.rs" +rust = "xai/responses/convert_to_responses_input.rs" status = "mapped" [xai.files.tests] @@ -201,7 +201,7 @@ issue = 275 [[xai.files]] upstream = "src/xai-provider.ts" -rust = "xai/mod.rs" +rust = "xai/provider.rs" status = "mapped" [xai.files.tests] @@ -213,42 +213,42 @@ sub-feature = [ [[xai.files]] upstream = "src/xai-speech-model-options.ts" -rust = "xai/speech.rs" +rust = "xai/speech_model.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/xai-speech-model.ts" -rust = "xai/speech.rs" +rust = "xai/speech_model.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/xai-transcription-model-options.ts" -rust = "xai/transcription.rs" +rust = "xai/transcription_model.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/xai-transcription-model.ts" -rust = "xai/transcription.rs" +rust = "xai/transcription_model.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/xai-video-model-options.ts" -rust = "xai/video.rs" +rust = "xai/video_model.rs" status = "mapped" [xai.files.tests] [[xai.files]] upstream = "src/xai-video-model.ts" -rust = "xai/video.rs" +rust = "xai/video_model.rs" status = "mapped" [xai.files.tests] diff --git a/aimux-providers/src/openai/responses/responses_convert.rs b/aimux-providers/src/openai/responses/responses_convert.rs index f9e0de143..262842c27 100644 --- a/aimux-providers/src/openai/responses/responses_convert.rs +++ b/aimux-providers/src/openai/responses/responses_convert.rs @@ -504,7 +504,7 @@ fn annotation_source(annotation: &Value, provider_key: &str) -> Result String { use std::sync::atomic::{AtomicU64, Ordering}; diff --git a/aimux-providers/src/xai/convert.rs b/aimux-providers/src/xai/convert.rs deleted file mode 100644 index 5441a8231..000000000 --- a/aimux-providers/src/xai/convert.rs +++ /dev/null @@ -1,90 +0,0 @@ -//! Helpers shared by the xAI Responses converter. -//! -//! Mirrors `supports-reasoning-effort.ts` and the provider-reference helper -//! used by upstream input conversion. - -use aimux_core::{ - error::AiMuxError, - language_model_message::FilePart, - shared::{FileBytes, FileData, SharedProviderReference}, -}; -use aimux_provider_utils::{ - MediaTypeData, detect_media_type, get_top_level_media_type, is_full_media_type, -}; - -// ── Reasoning effort ───────────────────────────────────────────────────────── - -/// Models that reject the `reasoning_effort` parameter. -/// Matches `^grok-4\.20(-\d{4})?-(non-)?reasoning$`. -fn is_model_without_reasoning_effort(model_id: &str) -> bool { - let rest = match model_id.strip_prefix("grok-4.20") { - Some(r) => r, - None => return false, - }; - let rest = if rest.len() >= 5 - && rest.starts_with('-') - && rest[1..5].chars().all(|c| c.is_ascii_digit()) - { - &rest[5..] - } else { - rest - }; - rest == "-reasoning" || rest == "-non-reasoning" -} - -/// Whether the model accepts the `reasoning_effort` parameter. -#[must_use] -pub fn supports_reasoning_effort(model_id: &str) -> bool { - !is_model_without_reasoning_effort(model_id) -} - -/// Resolve the provider-specific reference string from a file part's -/// `provider` map. -/// -/// # Errors -/// -/// Returns a `String` listing the available providers when the requested key -/// is absent. -pub fn resolve_provider_reference( - reference: &SharedProviderReference, - provider: &str, -) -> Result { - if let Some(value) = reference.get(provider) { - return Ok(value.clone()); - } - let mut available: Vec<&str> = reference.keys().map(String::as_str).collect(); - available.sort_unstable(); - Err(format!( - "No provider reference found for provider '{}'. Available providers: {}", - provider, - available.join(", ") - )) -} - -pub(super) fn resolve_full_media_type(part: &FilePart) -> Result { - let media_type = &part.media_type; - if is_full_media_type(media_type) { - return Ok(media_type.clone()); - } - let data = match &part.data { - FileData::Data { - data: FileBytes::Binary(bytes), - } => Some(MediaTypeData::Bytes(bytes)), - FileData::Data { - data: FileBytes::Base64(data), - } => Some(MediaTypeData::Base64(data)), - _ => None, - }; - let reason = if let Some(data) = data { - if let Some(detected) = detect_media_type(data, Some(get_top_level_media_type(media_type)))? - { - return Ok(detected.into()); - } - "it could not be auto-detected" - } else { - "it is not passed as inline bytes" - }; - Err(AiMuxError::UnsupportedFunctionality(format!( - "file of media type \"{media_type}\" must specify subtype since {reason}" - ))) -} diff --git a/aimux-providers/src/xai/error.rs b/aimux-providers/src/xai/error.rs new file mode 100644 index 000000000..13ee14f3d --- /dev/null +++ b/aimux-providers/src/xai/error.rs @@ -0,0 +1,222 @@ +//! The xAI error and response handlers (`xai-error.ts`): the failed-response +//! handler, stream error events, and the JSON error bodies xAI sends with a +//! successful status. + +use aimux_core::error::AiMuxError; +use serde::de::DeserializeOwned; +use serde_json::Value; + +pub(crate) fn xai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + if let (Some(code), Some(message)) = ( + data.get("code").and_then(Value::as_str), + data.get("error").and_then(Value::as_str), + ) { + return aimux_provider_utils::ProviderErrorParts { + message: format!("{code}: {message}"), + provider_code: Some(code.to_owned()), + }; + } + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .as_str() + .or_else(|| error.get("message").and_then(Value::as_str)) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("code") + .or_else(|| error.get("type")) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }), + } + }) +} + +pub(crate) fn xai_stream_error( + event: &Value, + url: &str, + request_body_values: Value, + response_headers: std::collections::HashMap, +) -> AiMuxError { + let error = event + .get("error") + .or_else(|| { + event + .get("response") + .and_then(|response| response.get("error")) + }) + .unwrap_or(event); + let message = error + .as_str() + .or_else(|| error.get("message").and_then(Value::as_str)) + .or_else(|| event.get("message").and_then(Value::as_str)) + .unwrap_or("xAI stream failed"); + let status_code = event + .get("status") + .or_else(|| event.get("code")) + .or_else(|| error.get("status")) + .or_else(|| error.get("code")) + .and_then(|value| match value { + Value::String(value) + if value.len() == 3 && value.bytes().all(|byte| byte.is_ascii_digit()) => + { + value.parse::().ok() + } + Value::Number(value) => value.as_u64().and_then(|status| u16::try_from(status).ok()), + _ => None, + }) + .filter(|status| (400..=599).contains(status)); + let provider_code = event + .get("code") + .or_else(|| error.get("code")) + .or_else(|| error.get("type")) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }); + let status_code = status_code.or(match provider_code.as_deref() { + Some("rate_limit_exceeded" | "rate_limit_error" | "insufficient_quota") => Some(429), + Some("api_error" | "internal_server_error" | "server_error") => Some(500), + Some("overloaded_error" | "service_unavailable") => Some(503), + Some("timeout" | "timeout_error") => Some(504), + Some("authentication_error" | "invalid_api_key") => Some(401), + Some("permission_error") => Some(403), + Some("not_found_error" | "model_not_found") => Some(404), + Some("bad_request" | "context_length_exceeded" | "invalid_request_error") => Some(400), + _ => None, + }); + let insufficient_quota = provider_code.as_deref() == Some("insufficient_quota"); + let mut result = aimux_provider_utils::stream_error_api_call( + message, + provider_code, + status_code, + event, + url, + request_body_values, + response_headers, + ); + if insufficient_quota && let AiMuxError::ApiCall(error) = &mut result { + error.is_retryable = false; + } + result +} + +pub(crate) fn xai_successful_response_handler() -> aimux_provider_utils::ResponseHandler +where + T: DeserializeOwned + Send + 'static, +{ + aimux_provider_utils::ResponseHandler::new(|input| async move { + let status = input.response.status().as_u16(); + let url = input.url.clone(); + let request_body_values = input.request_body_values.clone(); + let output = aimux_provider_utils::create_json_response_handler::() + .handle(input) + .await?; + let headers = output.response_headers.clone(); + let raw = output.value; + if let Some(error) = raw.get("error").filter(|value| !value.is_null()) { + let message = error + .as_str() + .or_else(|| error.get("message").and_then(Value::as_str)) + .unwrap_or("xAI request failed"); + let provider_code = raw + .get("code") + .or_else(|| error.get("code")) + .or_else(|| error.get("type")) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }); + return Err(AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + status_code: Some(status), + provider_code, + response_body: Some(raw.to_string()), + response_headers: Some(headers), + ..aimux_core::ApiCallError::new(message, url, request_body_values) + }))); + } + // Borrow `raw` instead of `from_value(raw.clone())`: the whole body is + // already a `Value` here and is handed back as `raw_value`, so the + // clone was a second full tree resident at peak. + let value = serde::Deserialize::deserialize(&raw).map_err(|error: serde_json::Error| { + AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + status_code: Some(status), + response_body: Some(raw.to_string()), + response_headers: Some(headers.clone()), + ..aimux_core::ApiCallError::new( + format!("Invalid JSON response: {error}"), + url, + request_body_values, + ) + })) + })?; + Ok(aimux_provider_utils::ResponseHandlerOutput { + value, + raw_value: Some(raw), + response_headers: output.response_headers, + }) + }) +} + +/// xAI occasionally returns a JSON error document with a successful status +/// where an SSE response was requested. Classify that response before handing +/// the body to the standard typed event-source handler. +pub(crate) fn xai_event_source_response_handler() +-> aimux_provider_utils::ResponseHandler>> +where + T: DeserializeOwned + Send + 'static, +{ + aimux_provider_utils::ResponseHandler::new(|input| async move { + let is_json = input + .response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.contains("application/json")); + if !is_json { + return aimux_provider_utils::create_event_source_response_handler::() + .handle(input) + .await; + } + + let status = input.response.status().as_u16(); + let url = input.url.clone(); + let request_body_values = input.request_body_values.clone(); + let output = aimux_provider_utils::create_json_response_handler::() + .handle(input) + .await?; + let headers = output.response_headers; + let raw = output.value; + let error = raw.get("error").filter(|value| !value.is_null()); + let message = error + .and_then(|value| { + value + .as_str() + .or_else(|| value.get("message").and_then(Value::as_str)) + }) + .unwrap_or("Expected an event stream but received JSON"); + let provider_code = raw + .get("code") + .or_else(|| error.and_then(|value| value.get("code"))) + .or_else(|| error.and_then(|value| value.get("type"))) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }); + Err(AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + status_code: Some(status), + provider_code, + response_body: Some(raw.to_string()), + response_headers: Some(headers), + ..aimux_core::ApiCallError::new(message, url, request_body_values) + }))) + }) + .streaming() +} diff --git a/aimux-providers/src/xai/mod.rs b/aimux-providers/src/xai/mod.rs index 65a291313..5b667e13f 100644 --- a/aimux-providers/src/xai/mod.rs +++ b/aimux-providers/src/xai/mod.rs @@ -8,7 +8,7 @@ //! //! The AI SDK's xAI package serves text through the Responses API only //! ([`XaiResponsesModel`], `provider()` = `xai.responses`), and -//! [`language_model`](Provider::language_model) returns it. There is no Chat +//! [`language_model`](aimux_core::provider::Provider::language_model) returns it. There is no Chat //! Completions model in this package; xAI's OpenAI-compatible endpoint is //! reachable through `create_openai_compatible` with `https://api.x.ai/v1`. //! @@ -16,455 +16,22 @@ //! ([`XaiTranscriptionModel`], HTTP and WebSocket streaming) and video //! ([`XaiVideoModel`]) models. Speech and transcription take no model id. -pub mod convert; +mod error; pub(crate) mod options; +mod provider; pub mod responses; -mod speech; -mod transcription; -mod video; +mod speech_model; +mod supports_reasoning_effort; +mod transcription_model; +mod video_model; +pub(crate) use provider::{API_KEY_ENV_VAR, DEFAULT_BASE_URL}; +pub use provider::{XAIProvider, XAIProviderSettings, create_xai, xai}; pub use responses::XaiResponsesModel; -pub use speech::XaiSpeechModel; -pub use transcription::XaiTranscriptionModel; -pub use video::XaiVideoModel; - -use std::sync::{Arc, OnceLock}; - -use futures::future::BoxFuture; -use serde::de::DeserializeOwned; -use serde_json::Value; - -use aimux_core::embedding_model::EmbeddingModel; -use aimux_core::error::AiMuxError; -use aimux_core::image_model::ImageModel; -use aimux_core::language_model::LanguageModel; -use aimux_core::model_catalogue::RuntimeModel; -use aimux_core::provider::{Provider, ProviderDiscovery}; -use aimux_core::speech_model::SpeechModel; -use aimux_core::transcription_model::TranscriptionModel; -use aimux_core::video_model::VideoModel; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable}; - -use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; - -pub(crate) const DEFAULT_BASE_URL: &str = "https://api.x.ai/v1"; -pub(crate) const API_KEY_ENV_VAR: &str = "XAI_API_KEY"; - -pub(crate) fn xai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { - aimux_provider_utils::create_json_error_response_handler(|data| { - if let (Some(code), Some(message)) = ( - data.get("code").and_then(Value::as_str), - data.get("error").and_then(Value::as_str), - ) { - return aimux_provider_utils::ProviderErrorParts { - message: format!("{code}: {message}"), - provider_code: Some(code.to_owned()), - }; - } - let error = data.get("error").unwrap_or(data); - aimux_provider_utils::ProviderErrorParts { - message: error - .as_str() - .or_else(|| error.get("message").and_then(Value::as_str)) - .unwrap_or_default() - .to_owned(), - provider_code: error - .get("code") - .or_else(|| error.get("type")) - .and_then(|value| match value { - Value::String(value) => Some(value.clone()), - Value::Number(value) => Some(value.to_string()), - _ => None, - }), - } - }) -} - -pub(crate) fn xai_stream_error( - event: &Value, - url: &str, - request_body_values: Value, - response_headers: std::collections::HashMap, -) -> AiMuxError { - let error = event - .get("error") - .or_else(|| { - event - .get("response") - .and_then(|response| response.get("error")) - }) - .unwrap_or(event); - let message = error - .as_str() - .or_else(|| error.get("message").and_then(Value::as_str)) - .or_else(|| event.get("message").and_then(Value::as_str)) - .unwrap_or("xAI stream failed"); - let status_code = event - .get("status") - .or_else(|| event.get("code")) - .or_else(|| error.get("status")) - .or_else(|| error.get("code")) - .and_then(|value| match value { - Value::String(value) - if value.len() == 3 && value.bytes().all(|byte| byte.is_ascii_digit()) => - { - value.parse::().ok() - } - Value::Number(value) => value.as_u64().and_then(|status| u16::try_from(status).ok()), - _ => None, - }) - .filter(|status| (400..=599).contains(status)); - let provider_code = event - .get("code") - .or_else(|| error.get("code")) - .or_else(|| error.get("type")) - .and_then(|value| match value { - Value::String(value) => Some(value.clone()), - Value::Number(value) => Some(value.to_string()), - _ => None, - }); - let status_code = status_code.or(match provider_code.as_deref() { - Some("rate_limit_exceeded" | "rate_limit_error" | "insufficient_quota") => Some(429), - Some("api_error" | "internal_server_error" | "server_error") => Some(500), - Some("overloaded_error" | "service_unavailable") => Some(503), - Some("timeout" | "timeout_error") => Some(504), - Some("authentication_error" | "invalid_api_key") => Some(401), - Some("permission_error") => Some(403), - Some("not_found_error" | "model_not_found") => Some(404), - Some("bad_request" | "context_length_exceeded" | "invalid_request_error") => Some(400), - _ => None, - }); - let insufficient_quota = provider_code.as_deref() == Some("insufficient_quota"); - let mut result = aimux_provider_utils::stream_error_api_call( - message, - provider_code, - status_code, - event, - url, - request_body_values, - response_headers, - ); - if insufficient_quota && let AiMuxError::ApiCall(error) = &mut result { - error.is_retryable = false; - } - result -} - -pub(crate) fn xai_successful_response_handler() -> aimux_provider_utils::ResponseHandler -where - T: DeserializeOwned + Send + 'static, -{ - aimux_provider_utils::ResponseHandler::new(|input| async move { - let status = input.response.status().as_u16(); - let url = input.url.clone(); - let request_body_values = input.request_body_values.clone(); - let output = aimux_provider_utils::create_json_response_handler::() - .handle(input) - .await?; - let headers = output.response_headers.clone(); - let raw = output.value; - if let Some(error) = raw.get("error").filter(|value| !value.is_null()) { - let message = error - .as_str() - .or_else(|| error.get("message").and_then(Value::as_str)) - .unwrap_or("xAI request failed"); - let provider_code = raw - .get("code") - .or_else(|| error.get("code")) - .or_else(|| error.get("type")) - .and_then(|value| match value { - Value::String(value) => Some(value.clone()), - Value::Number(value) => Some(value.to_string()), - _ => None, - }); - return Err(AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { - status_code: Some(status), - provider_code, - response_body: Some(raw.to_string()), - response_headers: Some(headers), - ..aimux_core::ApiCallError::new(message, url, request_body_values) - }))); - } - // Borrow `raw` instead of `from_value(raw.clone())`: the whole body is - // already a `Value` here and is handed back as `raw_value`, so the - // clone was a second full tree resident at peak. - let value = serde::Deserialize::deserialize(&raw).map_err(|error: serde_json::Error| { - AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { - status_code: Some(status), - response_body: Some(raw.to_string()), - response_headers: Some(headers.clone()), - ..aimux_core::ApiCallError::new( - format!("Invalid JSON response: {error}"), - url, - request_body_values, - ) - })) - })?; - Ok(aimux_provider_utils::ResponseHandlerOutput { - value, - raw_value: Some(raw), - response_headers: output.response_headers, - }) - }) -} - -/// xAI occasionally returns a JSON error document with a successful status -/// where an SSE response was requested. Classify that response before handing -/// the body to the standard typed event-source handler. -pub(crate) fn xai_event_source_response_handler() --> aimux_provider_utils::ResponseHandler>> -where - T: DeserializeOwned + Send + 'static, -{ - aimux_provider_utils::ResponseHandler::new(|input| async move { - let is_json = input - .response - .headers() - .get(reqwest::header::CONTENT_TYPE) - .and_then(|value| value.to_str().ok()) - .is_some_and(|value| value.contains("application/json")); - if !is_json { - return aimux_provider_utils::create_event_source_response_handler::() - .handle(input) - .await; - } - - let status = input.response.status().as_u16(); - let url = input.url.clone(); - let request_body_values = input.request_body_values.clone(); - let output = aimux_provider_utils::create_json_response_handler::() - .handle(input) - .await?; - let headers = output.response_headers; - let raw = output.value; - let error = raw.get("error").filter(|value| !value.is_null()); - let message = error - .and_then(|value| { - value - .as_str() - .or_else(|| value.get("message").and_then(Value::as_str)) - }) - .unwrap_or("Expected an event stream but received JSON"); - let provider_code = raw - .get("code") - .or_else(|| error.and_then(|value| value.get("code"))) - .or_else(|| error.and_then(|value| value.get("type"))) - .and_then(|value| match value { - Value::String(value) => Some(value.clone()), - Value::Number(value) => Some(value.to_string()), - _ => None, - }); - Err(AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { - status_code: Some(status), - provider_code, - response_body: Some(raw.to_string()), - response_headers: Some(headers), - ..aimux_core::ApiCallError::new(message, url, request_body_values) - }))) - }) - .streaming() -} - -/// Settings of [`create_xai`] (the AI SDK's `XaiProviderSettings`). -/// -/// Every field is optional. Nothing here is evaluated when the provider is -/// created except `base_url`; `api_key` and `headers` are -/// evaluated on every request. -#[derive(Clone, Default)] -pub struct XAIProviderSettings { - /// Base URL for the API calls. Default `https://api.x.ai/v1`; a trailing - /// slash is removed. - pub base_url: Option, - /// The API key. `None` loads `XAI_API_KEY` when a request is made and - /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An - /// explicit value is used as given, `""` included: it never falls back to - /// the environment. - pub api_key: Option, - /// Extra headers on every request. A `None` value removes the header, - /// including `Authorization`. Per-call headers win over these. - pub headers: Option, - /// The transport: a mock, a signing decorator, a proxy-aware client. - /// `None` uses the process default, resolved per request. - pub fetch: Option, - /// Opens the WebSocket of streaming transcription (the AI SDK's - /// `webSocket`). `None` uses the built-in tungstenite connector. - #[cfg(feature = "realtime")] - pub web_socket: Option>, -} - -impl std::fmt::Debug for XAIProviderSettings { - /// Never prints the key or header values. - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let mut debug = f.debug_struct("XAIProviderSettings"); - #[cfg(feature = "realtime")] - debug.field("web_socket", &self.web_socket.is_some()); - debug - .field("base_url", &self.base_url) - .field("api_key", &self.api_key.is_some()) - .field( - "headers", - &self.headers.as_ref().map(std::collections::HashMap::len), - ) - .field("fetch", &self.fetch.is_some()); - debug.finish() - } -} - -/// Create an xAI provider. -/// -/// # Errors -/// -/// Creation does not validate the URL or load the API key; request errors -/// surface when the model is called. -pub fn create_xai(settings: XAIProviderSettings) -> Result { - let base_url = settings - .base_url - .unwrap_or_else(|| DEFAULT_BASE_URL.to_string()); - let base_url = base_url.strip_suffix('/').unwrap_or(&base_url).to_string(); - Ok(XAIProvider { - base_url, - headers: ProviderHeaders::bearer( - Credential::explicit_or_env( - settings.api_key.map(Resolvable::Value), - API_KEY_ENV_VAR, - "xAI API key", - ), - Vec::new(), - settings.headers, - ) - .with_user_agent( - crate::shared::aisdk_pins::XAI.name, - crate::shared::aisdk_pins::XAI.version, - ), - fetch: Some(crate::shared::package_fetch( - settings.fetch, - crate::shared::aisdk_pins::XAI, - )), - #[cfg(feature = "realtime")] - web_socket: settings.web_socket, - }) -} - -/// The default provider: `create_xai` with default settings, created on first -/// use. Creating it reads nothing from the environment and cannot fail; a -/// missing key surfaces from the first request instead. -pub fn xai() -> &'static XAIProvider { - static DEFAULT: OnceLock = OnceLock::new(); - DEFAULT.get_or_init(|| { - create_xai(XAIProviderSettings::default()).expect("default xAI settings are always valid") - }) -} - -/// An xAI provider (the AI SDK's `XaiProvider`). Cheap to clone the models out -/// of; it holds no HTTP client. -pub struct XAIProvider { - base_url: String, - headers: ProviderHeaders, - fetch: Option, - #[cfg(feature = "realtime")] - web_socket: Option>, -} - -impl XAIProvider { - fn model_config(&self, method: &str) -> EndpointConfig { - EndpointConfig::fixed( - format!("xai.{method}"), - self.base_url.clone(), - self.headers.clone(), - self.fetch.clone(), - None, - ) - } - - /// A Responses model (e.g. `"grok-4"`); `provider()` is - /// `"xai.responses"`. - /// - /// Uses the xAI `/responses` endpoint with the Responses API wire format - /// (input items, reasoning objects, provider-executed tools, etc.). - #[must_use] - pub fn responses(&self, model_id: &str) -> XaiResponsesModel { - XaiResponsesModel::from_config(model_id.to_string(), self.model_config("responses")) - } - - /// A video model (e.g. `"grok-imagine-video"`); `provider()` is - /// `"xai.video"`. - #[must_use] - pub fn video(&self, model_id: &str) -> XaiVideoModel { - XaiVideoModel::from_config(model_id.to_string(), self.model_config("video")) - } - - /// The speech (TTS) model; `provider()` is `"xai.speech"`. It has no - /// model id. - #[must_use] - pub fn speech(&self) -> XaiSpeechModel { - XaiSpeechModel::from_config(self.model_config("speech")) - } - - /// The transcription (STT) model; `provider()` is `"xai.transcription"`. - /// It has no model id. - #[must_use] - pub fn transcription(&self) -> XaiTranscriptionModel { - XaiTranscriptionModel::from_config( - self.model_config("transcription"), - #[cfg(feature = "realtime")] - self.web_socket.clone(), - ) - } - - /// The provider as a function: the default language model for an id. The - /// AI SDK's callable provider; the same model as - /// [`responses`](Self::responses) and - /// [`language_model`](Provider::language_model). - #[must_use] - pub fn call(&self, model_id: &str) -> Arc { - Arc::new(self.responses(model_id)) - } -} - -impl Provider for XAIProvider { - fn discovery(&self) -> Option<&dyn ProviderDiscovery> { - Some(self) - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(self.call(model_id)) - } - - fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { - Err(AiMuxError::no_such_model(model_id, "embeddingModel")) - } - - fn image_model(&self, model_id: &str) -> Result, AiMuxError> { - Err(AiMuxError::no_such_model(model_id, "imageModel")) - } - - /// The model id is ignored: xAI has a single transcription model. - fn transcription_model( - &self, - _model_id: &str, - ) -> Option, AiMuxError>> { - Some(Ok(Arc::new(self.transcription()))) - } - - /// The model id is ignored: xAI has a single speech model. - fn speech_model(&self, _model_id: &str) -> Option, AiMuxError>> { - Some(Ok(Arc::new(self.speech()))) - } - - fn video_model(&self, model_id: &str) -> Option, AiMuxError>> { - Some(Ok(Arc::new(self.video(model_id)))) - } -} - -impl ProviderDiscovery for XAIProvider { - /// `GET {base_url}/models`: one exchange, no retry. - fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { - let config = self.model_config("models"); - Box::pin(async move { - crate::shared::list_data_models(&config, xai_failed_response_handler()).await - }) - } -} +pub use speech_model::XaiSpeechModel; +pub use supports_reasoning_effort::supports_reasoning_effort; +pub use transcription_model::XaiTranscriptionModel; +pub use video_model::XaiVideoModel; #[cfg(test)] mod tests; diff --git a/aimux-providers/src/xai/provider.rs b/aimux-providers/src/xai/provider.rs new file mode 100644 index 000000000..c3636fac7 --- /dev/null +++ b/aimux-providers/src/xai/provider.rs @@ -0,0 +1,224 @@ +//! The xAI provider factory and default instance (`xai-provider.ts`). + +use std::sync::{Arc, OnceLock}; + +use aimux_core::embedding_model::EmbeddingModel; +use aimux_core::error::AiMuxError; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::LanguageModel; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_core::speech_model::SpeechModel; +use aimux_core::transcription_model::TranscriptionModel; +use aimux_core::video_model::VideoModel; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable}; +use futures::future::BoxFuture; + +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; + +use super::error::xai_failed_response_handler; +use super::{XaiResponsesModel, XaiSpeechModel, XaiTranscriptionModel, XaiVideoModel}; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.x.ai/v1"; +pub(crate) const API_KEY_ENV_VAR: &str = "XAI_API_KEY"; + +/// Settings of [`create_xai`] (the AI SDK's `XaiProviderSettings`). +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct XAIProviderSettings { + /// Base URL for the API calls. Default `https://api.x.ai/v1`; a trailing + /// slash is removed. + pub base_url: Option, + /// The API key. `None` loads `XAI_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. + pub api_key: Option, + /// Extra headers on every request. A `None` value removes the header, + /// including `Authorization`. Per-call headers win over these. + pub headers: Option, + /// The transport: a mock, a signing decorator, a proxy-aware client. + /// `None` uses the process default, resolved per request. + pub fetch: Option, + /// Opens the WebSocket of streaming transcription (the AI SDK's + /// `webSocket`). `None` uses the built-in tungstenite connector. + #[cfg(feature = "realtime")] + pub web_socket: Option>, +} + +impl std::fmt::Debug for XAIProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut debug = f.debug_struct("XAIProviderSettings"); + #[cfg(feature = "realtime")] + debug.field("web_socket", &self.web_socket.is_some()); + debug + .field("base_url", &self.base_url) + .field("api_key", &self.api_key.is_some()) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("fetch", &self.fetch.is_some()); + debug.finish() + } +} + +/// Create an xAI provider. +/// +/// # Errors +/// +/// Creation does not validate the URL or load the API key; request errors +/// surface when the model is called. +pub fn create_xai(settings: XAIProviderSettings) -> Result { + let base_url = settings + .base_url + .unwrap_or_else(|| DEFAULT_BASE_URL.to_string()); + let base_url = base_url.strip_suffix('/').unwrap_or(&base_url).to_string(); + Ok(XAIProvider { + base_url, + headers: ProviderHeaders::bearer( + Credential::explicit_or_env( + settings.api_key.map(Resolvable::Value), + API_KEY_ENV_VAR, + "xAI API key", + ), + Vec::new(), + settings.headers, + ) + .with_user_agent( + crate::shared::aisdk_pins::XAI.name, + crate::shared::aisdk_pins::XAI.version, + ), + fetch: Some(crate::shared::package_fetch( + settings.fetch, + crate::shared::aisdk_pins::XAI, + )), + #[cfg(feature = "realtime")] + web_socket: settings.web_socket, + }) +} + +/// The default provider: `create_xai` with default settings, created on first +/// use. Creating it reads nothing from the environment and cannot fail; a +/// missing key surfaces from the first request instead. +pub fn xai() -> &'static XAIProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_xai(XAIProviderSettings::default()).expect("default xAI settings are always valid") + }) +} + +/// An xAI provider (the AI SDK's `XaiProvider`). Cheap to clone the models out +/// of; it holds no HTTP client. +pub struct XAIProvider { + base_url: String, + headers: ProviderHeaders, + fetch: Option, + #[cfg(feature = "realtime")] + web_socket: Option>, +} + +impl XAIProvider { + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("xai.{method}"), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) + } + + /// A Responses model (e.g. `"grok-4"`); `provider()` is + /// `"xai.responses"`. + /// + /// Uses the xAI `/responses` endpoint with the Responses API wire format + /// (input items, reasoning objects, provider-executed tools, etc.). + #[must_use] + pub fn responses(&self, model_id: &str) -> XaiResponsesModel { + XaiResponsesModel::from_config(model_id.to_string(), self.model_config("responses")) + } + + /// A video model (e.g. `"grok-imagine-video"`); `provider()` is + /// `"xai.video"`. + #[must_use] + pub fn video(&self, model_id: &str) -> XaiVideoModel { + XaiVideoModel::from_config(model_id.to_string(), self.model_config("video")) + } + + /// The speech (TTS) model; `provider()` is `"xai.speech"`. It has no + /// model id. + #[must_use] + pub fn speech(&self) -> XaiSpeechModel { + XaiSpeechModel::from_config(self.model_config("speech")) + } + + /// The transcription (STT) model; `provider()` is `"xai.transcription"`. + /// It has no model id. + #[must_use] + pub fn transcription(&self) -> XaiTranscriptionModel { + XaiTranscriptionModel::from_config( + self.model_config("transcription"), + #[cfg(feature = "realtime")] + self.web_socket.clone(), + ) + } + + /// The provider as a function: the default language model for an id. The + /// AI SDK's callable provider; the same model as + /// [`responses`](Self::responses) and + /// [`language_model`](Provider::language_model). + #[must_use] + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.responses(model_id)) + } +} + +impl Provider for XAIProvider { + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) + } + + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) + } + + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } + + /// The model id is ignored: xAI has a single transcription model. + fn transcription_model( + &self, + _model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.transcription()))) + } + + /// The model id is ignored: xAI has a single speech model. + fn speech_model(&self, _model_id: &str) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.speech()))) + } + + fn video_model(&self, model_id: &str) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.video(model_id)))) + } +} + +impl ProviderDiscovery for XAIProvider { + /// `GET {base_url}/models`: one exchange, no retry. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config("models"); + Box::pin(async move { + crate::shared::list_data_models(&config, xai_failed_response_handler()).await + }) + } +} diff --git a/aimux-providers/src/xai/responses/convert.rs b/aimux-providers/src/xai/responses/convert.rs deleted file mode 100644 index 0566f815c..000000000 --- a/aimux-providers/src/xai/responses/convert.rs +++ /dev/null @@ -1,1013 +0,0 @@ -//! Conversion functions for the xAI Responses API. -//! -//! Mirrors the TS sources: -//! - `convert-to-xai-responses-input.ts` -//! - `convert-xai-responses-usage.ts` -//! - `map-xai-responses-finish-reason.ts` -//! - `xai-responses-prepare-tools.ts` - -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, LanguageModelV4ToolChoice, ResponseFormat}; -use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; -use aimux_core::tool::Tool; -use aimux_core::types::{FinishReasonUnified, ReasoningEffort, Warning}; - -use aimux_provider_utils::get_top_level_media_type; -use serde_json::{Value, json}; - -use super::types::XaiResponsesUsage; -use crate::xai::convert::{ - resolve_full_media_type, resolve_provider_reference, supports_reasoning_effort, -}; - -// ── Finish reason ──────────────────────────────────────────────────────────── - -#[must_use] -pub fn map_xai_responses_finish_reason(s: &str) -> FinishReasonUnified { - match s { - "stop" | "completed" => FinishReasonUnified::Stop, - "length" | "max_output_tokens" => FinishReasonUnified::Length, - "tool_calls" | "function_call" => FinishReasonUnified::ToolCalls, - "content_filter" => FinishReasonUnified::ContentFilter, - _ => FinishReasonUnified::Other, - } -} - -// ── Usage conversion ───────────────────────────────────────────────────────── - -#[must_use] -pub fn convert_xai_responses_usage(usage: &XaiResponsesUsage) -> aimux_core::types::Usage { - let cache_read = usage - .input_tokens_details - .as_ref() - .and_then(|d| d.cached_tokens) - .unwrap_or(0); - let reasoning = usage - .output_tokens_details - .as_ref() - .and_then(|d| d.reasoning_tokens) - .unwrap_or(0); - - let input_includes_cached = cache_read <= usage.input_tokens; - let input_total = if input_includes_cached { - usage.input_tokens - } else { - usage.input_tokens + cache_read - }; - let input_no_cache = if input_includes_cached { - usage.input_tokens - cache_read - } else { - usage.input_tokens - }; - - aimux_core::types::Usage { - input_tokens: aimux_core::types::InputTokenUsage { - total: Some(input_total as u32), - no_cache: Some(input_no_cache as u32), - cache_read: Some(cache_read as u32), - cache_write: None, - }, - output_tokens: aimux_core::types::OutputTokenUsage { - total: Some(usage.output_tokens as u32), - text: Some((usage.output_tokens - reasoning) as u32), - reasoning: Some(reasoning as u32), - }, - // RFC-0015 P0-3: keep the raw provider usage payload. - raw: serde_json::to_value(usage) - .ok() - .and_then(|value| value.as_object().cloned()), - } -} - -// ── Tool preparation ───────────────────────────────────────────────────────── - -#[derive(Debug, Clone)] -pub struct PreparedResponsesTools { - pub tools: Option>, - pub tool_choice: Option, - pub tool_warnings: Vec, - /// Names of provider-executed tools, keyed by tool id (e.g. "xai.web_search"). - pub provider_tool_names: std::collections::HashMap, -} - -/// # Errors -/// Rejects provider-tool arguments that do not match the upstream schemas. -pub fn prepare_responses_tools( - tools: &Option>, - tool_choice: Option<&LanguageModelV4ToolChoice>, -) -> Result { - let non_empty = tools.as_ref().filter(|t| !t.is_empty()); - for tool in non_empty.into_iter().flatten() { - if let Tool::Provider(tool) = tool { - validate_provider_tool(tool)?; - } - } - let mut tool_warnings: Vec = Vec::new(); - let mut provider_tool_names: std::collections::HashMap = - std::collections::HashMap::new(); - - let tools_opt = match non_empty { - None => None, - Some(tools) => { - let xai_tools: Vec = tools - .iter() - .filter_map(|t| match t { - Tool::Function(ft) => { - let mut func = json!({ - "type": "function", - "name": ft.name, - "parameters": ft.input_schema, - }); - if let Some(ref desc) = ft.description { - func["description"] = json!(desc); - } - if let Some(strict) = ft.strict { - func["strict"] = json!(strict); - } - Some(func) - } - Tool::Provider(pt) => { - provider_tool_names.insert(pt.id.to_string(), pt.name.clone()); - prepare_provider_tool(pt, &mut tool_warnings) - } - }) - .collect(); - Some(xai_tools) - } - }; - - let tool_choice_opt = match (&tools_opt, tool_choice) { - (None, _) => None, - (Some(_), None) => None, - (Some(_), Some(tc)) => { - // Build a map of tool name -> is_provider_tool from the original tools. - let non_empty_tools = non_empty.unwrap(); - match tc { - LanguageModelV4ToolChoice::Auto => Some(json!("auto")), - LanguageModelV4ToolChoice::None => Some(json!("none")), - LanguageModelV4ToolChoice::Required => Some(json!("required")), - LanguageModelV4ToolChoice::Tool { tool_name } => { - let selected = non_empty_tools.iter().rev().find(|t| match t { - Tool::Function(ft) => ft.name == *tool_name, - Tool::Provider(pt) => pt.name == *tool_name, - }); - - match selected { - None => None, - Some(Tool::Function(_)) => { - Some(json!({ "type": "function", "name": tool_name })) - } - Some(Tool::Provider(pt)) => { - // Server-side tools cannot be forced via toolChoice. - tool_warnings.push(Warning::Unsupported { - feature: format!("toolChoice for server-side tool \"{}\"", pt.name), - details: None, - }); - None - } - } - } - } - } - }; - - Ok(PreparedResponsesTools { - tools: tools_opt, - tool_choice: tool_choice_opt, - tool_warnings, - provider_tool_names, - }) -} - -fn prepare_provider_tool( - pt: &aimux_core::tool::ProviderTool, - warnings: &mut Vec, -) -> Option { - let args = &pt.args; - match pt.id.as_str() { - "xai.web_search" => { - let mut tool = json!({ "type": "web_search" }); - if let Some(v) = args.get("allowedDomains") { - tool["allowed_domains"] = v.clone(); - } - if let Some(v) = args.get("excludedDomains") { - tool["excluded_domains"] = v.clone(); - } - if let Some(v) = args.get("enableImageSearch") { - tool["enable_image_search"] = v.clone(); - } - if let Some(v) = args.get("enableImageUnderstanding") { - tool["enable_image_understanding"] = v.clone(); - } - Some(tool) - } - "xai.x_search" => { - let mut tool = json!({ "type": "x_search" }); - if let Some(v) = args.get("allowedXHandles") { - tool["allowed_x_handles"] = v.clone(); - } - if let Some(v) = args.get("excludedXHandles") { - tool["excluded_x_handles"] = v.clone(); - } - if let Some(v) = args.get("fromDate") { - tool["from_date"] = v.clone(); - } - if let Some(v) = args.get("toDate") { - tool["to_date"] = v.clone(); - } - if let Some(v) = args.get("enableImageUnderstanding") { - tool["enable_image_understanding"] = v.clone(); - } - if let Some(v) = args.get("enableVideoUnderstanding") { - tool["enable_video_understanding"] = v.clone(); - } - Some(tool) - } - "xai.code_execution" => Some(json!({ "type": "code_interpreter" })), - "xai.view_image" => Some(json!({ "type": "view_image" })), - "xai.view_x_video" => Some(json!({ "type": "view_x_video" })), - "xai.image_generation" => { - let mut tool = json!({ "type": "image_generation" }); - if let Some(action) = args.get("action") { - tool["action"] = action.clone(); - } - Some(tool) - } - "xai.file_search" => { - let mut tool = json!({ "type": "file_search" }); - if let Some(v) = args.get("vectorStoreIds") { - tool["vector_store_ids"] = v.clone(); - } - if let Some(v) = args.get("maxNumResults") { - tool["max_num_results"] = v.clone(); - } - Some(tool) - } - "xai.mcp" => { - let mut tool = json!({ "type": "mcp" }); - if let Some(v) = args.get("serverUrl") { - tool["server_url"] = v.clone(); - } - if let Some(v) = args.get("serverLabel") { - tool["server_label"] = v.clone(); - } - if let Some(v) = args.get("serverDescription") { - tool["server_description"] = v.clone(); - } - if let Some(v) = args.get("allowedTools") { - tool["allowed_tools"] = v.clone(); - } - if let Some(v) = args.get("headers") { - tool["headers"] = v.clone(); - } - if let Some(v) = args.get("authorization") { - tool["authorization"] = v.clone(); - } - Some(tool) - } - _ => { - warnings.push(Warning::Unsupported { - feature: format!("provider-defined tool {}", pt.name), - details: None, - }); - None - } - } -} - -// ── Provider options helpers ───────────────────────────────────────────────── - -fn xai_option(options: &Option, key: &str) -> Option { - options - .as_ref() - .and_then(|m| crate::xai::options::xai_options(Some(m))) - .and_then(|o| o.get(key)) - .cloned() -} - -fn validate_value( - key: &str, - value: &Value, - valid: impl FnOnce(&Value) -> bool, -) -> Result<(), AiMuxError> { - if valid(value) { - Ok(()) - } else { - Err(AiMuxError::InvalidArgument(format!( - "Invalid xai option {key}: {value}" - ))) - } -} - -fn string_array(value: &Value, maximum: usize) -> bool { - value - .as_array() - .is_some_and(|items| items.len() <= maximum && items.iter().all(Value::is_string)) -} - -fn validate_response_options(options: &Option) -> Result<(), AiMuxError> { - for (key, value) in options - .as_ref() - .and_then(|options| crate::xai::options::xai_options(Some(options))) - .into_iter() - .flatten() - { - validate_value(key, value, |value| match key.as_str() { - "reasoningEffort" => value - .as_str() - .is_some_and(|v| ["none", "low", "medium", "high", "xhigh"].contains(&v)), - "reasoningSummary" => value - .as_str() - .is_some_and(|v| ["auto", "concise", "detailed"].contains(&v)), - "serviceTier" => value - .as_str() - .is_some_and(|v| ["default", "priority"].contains(&v)), - "topLogprobs" => value - .as_f64() - .is_some_and(|v| v.fract() == 0.0 && (0.0..=8.0).contains(&v)), - "minP" => value.as_f64().is_some_and(|v| (0.0..=1.0).contains(&v)), - "maxTurns" => value - .as_f64() - .is_some_and(|v| v.fract() == 0.0 && v.abs() <= 9_007_199_254_740_991.0), - "logprobs" | "store" | "parallelToolCalls" => value.is_boolean(), - "promptCacheKey" | "safetyIdentifier" | "user" | "previousResponseId" => { - value.is_string() - } - "include" => { - value.is_null() - || value.as_array().is_some_and(|items| { - items.iter().all(|item| { - item.as_str().is_some_and(|v| { - [ - "file_search_call.results", - "web_search_call.action.sources", - "code_interpreter_call.outputs", - "reasoning.encrypted_content", - "no_inline_citations", - ] - .contains(&v) - }) - }) - }) - } - _ => true, - })?; - } - Ok(()) -} - -fn validate_provider_tool(tool: &aimux_core::tool::ProviderTool) -> Result<(), AiMuxError> { - let id = tool.id.as_str(); - if ![ - "xai.web_search", - "xai.x_search", - "xai.image_generation", - "xai.file_search", - "xai.mcp", - ] - .contains(&id) - { - return Ok(()); - } - let args = &tool.args; - for (key, value) in args { - validate_value(key, value, |value| match (id, key.as_str()) { - ("xai.web_search", "allowedDomains" | "excludedDomains") => string_array(value, 5), - ("xai.x_search", "allowedXHandles" | "excludedXHandles") => string_array(value, 10), - ("xai.web_search", "enableImageSearch" | "enableImageUnderstanding") - | ("xai.x_search", "enableImageUnderstanding" | "enableVideoUnderstanding") => { - value.is_boolean() - } - ("xai.x_search", "fromDate" | "toDate") => value.is_string(), - ("xai.image_generation", "action") => value - .as_str() - .is_some_and(|v| ["auto", "generate", "edit"].contains(&v)), - ("xai.file_search", "vectorStoreIds") | ("xai.mcp", "allowedTools") => { - string_array(value, usize::MAX) - } - ("xai.file_search", "maxNumResults") => value.is_number(), - ("xai.mcp", "serverUrl" | "serverLabel" | "serverDescription" | "authorization") => { - value.is_string() - } - ("xai.mcp", "headers") => value - .as_object() - .is_some_and(|headers| headers.values().all(Value::is_string)), - _ => true, - })?; - } - for key in match id { - "xai.file_search" => &["vectorStoreIds"][..], - "xai.mcp" => &["serverUrl"][..], - _ => &[], - } { - if !args.contains_key(*key) { - return Err(AiMuxError::InvalidArgument(format!( - "Missing {key} for {}", - tool.name - ))); - } - } - Ok(()) -} - -// ── Input conversion ───────────────────────────────────────────────────────── - -/// Convert a [`LanguageModelPrompt`] into the xAI Responses `input` array, -/// returning warnings. -/// -/// # Errors -/// -/// Returns `AiMuxError::InvalidArgument` when a part cannot be converted. -pub fn convert_to_xai_responses_input( - prompt: &LanguageModelPrompt, -) -> Result<(Vec, Vec), AiMuxError> { - let mut input = Vec::new(); - let mut warnings: Vec = Vec::new(); - - for msg in prompt { - match msg { - LanguageModelMessage::System { content, .. } => { - input.push(json!({ "role": "system", "content": content })); - } - LanguageModelMessage::User { content, .. } => { - let mut content_parts = Vec::new(); - for part in content { - match part { - UserPart::Text(TextPart { text, .. }) => { - content_parts.push(json!({ "type": "input_text", "text": text })); - } - 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 - ))); - } - 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, Some(&b64), None)? - } - FileData::Data { - data: FileBytes::Base64(data), - } => convert_image_part(file, Some(data), None)?, - FileData::Url { url, .. } - if get_top_level_media_type(&file.media_type) == "image" => - { - convert_image_part(file, None, Some(url))? - } - FileData::Url { url, .. } => { - json!({ "type": "input_file", "file_url": url }) - } - FileData::Reference { reference } => { - let file_id = resolve_provider_reference( - reference, - crate::xai::options::NAMESPACE, - ) - .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 })); - } - LanguageModelMessage::Assistant { content, .. } => { - for part in content { - match part { - AssistantPart::Text(TextPart { - text, - provider_options, - }) => { - let id = provider_options - .as_ref() - .and_then(|po| crate::xai::options::xai_options(Some(po))) - .and_then(|x| x.get("itemId")) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let mut msg_obj = json!({ "role": "assistant", "content": text }); - if let Some(id) = id { - msg_obj["id"] = json!(id); - } - input.push(msg_obj); - } - AssistantPart::ToolCall(ToolCallPart { - tool_call_id, - tool_name, - input: tool_input, - provider_executed, - provider_options, - .. - }) => { - // Skip provider-executed tool calls. - let is_provider_executed = provider_executed.unwrap_or(false); - if is_provider_executed { - continue; - } - let id = provider_options - .as_ref() - .and_then(|po| crate::xai::options::xai_options(Some(po))) - .and_then(|x| x.get("itemId")) - .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 = tool_input.to_string(); - input.push(json!({ - "type": "function_call", - "id": item_id, - "call_id": tool_call_id, - "name": tool_name, - "arguments": arguments, - "status": "completed" - })); - } - AssistantPart::ToolResult(_) => {} - AssistantPart::Reasoning(ReasoningPart { - text, - provider_options, - .. - }) => { - let item_id = provider_options - .as_ref() - .and_then(|po| crate::xai::options::xai_options(Some(po))) - .and_then(|x| x.get("itemId")) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let encrypted_content = provider_options - .as_ref() - .and_then(|po| crate::xai::options::xai_options(Some(po))) - .and_then(|x| x.get("reasoningEncryptedContent")) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - - if item_id.is_some() || encrypted_content.is_some() { - let mut summary: Vec = Vec::new(); - if !text.is_empty() { - summary.push(json!({ "type": "summary_text", "text": text })); - } - let mut reasoning_obj = json!({ - "type": "reasoning", - "id": item_id.clone().unwrap_or_default(), - "summary": summary, - "status": "completed" - }); - if let Some(ec) = encrypted_content { - reasoning_obj["encrypted_content"] = json!(ec); - } - input.push(reasoning_obj); - } else { - warnings.push(Warning::Other { - message: "Reasoning parts without itemId or encrypted content cannot be sent back to xAI. Skipping.".to_string(), - }); - } - } - AssistantPart::File(_) - | AssistantPart::ReasoningFile(_) - | AssistantPart::Custom(_) => { - let kind = match part { - AssistantPart::File(_) => "file", - AssistantPart::ReasoningFile(_) => "reasoning-file", - _ => "custom", - }; - warnings.push(Warning::Other { - message: format!("xAI Responses API does not support {kind} in assistant messages"), - }); - } - } - } - } - LanguageModelMessage::Tool { content, .. } => { - for part in content { - let ToolPart::ToolResult(ToolResultPart { - tool_call_id, - output, - .. - }) = part - else { - continue; - }; - let output_value = convert_tool_output(output)?; - input.push(json!({ - "type": "function_call_output", - "call_id": tool_call_id, - "output": output_value - })); - } - } - } - } - - Ok((input, warnings)) -} - -fn convert_tool_output(output: &ToolResultOutput) -> Result { - match output { - ToolResultOutput::ExecutionDenied { reason, .. } => { - Ok(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 get_top_level_media_type(&part.media_type) == "image" => - { - let image_url = match &part.data { - FileData::Url { url, .. } => url.clone(), - FileData::Data { data } => { - use base64::Engine; - let encoded = match data { - FileBytes::Base64(data) => data.clone(), - FileBytes::Binary(bytes) => { - base64::engine::general_purpose::STANDARD.encode(bytes) - } - }; - format!("data:{};base64,{encoded}", resolve_full_media_type(part)?) - } - _ => continue, - }; - parts.push(json!({"type":"input_image", "image_url":image_url})); - } - _ => {} - } - } - Ok(json!(parts)) - } - _ => Ok(crate::openai::convert::tool_result_to_content(output)), - } -} - -fn convert_image_part( - file_part: &FilePart, - b64_data: Option<&str>, - url: Option<&str>, -) -> Result { - let media_type = &file_part.media_type; - if get_top_level_media_type(media_type) != "image" { - return Err(AiMuxError::UnsupportedFunctionality(format!( - "file part media type {media_type} as inline data (xAI Responses requires a URL or a Files API reference for non-image files)" - ))); - } - let image_url = if let Some(url_str) = url { - url_str.to_string() - } else if let Some(b64) = b64_data { - let full_mt = resolve_full_media_type(file_part)?; - format!("data:{full_mt};base64,{b64}") - } else { - String::new() - }; - - let mut part = json!({ - "type": "input_image", - "image_url": image_url, - }); - - // Image detail provider option. - if let Some(detail) = file_part - .provider_options - .as_ref() - .and_then(|po| crate::xai::options::xai_options(Some(po))) - .and_then(|x| x.get("imageDetail")) - { - validate_value("imageDetail", detail, |value| { - value - .as_str() - .is_some_and(|value| ["auto", "low", "high"].contains(&value)) - })?; - part["detail"] = detail.clone(); - } - - Ok(part) -} - -// ── Request body builder ───────────────────────────────────────────────────── - -#[derive(Debug, Clone)] -pub struct ResponsesRequestBodyResult { - pub body: Value, - pub warnings: Vec, - pub provider_tool_names: std::collections::HashMap, -} - -/// Build the xAI Responses request body, returning warnings and provider tool -/// names. -/// -/// # Errors -/// -/// Propagates conversion errors from `convert_to_xai_responses_input`. -pub fn build_responses_request_body( - model_id: &str, - options: &CallOptions, - stream: bool, -) -> Result { - let mut warnings: Vec = Vec::new(); - let xai_opts = &options.provider_options; - validate_response_options(xai_opts)?; - - if options.frequency_penalty.is_some() { - warnings.push(Warning::Unsupported { - feature: "frequencyPenalty".to_string(), - details: None, - }); - } - - if options.presence_penalty.is_some() { - warnings.push(Warning::Unsupported { - feature: "presencePenalty".to_string(), - details: None, - }); - } - - if options.stop_sequences.is_some() { - warnings.push(Warning::Unsupported { - feature: "stopSequences".to_string(), - details: None, - }); - } - - let (input, input_warnings) = convert_to_xai_responses_input(&options.prompt)?; - warnings.extend(input_warnings); - - let prepared = prepare_responses_tools(&options.tools, options.tool_choice.as_ref())?; - for tw in &prepared.tool_warnings { - warnings.push(tw.clone()); - } - - // Reasoning effort. - let mut reasoning_effort: Option = xai_option(xai_opts, "reasoningEffort").map(|v| { - v.as_str() - .map(std::string::ToString::to_string) - .unwrap_or(v.to_string()) - }); - - if reasoning_effort.is_none() && options.reasoning.is_some_and(ReasoningEffort::is_custom) { - let reasoning = options.reasoning.unwrap(); - if !supports_reasoning_effort(model_id) { - warnings.push(Warning::Unsupported { - feature: "reasoning".to_string(), - details: Some(format!( - "reasoning \"{reasoning}\" is not supported by this model." - )), - }); - } else if reasoning == ReasoningEffort::None { - reasoning_effort = Some("none".to_string()); - } else { - reasoning_effort = Some(match reasoning { - ReasoningEffort::Minimal => "low".to_string(), - ReasoningEffort::Low => "low".to_string(), - ReasoningEffort::Medium => "medium".to_string(), - ReasoningEffort::High => "high".to_string(), - ReasoningEffort::Xhigh if model_id == "grok-4.6" => "xhigh".to_string(), - ReasoningEffort::Xhigh => "high".to_string(), - _ => reasoning.to_string(), - }); - if reasoning_effort.as_deref() != Some(reasoning.to_string().as_str()) { - warnings.push(Warning::Compatibility { - feature: "reasoning".to_string(), - details: Some(format!( - "reasoning \"{reasoning}\" is not directly supported by this model. mapped to effort \"{}\".", - reasoning_effort.as_deref().unwrap_or_default() - )), - }); - } - } - } - - // Store option. - let store = xai_option(xai_opts, "store").and_then(|v| v.as_bool()); - - // Include array. - let mut include: Option> = xai_option(xai_opts, "include").and_then(|v| { - v.as_array().map(|arr| { - arr.iter() - .filter_map(|v| v.as_str().map(std::string::ToString::to_string)) - .collect() - }) - }); - - if store == Some(false) { - include = match include { - None => Some(vec!["reasoning.encrypted_content".to_string()]), - Some(mut v) => { - v.push("reasoning.encrypted_content".to_string()); - Some(v) - } - }; - } - - let previous_response_id = xai_option(xai_opts, "previousResponseId") - .and_then(|v| v.as_str().map(std::string::ToString::to_string)); - - // Logprobs. - let logprobs = xai_option(xai_opts, "logprobs") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - let top_logprobs = xai_option(xai_opts, "topLogprobs"); - - // Build body. - let mut body = json!({ "model": model_id, "input": input }); - - if stream { - body["stream"] = json!(true); - } - if let Some(max_tokens) = options.max_output_tokens { - body["max_output_tokens"] = json!(max_tokens); - } - if let Some(temp) = options.temperature { - body["temperature"] = json!(temp); - } - if let Some(tp) = options.top_p { - body["top_p"] = json!(tp); - } - if let Some(tk) = options.top_k { - body["top_k"] = json!(tk); - } - if let Some(seed) = options.seed { - body["seed"] = json!(seed); - } - if logprobs || top_logprobs.is_some() { - body["logprobs"] = json!(true); - } - if let Some(tlp) = top_logprobs { - body["top_logprobs"] = tlp; - } - - if let Some(effort) = &reasoning_effort { - body["reasoning"] = json!({ "effort": effort }); - } - for (option, parameter) in [ - ("minP", "min_p"), - ("maxTurns", "max_turns"), - ("parallelToolCalls", "parallel_tool_calls"), - ("promptCacheKey", "prompt_cache_key"), - ("safetyIdentifier", "safety_identifier"), - ("serviceTier", "service_tier"), - ("user", "user"), - ] { - if let Some(value) = xai_option(xai_opts, option) { - body[parameter] = value; - } - } - - // Store. - if store == Some(false) { - body["store"] = json!(false); - } - - // Include. - if let Some(inc) = &include { - body["include"] = json!(inc); - } - - // Previous response ID. - if let Some(prev_id) = &previous_response_id { - body["previous_response_id"] = json!(prev_id); - } - - // Response format. - if let Some(ref rf) = options.response_format - && let ResponseFormat::Json { - schema, - name, - description, - } = rf - { - let format = if schema.is_some() { - let mut fmt = json!({ "type": "json_schema", "strict": true }); - fmt["name"] = json!(name.clone().unwrap_or_else(|| "response".to_string())); - if let Some(desc) = description { - fmt["description"] = json!(desc); - } - if let Some(s) = schema { - fmt["schema"] = s.clone(); - } - fmt - } else { - json!({ "type": "json_object" }) - }; - body["text"] = json!({ "format": format }); - } - - // Tools. - if let Some(tools) = &prepared.tools { - body["tools"] = json!(tools); - } - if let Some(tc) = prepared.tool_choice { - body["tool_choice"] = tc; - } - - Ok(ResponsesRequestBodyResult { - body, - warnings, - provider_tool_names: prepared.provider_tool_names, - }) -} - -// ── Tool name resolution helpers ───────────────────────────────────────────── - -/// Sub-tool names that map to web_search. -pub const WEB_SEARCH_SUB_TOOLS: &[&str] = - &["web_search", "web_search_with_snippets", "browse_page"]; -/// Sub-tool names that map to x_search. -pub const X_SEARCH_SUB_TOOLS: &[&str] = &[ - "x_user_search", - "x_keyword_search", - "x_semantic_search", - "x_thread_fetch", -]; - -/// Resolve the tool name for a server-side tool call, using the user-provided -/// provider tool name if available, falling back to a default. -#[must_use] -pub fn resolve_tool_name( - part_type: &str, - part_name: Option<&str>, - provider_tool_names: &std::collections::HashMap, -) -> String { - let name = part_name.unwrap_or(""); - let default_web = provider_tool_names - .get("xai.web_search") - .cloned() - .unwrap_or_else(|| "web_search".to_string()); - let default_x = provider_tool_names - .get("xai.x_search") - .cloned() - .unwrap_or_else(|| "x_search".to_string()); - let default_code = provider_tool_names - .get("xai.code_execution") - .cloned() - .unwrap_or_else(|| "code_execution".to_string()); - let default_mcp = provider_tool_names - .get("xai.mcp") - .cloned() - .unwrap_or_else(|| part_name.unwrap_or("mcp").to_string()); - let default_file_search = provider_tool_names - .get("xai.file_search") - .cloned() - .unwrap_or_else(|| "file_search".to_string()); - - if WEB_SEARCH_SUB_TOOLS.contains(&name) || part_type == "web_search_call" { - default_web - } else if X_SEARCH_SUB_TOOLS.contains(&name) || part_type == "x_search_call" { - default_x - } else if name == "code_execution" - || part_type == "code_interpreter_call" - || part_type == "code_execution_call" - { - default_code - } else if part_type == "mcp_call" { - default_mcp - } else if part_type == "file_search_call" { - default_file_search - } else if part_type == "image_generation_call" { - provider_tool_names - .get("xai.image_generation") - .cloned() - .unwrap_or_else(|| "image_generation".to_string()) - } else if !name.is_empty() { - name.to_string() - } else { - String::new() - } -} - -/// Get the tool input for a server-side tool call. -#[must_use] -pub fn get_tool_input(part_type: &str, part: &Value) -> String { - match part_type { - "custom_tool_call" => part - .get("input") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(), - "mcp_call" => part - .get("arguments") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(), - _ => part - .get("arguments") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(), - } -} diff --git a/aimux-providers/src/xai/responses/convert_responses_usage.rs b/aimux-providers/src/xai/responses/convert_responses_usage.rs new file mode 100644 index 000000000..984c80584 --- /dev/null +++ b/aimux-providers/src/xai/responses/convert_responses_usage.rs @@ -0,0 +1,47 @@ +//! `convert-xai-responses-usage.ts`. + +use super::responses_api::XaiResponsesUsage; + +#[must_use] +pub fn convert_xai_responses_usage(usage: &XaiResponsesUsage) -> aimux_core::types::Usage { + let cache_read = usage + .input_tokens_details + .as_ref() + .and_then(|d| d.cached_tokens) + .unwrap_or(0); + let reasoning = usage + .output_tokens_details + .as_ref() + .and_then(|d| d.reasoning_tokens) + .unwrap_or(0); + + let input_includes_cached = cache_read <= usage.input_tokens; + let input_total = if input_includes_cached { + usage.input_tokens + } else { + usage.input_tokens + cache_read + }; + let input_no_cache = if input_includes_cached { + usage.input_tokens - cache_read + } else { + usage.input_tokens + }; + + aimux_core::types::Usage { + input_tokens: aimux_core::types::InputTokenUsage { + total: Some(input_total as u32), + no_cache: Some(input_no_cache as u32), + cache_read: Some(cache_read as u32), + cache_write: None, + }, + output_tokens: aimux_core::types::OutputTokenUsage { + total: Some(usage.output_tokens as u32), + text: Some((usage.output_tokens - reasoning) as u32), + reasoning: Some(reasoning as u32), + }, + // RFC-0015 P0-3: keep the raw provider usage payload. + raw: serde_json::to_value(usage) + .ok() + .and_then(|value| value.as_object().cloned()), + } +} diff --git a/aimux-providers/src/xai/responses/convert_to_responses_input.rs b/aimux-providers/src/xai/responses/convert_to_responses_input.rs new file mode 100644 index 000000000..bb2749ad8 --- /dev/null +++ b/aimux-providers/src/xai/responses/convert_to_responses_input.rs @@ -0,0 +1,350 @@ +//! Prompt to the Responses `input` array (`convert-to-xai-responses-input.ts`), +//! with the image `detail` option of `xai-file-part-options.ts`. + +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::shared::{FileBytes, FileData, SharedProviderReference}; +use aimux_core::types::Warning; +use aimux_provider_utils::{ + MediaTypeData, detect_media_type, get_top_level_media_type, is_full_media_type, +}; +use serde_json::{Value, json}; + +use super::responses_language_model_options::validate_value; + +/// Convert a [`LanguageModelPrompt`] into the xAI Responses `input` array, +/// returning warnings. +/// +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when a part cannot be converted. +pub fn convert_to_xai_responses_input( + prompt: &LanguageModelPrompt, +) -> Result<(Vec, Vec), AiMuxError> { + let mut input = Vec::new(); + let mut warnings: Vec = Vec::new(); + + for msg in prompt { + match msg { + LanguageModelMessage::System { content, .. } => { + input.push(json!({ "role": "system", "content": content })); + } + LanguageModelMessage::User { content, .. } => { + let mut content_parts = Vec::new(); + for part in content { + match part { + UserPart::Text(TextPart { text, .. }) => { + content_parts.push(json!({ "type": "input_text", "text": text })); + } + 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 + ))); + } + 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, Some(&b64), None)? + } + FileData::Data { + data: FileBytes::Base64(data), + } => convert_image_part(file, Some(data), None)?, + FileData::Url { url, .. } + if get_top_level_media_type(&file.media_type) == "image" => + { + convert_image_part(file, None, Some(url))? + } + FileData::Url { url, .. } => { + json!({ "type": "input_file", "file_url": url }) + } + FileData::Reference { reference } => { + let file_id = resolve_provider_reference( + reference, + crate::xai::options::NAMESPACE, + ) + .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 })); + } + LanguageModelMessage::Assistant { content, .. } => { + for part in content { + match part { + AssistantPart::Text(TextPart { + text, + provider_options, + }) => { + let id = provider_options + .as_ref() + .and_then(|po| crate::xai::options::xai_options(Some(po))) + .and_then(|x| x.get("itemId")) + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let mut msg_obj = json!({ "role": "assistant", "content": text }); + if let Some(id) = id { + msg_obj["id"] = json!(id); + } + input.push(msg_obj); + } + AssistantPart::ToolCall(ToolCallPart { + tool_call_id, + tool_name, + input: tool_input, + provider_executed, + provider_options, + .. + }) => { + // Skip provider-executed tool calls. + let is_provider_executed = provider_executed.unwrap_or(false); + if is_provider_executed { + continue; + } + let id = provider_options + .as_ref() + .and_then(|po| crate::xai::options::xai_options(Some(po))) + .and_then(|x| x.get("itemId")) + .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 = tool_input.to_string(); + input.push(json!({ + "type": "function_call", + "id": item_id, + "call_id": tool_call_id, + "name": tool_name, + "arguments": arguments, + "status": "completed" + })); + } + AssistantPart::ToolResult(_) => {} + AssistantPart::Reasoning(ReasoningPart { + text, + provider_options, + .. + }) => { + let item_id = provider_options + .as_ref() + .and_then(|po| crate::xai::options::xai_options(Some(po))) + .and_then(|x| x.get("itemId")) + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let encrypted_content = provider_options + .as_ref() + .and_then(|po| crate::xai::options::xai_options(Some(po))) + .and_then(|x| x.get("reasoningEncryptedContent")) + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + + if item_id.is_some() || encrypted_content.is_some() { + let mut summary: Vec = Vec::new(); + if !text.is_empty() { + summary.push(json!({ "type": "summary_text", "text": text })); + } + let mut reasoning_obj = json!({ + "type": "reasoning", + "id": item_id.clone().unwrap_or_default(), + "summary": summary, + "status": "completed" + }); + if let Some(ec) = encrypted_content { + reasoning_obj["encrypted_content"] = json!(ec); + } + input.push(reasoning_obj); + } else { + warnings.push(Warning::Other { + message: "Reasoning parts without itemId or encrypted content cannot be sent back to xAI. Skipping.".to_string(), + }); + } + } + AssistantPart::File(_) + | AssistantPart::ReasoningFile(_) + | AssistantPart::Custom(_) => { + let kind = match part { + AssistantPart::File(_) => "file", + AssistantPart::ReasoningFile(_) => "reasoning-file", + _ => "custom", + }; + warnings.push(Warning::Other { + message: format!("xAI Responses API does not support {kind} in assistant messages"), + }); + } + } + } + } + LanguageModelMessage::Tool { content, .. } => { + for part in content { + let ToolPart::ToolResult(ToolResultPart { + tool_call_id, + output, + .. + }) = part + else { + continue; + }; + let output_value = convert_tool_output(output)?; + input.push(json!({ + "type": "function_call_output", + "call_id": tool_call_id, + "output": output_value + })); + } + } + } + } + + Ok((input, warnings)) +} + +fn convert_tool_output(output: &ToolResultOutput) -> Result { + match output { + ToolResultOutput::ExecutionDenied { reason, .. } => { + Ok(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 get_top_level_media_type(&part.media_type) == "image" => + { + let image_url = match &part.data { + FileData::Url { url, .. } => url.clone(), + FileData::Data { data } => { + use base64::Engine; + let encoded = match data { + FileBytes::Base64(data) => data.clone(), + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) + } + }; + format!("data:{};base64,{encoded}", resolve_full_media_type(part)?) + } + _ => continue, + }; + parts.push(json!({"type":"input_image", "image_url":image_url})); + } + _ => {} + } + } + Ok(json!(parts)) + } + _ => Ok(crate::openai::convert::tool_result_to_content(output)), + } +} + +fn convert_image_part( + file_part: &FilePart, + b64_data: Option<&str>, + url: Option<&str>, +) -> Result { + let media_type = &file_part.media_type; + if get_top_level_media_type(media_type) != "image" { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "file part media type {media_type} as inline data (xAI Responses requires a URL or a Files API reference for non-image files)" + ))); + } + let image_url = if let Some(url_str) = url { + url_str.to_string() + } else if let Some(b64) = b64_data { + let full_mt = resolve_full_media_type(file_part)?; + format!("data:{full_mt};base64,{b64}") + } else { + String::new() + }; + + let mut part = json!({ + "type": "input_image", + "image_url": image_url, + }); + + // Image detail provider option. + if let Some(detail) = file_part + .provider_options + .as_ref() + .and_then(|po| crate::xai::options::xai_options(Some(po))) + .and_then(|x| x.get("imageDetail")) + { + validate_value("imageDetail", detail, |value| { + value + .as_str() + .is_some_and(|value| ["auto", "low", "high"].contains(&value)) + })?; + part["detail"] = detail.clone(); + } + + Ok(part) +} + +/// Resolve the provider-specific reference string from a file part's +/// `provider` map. +/// +/// # Errors +/// +/// Returns a `String` listing the available providers when the requested key +/// is absent. +pub fn resolve_provider_reference( + reference: &SharedProviderReference, + provider: &str, +) -> Result { + if let Some(value) = reference.get(provider) { + return Ok(value.clone()); + } + let mut available: Vec<&str> = reference.keys().map(String::as_str).collect(); + available.sort_unstable(); + Err(format!( + "No provider reference found for provider '{}'. Available providers: {}", + provider, + available.join(", ") + )) +} + +fn resolve_full_media_type(part: &FilePart) -> Result { + let media_type = &part.media_type; + if is_full_media_type(media_type) { + return Ok(media_type.clone()); + } + let data = match &part.data { + FileData::Data { + data: FileBytes::Binary(bytes), + } => Some(MediaTypeData::Bytes(bytes)), + FileData::Data { + data: FileBytes::Base64(data), + } => Some(MediaTypeData::Base64(data)), + _ => None, + }; + let reason = if let Some(data) = data { + if let Some(detected) = detect_media_type(data, Some(get_top_level_media_type(media_type)))? + { + return Ok(detected.into()); + } + "it could not be auto-detected" + } else { + "it is not passed as inline bytes" + }; + Err(AiMuxError::UnsupportedFunctionality(format!( + "file of media type \"{media_type}\" must specify subtype since {reason}" + ))) +} diff --git a/aimux-providers/src/xai/responses/map_responses_finish_reason.rs b/aimux-providers/src/xai/responses/map_responses_finish_reason.rs new file mode 100644 index 000000000..1d1cc5135 --- /dev/null +++ b/aimux-providers/src/xai/responses/map_responses_finish_reason.rs @@ -0,0 +1,14 @@ +//! `map-xai-responses-finish-reason.ts`. + +use aimux_core::types::FinishReasonUnified; + +#[must_use] +pub fn map_xai_responses_finish_reason(s: &str) -> FinishReasonUnified { + match s { + "stop" | "completed" => FinishReasonUnified::Stop, + "length" | "max_output_tokens" => FinishReasonUnified::Length, + "tool_calls" | "function_call" => FinishReasonUnified::ToolCalls, + "content_filter" => FinishReasonUnified::ContentFilter, + _ => FinishReasonUnified::Other, + } +} diff --git a/aimux-providers/src/xai/responses/mod.rs b/aimux-providers/src/xai/responses/mod.rs index 5c4f3700b..60d576b81 100644 --- a/aimux-providers/src/xai/responses/mod.rs +++ b/aimux-providers/src/xai/responses/mod.rs @@ -1,1220 +1,22 @@ -//! xAI Responses API language model. -//! -//! Implements `LanguageModel` against xAI's `/responses` endpoint. -//! Mirrors the TS `XaiResponsesLanguageModel`. -//! -//! Key differences from the chat model: -//! - Uses `input` (not `messages`) with Responses API format -//! - `reasoning: { effort, summary }` object (not `reasoning_effort` string) -//! - `text.format` for structured output (not `response_format`) -//! - Provider-executed tools (web_search, x_search, code_execution, etc.) -//! - ~65 streaming event types (response.created, response.output_text.delta, etc.) -//! - `cost_in_usd_ticks` in providerMetadata -//! - Citations via annotations → sources - -pub mod convert; -pub mod types; - -use aimux_core::tool::RawToolCall; -use aimux_core::tool::ToolResult; -use std::collections::HashMap; -use std::sync::atomic::{AtomicU64, Ordering}; - -use aimux_provider_utils::{Transformer, pipe_through}; - -use crate::stream::{Controller, Emit, FINISH_OTHER, StreamContext}; -use async_trait::async_trait; -use futures::StreamExt; -use serde_json::{Value, json}; - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::{LanguageModel, SupportedUrls}; -use aimux_core::options::CallOptions; -use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; -use aimux_core::stream_part::StreamPart; -use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; - -use crate::shared::EndpointConfig; -use convert::{ - build_responses_request_body, convert_xai_responses_usage, get_tool_input, - map_xai_responses_finish_reason, resolve_tool_name, +//! The xAI Responses API (`src/responses/` of `@ai-sdk/xai`). + +mod convert_responses_usage; +mod convert_to_responses_input; +mod map_responses_finish_reason; +mod responses_api; +mod responses_language_model; +mod responses_language_model_options; +mod responses_prepare_tools; + +pub use convert_responses_usage::convert_xai_responses_usage; +pub use convert_to_responses_input::{convert_to_xai_responses_input, resolve_provider_reference}; +pub use map_responses_finish_reason::map_xai_responses_finish_reason; +pub use responses_api::{ + XaiResponsesInputTokensDetails, XaiResponsesOutputTokensDetails, XaiResponsesResponse, + XaiResponsesUsage, event_type, }; - -/// Global counter for generating unique source IDs. -static SOURCE_ID_COUNTER: AtomicU64 = AtomicU64::new(0); - -fn generate_source_id() -> String { - format!("id-{}", SOURCE_ID_COUNTER.fetch_add(1, Ordering::SeqCst)) -} - -fn zero_usage() -> Usage { - Usage { - input_tokens: aimux_core::types::InputTokenUsage { - total: Some(0), - no_cache: Some(0), - cache_read: Some(0), - cache_write: Some(0), - }, - output_tokens: aimux_core::types::OutputTokenUsage { - total: Some(0), - text: Some(0), - reasoning: Some(0), - }, - raw: None, - } -} - -fn response_provider_metadata( - cost: Option, - service_tier: Option<&str>, - prompt_cache_key: Option<&str>, - safety_identifier: Option<&str>, -) -> Option { - let mut metadata = serde_json::Map::new(); - if let Some(cost) = cost { - metadata.insert("costInUsdTicks".into(), json!(cost)); - } - for (key, value) in [ - ("serviceTier", service_tier), - ("promptCacheKey", prompt_cache_key), - ("safetyIdentifier", safety_identifier), - ] { - if let Some(value) = value { - metadata.insert(key.into(), json!(value)); - } - } - (!metadata.is_empty()).then(|| crate::xai::options::xai_metadata(Value::Object(metadata))) -} - -fn map_web_search_action(action: &Value) -> Value { - let Some(kind) = action.get("type").and_then(Value::as_str) else { - return json!({}); - }; - let fields: &[&str] = match kind { - "search" => &["query"], - "open_page" => &["url"], - "find_in_page" => &["url", "pattern"], - _ => return json!({}), - }; - if fields.iter().any(|key| { - action - .get(key) - .is_some_and(|v| !v.is_null() && !v.is_string()) - }) || action - .get("sources") - .is_some_and(|v| !v.is_null() && !v.is_array()) - || (kind == "search" - && action.get("queries").is_some_and(|v| { - !v.is_null() && !v.as_array().is_some_and(|a| a.iter().all(Value::is_string)) - })) - { - return json!({}); - } - let mapped_type = match kind { - "open_page" => "openPage", - "find_in_page" => "findInPage", - _ => kind, - }; - let mut mapped_action = json!({ "type": mapped_type }); - for key in fields { - if let Some(value) = action.get(key) - && (kind != "search" || !value.is_null()) - { - mapped_action[key] = value.clone(); - } - } - if kind == "search" - && let Some(queries) = action.get("queries").filter(|v| !v.is_null()) - { - mapped_action["queries"] = queries.clone(); - } - let mut result = json!({ "action": mapped_action }); - if let Some(sources) = action.get("sources").and_then(Value::as_array) { - let sources: Vec = sources - .iter() - .filter(|source| source.get("type").and_then(Value::as_str) == Some("url")) - .filter_map(|source| source.get("url").and_then(Value::as_str)) - .map(|url| json!({ "type": "url", "url": url })) - .collect(); - if !sources.is_empty() { - result["sources"] = json!(sources); - } - } - result -} - -fn image_generation_result(part: &Value, tool_name: String) -> ToolResult { - let (result, is_error) = match part.get("result").filter(|value| !value.is_null()) { - Some(result) => { - let mut output = json!({ "result": result }); - if let Some(prompt) = part.get("prompt").filter(|value| !value.is_null()) { - output["prompt"] = prompt.clone(); - } - (output, None) - } - None => ( - json!(format!( - "Image generation failed (status: {}).", - part.get("status").and_then(Value::as_str).unwrap_or("") - )), - Some(true), - ), - }; - ToolResult { - tool_call_id: part - .get("id") - .and_then(Value::as_str) - .unwrap_or("") - .to_string(), - tool_name, - result, - is_error, - preliminary: None, - dynamic: None, - provider_metadata: None, - } -} - -/// An xAI Responses language model (Grok). -/// -/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the -/// process-wide shared `Client` internally (RFC-0009 §4.1). -pub struct XaiResponsesModel { - model_id: String, - config: EndpointConfig, -} - -impl XaiResponsesModel { - pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { - Self { model_id, config } - } -} - -#[async_trait] -impl LanguageModel for XaiResponsesModel { - fn provider(&self) -> &str { - &self.config.provider - } - - fn model_id(&self) -> &str { - &self.model_id - } - - /// Image, PDF and text URLs go to the API as they are, as `input_image` - /// or `input_file` (`xaiResponsesSupportedUrls`, - /// xai-responses-language-model.ts:142-149, :181). - async fn supported_urls(&self) -> SupportedUrls { - let http = regex::Regex::new(r"^https?://.*$").expect("static pattern"); - SupportedUrls( - ["image/*", "application/pdf", "text/*"] - .into_iter() - .map(|media_type| (media_type.to_string(), vec![http.clone()])) - .collect(), - ) - } - - async fn do_generate(&self, options: &CallOptions) -> Result { - let exchange = self.config.exchange(options.headers.as_ref()).await?; - let request_result = build_responses_request_body(&self.model_id, options, false)?; - let body = exchange.transform_body(request_result.body); - let provider_tool_names = request_result.provider_tool_names; - - let resp = aimux_provider_utils::post_json_to_api( - exchange.request(exchange.url("/responses"), options), - body.clone(), - super::xai_successful_response_handler::(), - super::xai_failed_response_handler(), - ) - .await?; - - let response_headers = resp.response_headers; - - let raw_value = resp.raw_value.unwrap_or(Value::Null); - let data = resp.value; - - let mut content: Vec = Vec::new(); - let mut has_function_call = false; - - for part in &data.output { - let part_type = part.get("type").and_then(|v| v.as_str()).unwrap_or(""); - - if part_type == "image_generation_call" { - let tool_name = resolve_tool_name(part_type, None, &provider_tool_names); - content.push(GenerateContent::ToolCall(RawToolCall { - tool_call_id: part - .get("id") - .and_then(Value::as_str) - .unwrap_or("") - .to_string(), - tool_name: tool_name.clone(), - input: "{}".to_string(), - provider_executed: Some(true), - dynamic: None, - provider_metadata: None, - })); - content.push(GenerateContent::ToolResult(image_generation_result( - part, tool_name, - ))); - continue; - } - - // ── file_search_call ── - if part_type == "file_search_call" { - let tool_name = resolve_tool_name(part_type, None, &provider_tool_names); - let part_id = part.get("id").and_then(|v| v.as_str()).unwrap_or(""); - - content.push(GenerateContent::ToolCall(RawToolCall { - tool_call_id: part_id.to_string(), - tool_name: tool_name.clone(), - input: String::new(), - provider_executed: Some(true), - dynamic: None, - provider_metadata: None, - })); - - let queries = part - .get("queries") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - let results = - part.get("results") - .and_then(|v| v.as_array()) - .map(|arr| { - json!(arr.iter().map(|r| json!({ - "fileId": r.get("file_id").cloned().unwrap_or(Value::Null), - "filename": r.get("filename").cloned().unwrap_or(Value::Null), - "score": r.get("score").cloned().unwrap_or(Value::Null), - "text": r.get("text").cloned().unwrap_or(Value::Null), - })).collect::>()) - }) - .unwrap_or(Value::Null); - - content.push(GenerateContent::ToolResult(ToolResult { - tool_call_id: part_id.to_string(), - tool_name, - result: json!({ "queries": queries, "results": results }), - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: None, - })); - continue; - } - - // ── Server-side tool calls ── - if matches!( - part_type, - "web_search_call" - | "x_search_call" - | "code_interpreter_call" - | "code_execution_call" - | "view_image_call" - | "view_x_video_call" - | "custom_tool_call" - | "mcp_call" - ) { - let part_name = part.get("name").and_then(|v| v.as_str()); - let tool_name = resolve_tool_name(part_type, part_name, &provider_tool_names); - let tool_input = get_tool_input(part_type, part); - let part_id = part.get("id").and_then(|v| v.as_str()).unwrap_or(""); - - content.push(GenerateContent::ToolCall(RawToolCall { - tool_call_id: part_id.to_string(), - tool_name: tool_name.clone(), - input: tool_input, - provider_executed: Some(true), - dynamic: None, - provider_metadata: None, - })); - if part_type == "web_search_call" { - content.push(GenerateContent::ToolResult(ToolResult { - tool_call_id: part_id.to_string(), - tool_name, - result: map_web_search_action(part.get("action").unwrap_or(&Value::Null)), - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: None, - })); - } - continue; - } - - match part_type { - "message" => { - let content_parts = part - .get("content") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - for cp in &content_parts { - if let Some(text) = cp.get("text").and_then(|v| v.as_str()) - && !text.is_empty() - { - content.push(GenerateContent::Text { - text: text.to_string(), - provider_metadata: None, - }); - } - if let Some(annotations) = cp.get("annotations").and_then(|v| v.as_array()) - { - for ann in annotations { - if ann.get("type").and_then(|v| v.as_str()) == Some("url_citation") - { - let url = ann.get("url").and_then(|v| v.as_str()).unwrap_or(""); - let title = ann - .get("title") - .and_then(|v| v.as_str()) - .unwrap_or(url) - .to_string(); - content.push(GenerateContent::Source(Source::Url { - id: generate_source_id(), - url: url.to_string(), - title: Some(title), - provider_metadata: None, - })); - } - } - } - } - } - "function_call" => { - has_function_call = true; - let call_id = part.get("call_id").and_then(|v| v.as_str()).unwrap_or(""); - let name = part.get("name").and_then(|v| v.as_str()).unwrap_or(""); - let arguments = part.get("arguments").and_then(|v| v.as_str()).unwrap_or(""); - let input = arguments.to_string(); - content.push(GenerateContent::ToolCall(RawToolCall { - tool_call_id: call_id.to_string(), - tool_name: name.to_string(), - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - } - "reasoning" => { - let summary = part.get("summary").and_then(|v| v.as_array()); - let reasoning_content = part.get("content").and_then(|v| v.as_array()); - - let texts: Vec = if let Some(summary) = summary { - if !summary.is_empty() { - summary - .iter() - .filter_map(|s| { - s.get("text") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .filter(|t| !t.is_empty()) - .collect() - } else { - reasoning_content - .map(|c| { - c.iter() - .filter_map(|s| { - s.get("text") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .filter(|t| !t.is_empty()) - .collect() - }) - .unwrap_or_default() - } - } else { - Vec::new() - }; - - let reasoning_text = texts.join(""); - let encrypted_content = part - .get("encrypted_content") - .and_then(|v| v.as_str()) - .filter(|v| !v.is_empty()); - let item_id = part - .get("id") - .and_then(|v| v.as_str()) - .filter(|v| !v.is_empty()); - - if !reasoning_text.is_empty() || encrypted_content.is_some() { - let mut meta = json!({}); - if let Some(ec) = encrypted_content { - meta["reasoningEncryptedContent"] = json!(ec); - } - if let Some(id) = item_id { - meta["itemId"] = json!(id); - } - content.push(GenerateContent::Reasoning(ReasoningOutput { - text: reasoning_text, - provider_metadata: (encrypted_content.is_some() || item_id.is_some()) - .then(|| crate::xai::options::xai_metadata(meta)), - })); - } - } - _ => {} - } - } - - let finish_reason = FinishReason { - unified: if has_function_call { - FinishReasonUnified::ToolCalls - } else { - data.status - .as_deref() - .map(map_xai_responses_finish_reason) - .unwrap_or(FinishReasonUnified::Other) - }, - raw: data.status.clone(), - }; - - let usage = data.usage.as_ref().map_or_else(zero_usage, |u| { - let mut usage = convert_xai_responses_usage(u); - usage.raw = raw_value - .get("usage") - .and_then(|value| value.as_object().cloned()); - usage - }); - let provider_metadata = response_provider_metadata( - data.usage.as_ref().and_then(|u| u.cost_in_usd_ticks), - data.service_tier.as_deref(), - data.prompt_cache_key.as_deref(), - data.safety_identifier.as_deref(), - ); - - let timestamp = data - .created_at - .and_then(|c| aimux_core::Timestamp::from_unix_seconds(c as f64)); - - Ok(GenerateResult { - content, - finish_reason, - usage, - warnings: request_result.warnings, - provider_metadata, - response: Some(aimux_core::shared::ResponseInfo { - id: data.id, - timestamp, - model_id: data.model, - headers: Some(response_headers), - body: Some(raw_value), - }), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - }) - } - - async fn do_stream(&self, options: &CallOptions) -> Result { - let exchange = self.config.exchange(options.headers.as_ref()).await?; - let request_result = build_responses_request_body(&self.model_id, options, true)?; - let body = exchange.transform_body(request_result.body); - let warnings = request_result.warnings; - let provider_tool_names = request_result.provider_tool_names; - let endpoint = exchange.url("/responses"); - - let resp = aimux_provider_utils::post_json_to_api( - exchange.request(endpoint.clone(), options), - body.clone(), - super::xai_event_source_response_handler::(), - super::xai_failed_response_handler(), - ) - .await?; - - let response_headers = resp.response_headers; - - let mut sse_stream = resp.value; - let first_event = match sse_stream.next().await { - Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), - first => first, - }; - if let Some(Ok(event)) = &first_event - && types::event_type(event) == "error" - { - return Err(super::xai_stream_error( - event, - &endpoint, - body.clone(), - response_headers.clone(), - )); - } - let stream = pipe_through( - futures::stream::iter(first_event).chain(sse_stream), - XaiResponsesStream { - ctx: StreamContext::new(options, warnings, endpoint, &body, &response_headers), - provider_tool_names, - ..XaiResponsesStream::default() - }, - ); - - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) - } -} - -/// The `TransformStream` of the xAI Responses `doStream`. -#[derive(Default)] -struct XaiResponsesStream { - ctx: StreamContext, - provider_tool_names: HashMap, - final_usage: Option, - /// `other` until the response says otherwise. - final_finish_reason: Option, - cost_in_usd_ticks: Option, - service_tier: Option, - prompt_cache_key: Option, - safety_identifier: Option, - has_function_call: bool, - metadata_emitted: bool, - /// Open text blocks. - content_blocks: Vec, - seen_tool_calls: std::collections::HashSet, - active_reasoning: HashMap, - /// Ongoing function calls: output_index -> (tool_call_id, tool_name). - ongoing_tool_calls: HashMap, - /// Set by an `error` event: later chunks are ignored, the stream still - /// finishes. - finished: bool, -} - -impl Transformer for XaiResponsesStream { - type Input = Result; - type Output = StreamPart; - - fn start(&mut self, controller: &mut Controller) { - self.ctx.start(controller); - } - - fn transform(&mut self, event: Result, controller: &mut Controller) { - self.transform_event(event, controller); - if self.finished { - // An error event ends the response: the rest of the body is not - // read and the stream finishes now. - self.finish(controller); - controller.terminate(); - } - } - - fn flush(mut self, controller: &mut Controller) { - self.finish(controller); - } -} - -impl XaiResponsesStream { - /// Handle one parsed SSE event. - fn transform_event(&mut self, event: Result, controller: &mut Controller) { - let Self { - ctx: - StreamContext { - include_raw_chunks, - url: ref stream_error_url, - request_body: ref stream_request_body, - response_headers: ref stream_response_headers, - .. - }, - ref provider_tool_names, - ref mut final_usage, - ref mut final_finish_reason, - ref mut cost_in_usd_ticks, - ref mut service_tier, - ref mut prompt_cache_key, - ref mut safety_identifier, - ref mut has_function_call, - ref mut metadata_emitted, - ref mut content_blocks, - ref mut seen_tool_calls, - ref mut active_reasoning, - ref mut ongoing_tool_calls, - ref mut finished, - .. - } = *self; - match event { - Ok(parsed) => { - if include_raw_chunks { - controller.enqueue(StreamPart::Raw { - raw_value: parsed.clone(), - }); - } - let event_type = types::event_type(&parsed); - - // ── response.created / response.in_progress ── - if event_type == "response.created" || event_type == "response.in_progress" { - if !*metadata_emitted { - *metadata_emitted = true; - let response = parsed.get("response").cloned().unwrap_or(Value::Null); - let id = response - .get("id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let model = response - .get("model") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let created_at = response - .get("created_at") - .and_then(serde_json::Value::as_u64); - let timestamp = created_at - .and_then(|c| aimux_core::Timestamp::from_unix_seconds(c as f64)); - controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { - id, - timestamp, - model_id: model, - })); - } - return; - } - - // ── reasoning_summary_part.added ── - if event_type == "response.reasoning_summary_part.added" { - let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("reasoning-{item_id}"); - if !active_reasoning.contains_key(item_id) { - active_reasoning.insert(item_id.to_string(), ()); - controller.enqueue(StreamPart::ReasoningStart { - id: block_id, - provider_metadata: Some(crate::xai::options::xai_metadata( - json!({ "itemId": item_id }), - )), - }); - } - return; - } - - // ── reasoning_summary_text.delta ── - if event_type == "response.reasoning_summary_text.delta" { - let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); - let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("reasoning-{item_id}"); - controller.enqueue(StreamPart::ReasoningDelta { - id: block_id, - delta: delta.to_string(), - provider_metadata: Some(crate::xai::options::xai_metadata( - json!({ "itemId": item_id }), - )), - }); - return; - } - - // ── reasoning_summary_text.done ── - if event_type == "response.reasoning_summary_text.done" { - return; - } - - // ── reasoning_text.delta ── - if event_type == "response.reasoning_text.delta" { - let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); - let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("reasoning-{item_id}"); - if !active_reasoning.contains_key(item_id) { - active_reasoning.insert(item_id.to_string(), ()); - controller.enqueue(StreamPart::ReasoningStart { - id: block_id.clone(), - provider_metadata: Some(crate::xai::options::xai_metadata( - json!({ "itemId": item_id }), - )), - }); - } - controller.enqueue(StreamPart::ReasoningDelta { - id: block_id, - delta: delta.to_string(), - provider_metadata: Some(crate::xai::options::xai_metadata( - json!({ "itemId": item_id }), - )), - }); - return; - } - - // ── reasoning_text.done ── - if event_type == "response.reasoning_text.done" { - return; - } - - // ── output_text.delta ── - if event_type == "response.output_text.delta" { - let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); - let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("text-{item_id}"); - if !content_blocks.contains(&block_id) { - content_blocks.push(block_id.clone()); - controller.text_start(block_id.clone()); - } - controller.text_delta(block_id, delta.to_string()); - return; - } - - // ── output_text.done ── - if event_type == "response.output_text.done" { - if let Some(annotations) = parsed.get("annotations").and_then(|v| v.as_array()) - { - for ann in annotations { - if ann.get("type").and_then(|v| v.as_str()) == Some("url_citation") { - let url = ann.get("url").and_then(|v| v.as_str()).unwrap_or(""); - let title = ann - .get("title") - .and_then(|v| v.as_str()) - .unwrap_or(url) - .to_string(); - controller.enqueue(StreamPart::Source(Source::Url { - id: generate_source_id(), - url: url.to_string(), - title: Some(title), - provider_metadata: None, - })); - } - } - } - return; - } - - // ── output_text.annotation.added ── - if event_type == "response.output_text.annotation.added" { - let annotation = parsed.get("annotation").cloned().unwrap_or(Value::Null); - if annotation.get("type").and_then(|v| v.as_str()) == Some("url_citation") { - let url = annotation.get("url").and_then(|v| v.as_str()).unwrap_or(""); - let title = annotation - .get("title") - .and_then(|v| v.as_str()) - .unwrap_or(url) - .to_string(); - controller.enqueue(StreamPart::Source(Source::Url { - id: generate_source_id(), - url: url.to_string(), - title: Some(title), - provider_metadata: None, - })); - } - return; - } - - // ── response.done / response.completed / response.incomplete ── - if event_type == "response.done" - || event_type == "response.completed" - || event_type == "response.incomplete" - { - let response = parsed.get("response").cloned().unwrap_or(Value::Null); - if let Some(usage) = response.get("usage") - && let Ok(u) = - serde_json::from_value::(usage.clone()) - { - let mut converted = convert_xai_responses_usage(&u); - converted.raw = usage.as_object().cloned(); - *final_usage = Some(converted); - *cost_in_usd_ticks = u.cost_in_usd_ticks; - } - - *service_tier = response - .get("service_tier") - .and_then(Value::as_str) - .map(str::to_string); - *prompt_cache_key = response - .get("prompt_cache_key") - .and_then(Value::as_str) - .map(str::to_string); - *safety_identifier = response - .get("safety_identifier") - .and_then(Value::as_str) - .map(str::to_string); - - if event_type == "response.incomplete" { - let reason = response - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|v| v.as_str()); - *final_finish_reason = Some(FinishReason { - unified: reason - .map(map_xai_responses_finish_reason) - .unwrap_or(FinishReasonUnified::Other), - raw: reason - .map(std::string::ToString::to_string) - .or(Some("incomplete".to_string())), - }); - } else if let Some(status) = response.get("status").and_then(|v| v.as_str()) { - *final_finish_reason = Some(FinishReason { - unified: if *has_function_call { - FinishReasonUnified::ToolCalls - } else { - map_xai_responses_finish_reason(status) - }, - raw: Some(status.to_string()), - }); - } - return; - } - - // ── response.failed ── - if event_type == "response.failed" { - let response = parsed.get("response").cloned().unwrap_or(Value::Null); - let reason = response - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|v| v.as_str()); - *final_finish_reason = Some(FinishReason { - unified: reason - .map(map_xai_responses_finish_reason) - .unwrap_or(FinishReasonUnified::Error), - raw: reason - .map(std::string::ToString::to_string) - .or(Some("error".to_string())), - }); - if let Some(usage) = response.get("usage") - && let Ok(u) = - serde_json::from_value::(usage.clone()) - { - let mut converted = convert_xai_responses_usage(&u); - converted.raw = usage.as_object().cloned(); - *final_usage = Some(converted); - } - if response.get("error").is_some_and(|error| !error.is_null()) { - controller.enqueue(StreamPart::Error { - error: super::xai_stream_error( - &parsed, - stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ), - }); - } - return; - } - - // ── error event ── - if event_type == "error" { - controller.enqueue(StreamPart::Error { - error: super::xai_stream_error( - &parsed, - stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ), - }); - *final_finish_reason = Some(FinishReason { - unified: FinishReasonUnified::Error, - raw: Some("error".to_string()), - }); - *finished = true; - return; - } - - // ── custom_tool_call_input.delta / .done ── - if event_type == "response.custom_tool_call_input.delta" - || event_type == "response.custom_tool_call_input.done" - { - return; - } - - // ── function_call_arguments.delta ── - if event_type == "response.function_call_arguments.delta" { - let output_index = parsed - .get("output_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0); - let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); - if let Some((tool_call_id, _)) = ongoing_tool_calls.get(&output_index) { - controller.tool_input_delta(tool_call_id.clone(), delta.to_string()); - } - return; - } - - // ── function_call_arguments.done ── - if event_type == "response.function_call_arguments.done" { - return; - } - - // ── output_item.added / output_item.done ── - if event_type == "response.output_item.added" - || event_type == "response.output_item.done" - || matches!( - event_type, - "response.image_generation_call.in_progress" - | "response.image_generation_call.generating" - | "response.image_generation_call.completed" - ) - { - let item = if event_type.starts_with("response.image_generation_call.") { - json!({ "type": "image_generation_call", "id": parsed.get("item_id") }) - } else { - parsed.get("item").cloned().unwrap_or(Value::Null) - }; - let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); - let output_index = parsed - .get("output_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0); - - // ── reasoning item ── - if item_type == "reasoning" { - if event_type == "response.output_item.done" { - let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("reasoning-{part_id}"); - let encrypted = item - .get("encrypted_content") - .and_then(|v| v.as_str()) - .filter(|v| !v.is_empty()); - - if !active_reasoning.contains_key(part_id) { - active_reasoning.insert(part_id.to_string(), ()); - controller.enqueue(StreamPart::ReasoningStart { - id: block_id.clone(), - provider_metadata: Some(crate::xai::options::xai_metadata( - if part_id.is_empty() { - json!({}) - } else { - json!({ "itemId": part_id }) - }, - )), - }); - } - - let mut meta = json!({}); - if !part_id.is_empty() { - meta["itemId"] = json!(part_id); - } - if let Some(ec) = encrypted { - meta["reasoningEncryptedContent"] = json!(ec); - } - controller.enqueue(StreamPart::ReasoningEnd { - id: block_id, - provider_metadata: Some(crate::xai::options::xai_metadata(json!( - meta - ))), - }); - active_reasoning.remove(part_id); - } - return; - } - - // ── file_search_call item ── - if item_type == "file_search_call" || item_type == "image_generation_call" { - let tool_name = resolve_tool_name(item_type, None, provider_tool_names); - let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); - let input = if item_type == "image_generation_call" { - "{}" - } else { - "" - }; - - if !seen_tool_calls.contains(part_id) { - seen_tool_calls.insert(part_id.to_string()); - controller.tool_input_start(part_id.to_string(), tool_name.clone()); - controller.tool_input_delta(part_id.to_string(), input.to_string()); - controller.tool_input_end(part_id.to_string()); - controller.enqueue(StreamPart::ToolCall(RawToolCall { - tool_call_id: part_id.to_string(), - tool_name: tool_name.clone(), - input: input.to_string(), - provider_executed: Some(true), - dynamic: None, - provider_metadata: None, - })); - } - - if event_type == "response.output_item.done" { - if item_type == "image_generation_call" { - controller.enqueue(StreamPart::ToolResult( - image_generation_result(&item, tool_name), - )); - return; - } - let queries = item - .get("queries") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - let results = item.get("results").and_then(|v| v.as_array()).map(|arr| { - json!(arr.iter().map(|r| json!({ - "fileId": r.get("file_id").cloned().unwrap_or(Value::Null), - "filename": r.get("filename").cloned().unwrap_or(Value::Null), - "score": r.get("score").cloned().unwrap_or(Value::Null), - "text": r.get("text").cloned().unwrap_or(Value::Null), - })).collect::>()) - }).unwrap_or(Value::Null); - controller.enqueue(StreamPart::ToolResult(ToolResult { - tool_call_id: part_id.to_string(), - tool_name, - result: json!({ "queries": queries, "results": results }), - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: None, - })); - } - return; - } - - // ── Server-side tool calls ── - if matches!( - item_type, - "web_search_call" - | "x_search_call" - | "code_interpreter_call" - | "code_execution_call" - | "view_image_call" - | "view_x_video_call" - | "custom_tool_call" - | "mcp_call" - ) { - let item_name = item.get("name").and_then(|v| v.as_str()); - let tool_name = - resolve_tool_name(item_type, item_name, provider_tool_names); - let tool_input = get_tool_input(item_type, &item); - let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); - - let should_emit = if item_type == "custom_tool_call" { - event_type == "response.output_item.done" - } else { - !seen_tool_calls.contains(part_id) - }; - - if should_emit && !seen_tool_calls.contains(part_id) { - seen_tool_calls.insert(part_id.to_string()); - controller.tool_input_start(part_id.to_string(), tool_name.clone()); - controller.tool_input_delta(part_id.to_string(), tool_input.clone()); - controller.tool_input_end(part_id.to_string()); - controller.enqueue(StreamPart::ToolCall(RawToolCall { - tool_call_id: part_id.to_string(), - tool_name: tool_name.clone(), - input: tool_input, - provider_executed: Some(true), - dynamic: None, - provider_metadata: None, - })); - } - - if event_type == "response.output_item.done" { - controller.enqueue(StreamPart::ToolResult(ToolResult { - tool_call_id: part_id.to_string(), - tool_name, - result: if item_type == "web_search_call" { - map_web_search_action( - item.get("action").unwrap_or(&Value::Null), - ) - } else { - json!({}) - }, - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: None, - })); - } - return; - } - - // ── message item ── - if item_type == "message" { - let content_parts = item - .get("content") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - for cp in &content_parts { - if let Some(text) = cp.get("text").and_then(|v| v.as_str()) - && !text.is_empty() - { - let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("text-{part_id}"); - if !content_blocks.contains(&block_id) { - content_blocks.push(block_id.clone()); - controller.text_start(block_id.clone()); - controller.text_delta(block_id, text.to_string()); - } - } - if let Some(annotations) = - cp.get("annotations").and_then(|v| v.as_array()) - { - for ann in annotations { - if ann.get("type").and_then(|v| v.as_str()) - == Some("url_citation") - { - let url = - ann.get("url").and_then(|v| v.as_str()).unwrap_or(""); - let title = ann - .get("title") - .and_then(|v| v.as_str()) - .unwrap_or(url) - .to_string(); - controller.enqueue(StreamPart::Source(Source::Url { - id: generate_source_id(), - url: url.to_string(), - title: Some(title), - provider_metadata: None, - })); - } - } - } - } - return; - } - - // ── function_call item ── - if item_type == "function_call" { - let call_id = item.get("call_id").and_then(|v| v.as_str()).unwrap_or(""); - let name = item.get("name").and_then(|v| v.as_str()).unwrap_or(""); - - if event_type == "response.output_item.added" { - ongoing_tool_calls - .insert(output_index, (call_id.to_string(), name.to_string())); - controller.tool_input_start(call_id.to_string(), name.to_string()); - } else { - // output_item.done - *has_function_call = true; - ongoing_tool_calls.remove(&output_index); - let arguments = - item.get("arguments").and_then(|v| v.as_str()).unwrap_or(""); - controller.tool_input_end(call_id.to_string()); - let input = arguments.to_string(); - controller.enqueue(StreamPart::ToolCall(RawToolCall { - tool_call_id: call_id.to_string(), - tool_name: name.to_string(), - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - } - } - } - - // All other event types (web_search_call.in_progress, etc.) are ignored. - } - Err(error) => { - if error.is_recoverable_stream_error() { - if include_raw_chunks { - controller.enqueue(StreamPart::Raw { - raw_value: Value::Null, - }); - } - controller.enqueue(StreamPart::Error { error }); - } else { - controller.error(error); - } - } - } - } - - /// Close the open text blocks and emit `Finish`: at the end of the - /// body, or right after an error event. - fn finish(&mut self, controller: &mut Controller) { - // Close any remaining open text blocks. - for block_id in std::mem::take(&mut self.content_blocks) { - controller.text_end(block_id); - } - - // Final part: Finish. - let provider_meta = response_provider_metadata( - self.cost_in_usd_ticks, - self.service_tier.as_deref(), - self.prompt_cache_key.as_deref(), - self.safety_identifier.as_deref(), - ); - - controller.enqueue(StreamPart::Finish { - finish_reason: self.final_finish_reason.clone().unwrap_or(FINISH_OTHER), - usage: self.final_usage.take().unwrap_or_else(zero_usage), - provider_metadata: provider_meta, - }); - } -} +pub use responses_language_model::{ + ResponsesRequestBodyResult, WEB_SEARCH_SUB_TOOLS, X_SEARCH_SUB_TOOLS, XaiResponsesModel, + build_responses_request_body, get_tool_input, resolve_tool_name, +}; +pub use responses_prepare_tools::{PreparedResponsesTools, prepare_responses_tools}; diff --git a/aimux-providers/src/xai/responses/types.rs b/aimux-providers/src/xai/responses/responses_api.rs similarity index 100% rename from aimux-providers/src/xai/responses/types.rs rename to aimux-providers/src/xai/responses/responses_api.rs diff --git a/aimux-providers/src/xai/responses/responses_language_model.rs b/aimux-providers/src/xai/responses/responses_language_model.rs new file mode 100644 index 000000000..ea71f7bbd --- /dev/null +++ b/aimux-providers/src/xai/responses/responses_language_model.rs @@ -0,0 +1,1536 @@ +//! The xAI Responses language model (`xai-responses-language-model.ts`). +//! +//! Implements `LanguageModel` against xAI's `/responses` endpoint. +//! Mirrors the TS `XaiResponsesLanguageModel`. +//! +//! Key differences from the chat model: +//! - Uses `input` (not `messages`) with Responses API format +//! - `reasoning: { effort, summary }` object (not `reasoning_effort` string) +//! - `text.format` for structured output (not `response_format`) +//! - Provider-executed tools (web_search, x_search, code_execution, etc.) +//! - ~65 streaming event types (response.created, response.output_text.delta, etc.) +//! - `cost_in_usd_ticks` in providerMetadata +//! - Citations via annotations → sources +//! +//! The request body (`getArgs`) and the tool-name helpers of `doGenerate` +//! live here too; `get-response-metadata.ts` is inlined in both calls. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::options::{CallOptions, ResponseFormat}; +use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::{RawToolCall, ToolResult}; +use aimux_core::types::{ + FinishReason, FinishReasonUnified, ReasoningEffort, ResponseMetadata, Usage, Warning, +}; +use aimux_provider_utils::{Transformer, pipe_through}; +use async_trait::async_trait; +use futures::StreamExt; +use serde_json::{Value, json}; + +use crate::shared::EndpointConfig; +use crate::stream::{Controller, Emit, FINISH_OTHER, StreamContext}; +use crate::xai::supports_reasoning_effort; + +use super::convert_responses_usage::convert_xai_responses_usage; +use super::convert_to_responses_input::convert_to_xai_responses_input; +use super::map_responses_finish_reason::map_xai_responses_finish_reason; +use super::responses_api as types; +use super::responses_language_model_options::{validate_response_options, xai_option}; +use super::responses_prepare_tools::prepare_responses_tools; + +/// Global counter for generating unique source IDs. +static SOURCE_ID_COUNTER: AtomicU64 = AtomicU64::new(0); + +fn generate_source_id() -> String { + format!("id-{}", SOURCE_ID_COUNTER.fetch_add(1, Ordering::SeqCst)) +} + +fn zero_usage() -> Usage { + Usage { + input_tokens: aimux_core::types::InputTokenUsage { + total: Some(0), + no_cache: Some(0), + cache_read: Some(0), + cache_write: Some(0), + }, + output_tokens: aimux_core::types::OutputTokenUsage { + total: Some(0), + text: Some(0), + reasoning: Some(0), + }, + raw: None, + } +} + +fn response_provider_metadata( + cost: Option, + service_tier: Option<&str>, + prompt_cache_key: Option<&str>, + safety_identifier: Option<&str>, +) -> Option { + let mut metadata = serde_json::Map::new(); + if let Some(cost) = cost { + metadata.insert("costInUsdTicks".into(), json!(cost)); + } + for (key, value) in [ + ("serviceTier", service_tier), + ("promptCacheKey", prompt_cache_key), + ("safetyIdentifier", safety_identifier), + ] { + if let Some(value) = value { + metadata.insert(key.into(), json!(value)); + } + } + (!metadata.is_empty()).then(|| crate::xai::options::xai_metadata(Value::Object(metadata))) +} + +fn map_web_search_action(action: &Value) -> Value { + let Some(kind) = action.get("type").and_then(Value::as_str) else { + return json!({}); + }; + let fields: &[&str] = match kind { + "search" => &["query"], + "open_page" => &["url"], + "find_in_page" => &["url", "pattern"], + _ => return json!({}), + }; + if fields.iter().any(|key| { + action + .get(key) + .is_some_and(|v| !v.is_null() && !v.is_string()) + }) || action + .get("sources") + .is_some_and(|v| !v.is_null() && !v.is_array()) + || (kind == "search" + && action.get("queries").is_some_and(|v| { + !v.is_null() && !v.as_array().is_some_and(|a| a.iter().all(Value::is_string)) + })) + { + return json!({}); + } + let mapped_type = match kind { + "open_page" => "openPage", + "find_in_page" => "findInPage", + _ => kind, + }; + let mut mapped_action = json!({ "type": mapped_type }); + for key in fields { + if let Some(value) = action.get(key) + && (kind != "search" || !value.is_null()) + { + mapped_action[key] = value.clone(); + } + } + if kind == "search" + && let Some(queries) = action.get("queries").filter(|v| !v.is_null()) + { + mapped_action["queries"] = queries.clone(); + } + let mut result = json!({ "action": mapped_action }); + if let Some(sources) = action.get("sources").and_then(Value::as_array) { + let sources: Vec = sources + .iter() + .filter(|source| source.get("type").and_then(Value::as_str) == Some("url")) + .filter_map(|source| source.get("url").and_then(Value::as_str)) + .map(|url| json!({ "type": "url", "url": url })) + .collect(); + if !sources.is_empty() { + result["sources"] = json!(sources); + } + } + result +} + +fn image_generation_result(part: &Value, tool_name: String) -> ToolResult { + let (result, is_error) = match part.get("result").filter(|value| !value.is_null()) { + Some(result) => { + let mut output = json!({ "result": result }); + if let Some(prompt) = part.get("prompt").filter(|value| !value.is_null()) { + output["prompt"] = prompt.clone(); + } + (output, None) + } + None => ( + json!(format!( + "Image generation failed (status: {}).", + part.get("status").and_then(Value::as_str).unwrap_or("") + )), + Some(true), + ), + }; + ToolResult { + tool_call_id: part + .get("id") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(), + tool_name, + result, + is_error, + preliminary: None, + dynamic: None, + provider_metadata: None, + } +} + +/// An xAI Responses language model (Grok). +/// +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the +/// process-wide shared `Client` internally (RFC-0009 §4.1). +pub struct XaiResponsesModel { + model_id: String, + config: EndpointConfig, +} + +impl XaiResponsesModel { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { + Self { model_id, config } + } +} + +#[async_trait] +impl LanguageModel for XaiResponsesModel { + fn provider(&self) -> &str { + &self.config.provider + } + + fn model_id(&self) -> &str { + &self.model_id + } + + /// Image, PDF and text URLs go to the API as they are, as `input_image` + /// or `input_file` (`xaiResponsesSupportedUrls`, + /// xai-responses-language-model.ts:142-149, :181). + async fn supported_urls(&self) -> SupportedUrls { + let http = regex::Regex::new(r"^https?://.*$").expect("static pattern"); + SupportedUrls( + ["image/*", "application/pdf", "text/*"] + .into_iter() + .map(|media_type| (media_type.to_string(), vec![http.clone()])) + .collect(), + ) + } + + async fn do_generate(&self, options: &CallOptions) -> Result { + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let request_result = build_responses_request_body(&self.model_id, options, false)?; + let body = exchange.transform_body(request_result.body); + let provider_tool_names = request_result.provider_tool_names; + + let resp = aimux_provider_utils::post_json_to_api( + exchange.request(exchange.url("/responses"), options), + body.clone(), + crate::xai::error::xai_successful_response_handler::(), + crate::xai::error::xai_failed_response_handler(), + ) + .await?; + + let response_headers = resp.response_headers; + + let raw_value = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; + + let mut content: Vec = Vec::new(); + let mut has_function_call = false; + + for part in &data.output { + let part_type = part.get("type").and_then(|v| v.as_str()).unwrap_or(""); + + if part_type == "image_generation_call" { + let tool_name = resolve_tool_name(part_type, None, &provider_tool_names); + content.push(GenerateContent::ToolCall(RawToolCall { + tool_call_id: part + .get("id") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(), + tool_name: tool_name.clone(), + input: "{}".to_string(), + provider_executed: Some(true), + dynamic: None, + provider_metadata: None, + })); + content.push(GenerateContent::ToolResult(image_generation_result( + part, tool_name, + ))); + continue; + } + + // ── file_search_call ── + if part_type == "file_search_call" { + let tool_name = resolve_tool_name(part_type, None, &provider_tool_names); + let part_id = part.get("id").and_then(|v| v.as_str()).unwrap_or(""); + + content.push(GenerateContent::ToolCall(RawToolCall { + tool_call_id: part_id.to_string(), + tool_name: tool_name.clone(), + input: String::new(), + provider_executed: Some(true), + dynamic: None, + provider_metadata: None, + })); + + let queries = part + .get("queries") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + let results = + part.get("results") + .and_then(|v| v.as_array()) + .map(|arr| { + json!(arr.iter().map(|r| json!({ + "fileId": r.get("file_id").cloned().unwrap_or(Value::Null), + "filename": r.get("filename").cloned().unwrap_or(Value::Null), + "score": r.get("score").cloned().unwrap_or(Value::Null), + "text": r.get("text").cloned().unwrap_or(Value::Null), + })).collect::>()) + }) + .unwrap_or(Value::Null); + + content.push(GenerateContent::ToolResult(ToolResult { + tool_call_id: part_id.to_string(), + tool_name, + result: json!({ "queries": queries, "results": results }), + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: None, + })); + continue; + } + + // ── Server-side tool calls ── + if matches!( + part_type, + "web_search_call" + | "x_search_call" + | "code_interpreter_call" + | "code_execution_call" + | "view_image_call" + | "view_x_video_call" + | "custom_tool_call" + | "mcp_call" + ) { + let part_name = part.get("name").and_then(|v| v.as_str()); + let tool_name = resolve_tool_name(part_type, part_name, &provider_tool_names); + let tool_input = get_tool_input(part_type, part); + let part_id = part.get("id").and_then(|v| v.as_str()).unwrap_or(""); + + content.push(GenerateContent::ToolCall(RawToolCall { + tool_call_id: part_id.to_string(), + tool_name: tool_name.clone(), + input: tool_input, + provider_executed: Some(true), + dynamic: None, + provider_metadata: None, + })); + if part_type == "web_search_call" { + content.push(GenerateContent::ToolResult(ToolResult { + tool_call_id: part_id.to_string(), + tool_name, + result: map_web_search_action(part.get("action").unwrap_or(&Value::Null)), + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: None, + })); + } + continue; + } + + match part_type { + "message" => { + let content_parts = part + .get("content") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + for cp in &content_parts { + if let Some(text) = cp.get("text").and_then(|v| v.as_str()) + && !text.is_empty() + { + content.push(GenerateContent::Text { + text: text.to_string(), + provider_metadata: None, + }); + } + if let Some(annotations) = cp.get("annotations").and_then(|v| v.as_array()) + { + for ann in annotations { + if ann.get("type").and_then(|v| v.as_str()) == Some("url_citation") + { + let url = ann.get("url").and_then(|v| v.as_str()).unwrap_or(""); + let title = ann + .get("title") + .and_then(|v| v.as_str()) + .unwrap_or(url) + .to_string(); + content.push(GenerateContent::Source(Source::Url { + id: generate_source_id(), + url: url.to_string(), + title: Some(title), + provider_metadata: None, + })); + } + } + } + } + } + "function_call" => { + has_function_call = true; + let call_id = part.get("call_id").and_then(|v| v.as_str()).unwrap_or(""); + let name = part.get("name").and_then(|v| v.as_str()).unwrap_or(""); + let arguments = part.get("arguments").and_then(|v| v.as_str()).unwrap_or(""); + let input = arguments.to_string(); + content.push(GenerateContent::ToolCall(RawToolCall { + tool_call_id: call_id.to_string(), + tool_name: name.to_string(), + input, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); + } + "reasoning" => { + let summary = part.get("summary").and_then(|v| v.as_array()); + let reasoning_content = part.get("content").and_then(|v| v.as_array()); + + let texts: Vec = if let Some(summary) = summary { + if !summary.is_empty() { + summary + .iter() + .filter_map(|s| { + s.get("text") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string) + }) + .filter(|t| !t.is_empty()) + .collect() + } else { + reasoning_content + .map(|c| { + c.iter() + .filter_map(|s| { + s.get("text") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string) + }) + .filter(|t| !t.is_empty()) + .collect() + }) + .unwrap_or_default() + } + } else { + Vec::new() + }; + + let reasoning_text = texts.join(""); + let encrypted_content = part + .get("encrypted_content") + .and_then(|v| v.as_str()) + .filter(|v| !v.is_empty()); + let item_id = part + .get("id") + .and_then(|v| v.as_str()) + .filter(|v| !v.is_empty()); + + if !reasoning_text.is_empty() || encrypted_content.is_some() { + let mut meta = json!({}); + if let Some(ec) = encrypted_content { + meta["reasoningEncryptedContent"] = json!(ec); + } + if let Some(id) = item_id { + meta["itemId"] = json!(id); + } + content.push(GenerateContent::Reasoning(ReasoningOutput { + text: reasoning_text, + provider_metadata: (encrypted_content.is_some() || item_id.is_some()) + .then(|| crate::xai::options::xai_metadata(meta)), + })); + } + } + _ => {} + } + } + + let finish_reason = FinishReason { + unified: if has_function_call { + FinishReasonUnified::ToolCalls + } else { + data.status + .as_deref() + .map(map_xai_responses_finish_reason) + .unwrap_or(FinishReasonUnified::Other) + }, + raw: data.status.clone(), + }; + + let usage = data.usage.as_ref().map_or_else(zero_usage, |u| { + let mut usage = convert_xai_responses_usage(u); + usage.raw = raw_value + .get("usage") + .and_then(|value| value.as_object().cloned()); + usage + }); + let provider_metadata = response_provider_metadata( + data.usage.as_ref().and_then(|u| u.cost_in_usd_ticks), + data.service_tier.as_deref(), + data.prompt_cache_key.as_deref(), + data.safety_identifier.as_deref(), + ); + + let timestamp = data + .created_at + .and_then(|c| aimux_core::Timestamp::from_unix_seconds(c as f64)); + + Ok(GenerateResult { + content, + finish_reason, + usage, + warnings: request_result.warnings, + provider_metadata, + response: Some(aimux_core::shared::ResponseInfo { + id: data.id, + timestamp, + model_id: data.model, + headers: Some(response_headers), + body: Some(raw_value), + }), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + }) + } + + async fn do_stream(&self, options: &CallOptions) -> Result { + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let request_result = build_responses_request_body(&self.model_id, options, true)?; + let body = exchange.transform_body(request_result.body); + let warnings = request_result.warnings; + let provider_tool_names = request_result.provider_tool_names; + let endpoint = exchange.url("/responses"); + + let resp = aimux_provider_utils::post_json_to_api( + exchange.request(endpoint.clone(), options), + body.clone(), + crate::xai::error::xai_event_source_response_handler::(), + crate::xai::error::xai_failed_response_handler(), + ) + .await?; + + let response_headers = resp.response_headers; + + let mut sse_stream = resp.value; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first => first, + }; + if let Some(Ok(event)) = &first_event + && types::event_type(event) == "error" + { + return Err(crate::xai::error::xai_stream_error( + event, + &endpoint, + body.clone(), + response_headers.clone(), + )); + } + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + XaiResponsesStream { + ctx: StreamContext::new(options, warnings, endpoint, &body, &response_headers), + provider_tool_names, + ..XaiResponsesStream::default() + }, + ); + + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + response: Some(aimux_core::shared::StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} + +/// The `TransformStream` of the xAI Responses `doStream`. +#[derive(Default)] +struct XaiResponsesStream { + ctx: StreamContext, + provider_tool_names: HashMap, + final_usage: Option, + /// `other` until the response says otherwise. + final_finish_reason: Option, + cost_in_usd_ticks: Option, + service_tier: Option, + prompt_cache_key: Option, + safety_identifier: Option, + has_function_call: bool, + metadata_emitted: bool, + /// Open text blocks. + content_blocks: Vec, + seen_tool_calls: std::collections::HashSet, + active_reasoning: HashMap, + /// Ongoing function calls: output_index -> (tool_call_id, tool_name). + ongoing_tool_calls: HashMap, + /// Set by an `error` event: later chunks are ignored, the stream still + /// finishes. + finished: bool, +} + +impl Transformer for XaiResponsesStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } + + fn transform(&mut self, event: Result, controller: &mut Controller) { + self.transform_event(event, controller); + if self.finished { + // An error event ends the response: the rest of the body is not + // read and the stream finishes now. + self.finish(controller); + controller.terminate(); + } + } + + fn flush(mut self, controller: &mut Controller) { + self.finish(controller); + } +} + +impl XaiResponsesStream { + /// Handle one parsed SSE event. + fn transform_event(&mut self, event: Result, controller: &mut Controller) { + let Self { + ctx: + StreamContext { + include_raw_chunks, + url: ref stream_error_url, + request_body: ref stream_request_body, + response_headers: ref stream_response_headers, + .. + }, + ref provider_tool_names, + ref mut final_usage, + ref mut final_finish_reason, + ref mut cost_in_usd_ticks, + ref mut service_tier, + ref mut prompt_cache_key, + ref mut safety_identifier, + ref mut has_function_call, + ref mut metadata_emitted, + ref mut content_blocks, + ref mut seen_tool_calls, + ref mut active_reasoning, + ref mut ongoing_tool_calls, + ref mut finished, + .. + } = *self; + match event { + Ok(parsed) => { + if include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: parsed.clone(), + }); + } + let event_type = types::event_type(&parsed); + + // ── response.created / response.in_progress ── + if event_type == "response.created" || event_type == "response.in_progress" { + if !*metadata_emitted { + *metadata_emitted = true; + let response = parsed.get("response").cloned().unwrap_or(Value::Null); + let id = response + .get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let model = response + .get("model") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let created_at = response + .get("created_at") + .and_then(serde_json::Value::as_u64); + let timestamp = created_at + .and_then(|c| aimux_core::Timestamp::from_unix_seconds(c as f64)); + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id, + timestamp, + model_id: model, + })); + } + return; + } + + // ── reasoning_summary_part.added ── + if event_type == "response.reasoning_summary_part.added" { + let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("reasoning-{item_id}"); + if !active_reasoning.contains_key(item_id) { + active_reasoning.insert(item_id.to_string(), ()); + controller.enqueue(StreamPart::ReasoningStart { + id: block_id, + provider_metadata: Some(crate::xai::options::xai_metadata( + json!({ "itemId": item_id }), + )), + }); + } + return; + } + + // ── reasoning_summary_text.delta ── + if event_type == "response.reasoning_summary_text.delta" { + let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); + let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("reasoning-{item_id}"); + controller.enqueue(StreamPart::ReasoningDelta { + id: block_id, + delta: delta.to_string(), + provider_metadata: Some(crate::xai::options::xai_metadata( + json!({ "itemId": item_id }), + )), + }); + return; + } + + // ── reasoning_summary_text.done ── + if event_type == "response.reasoning_summary_text.done" { + return; + } + + // ── reasoning_text.delta ── + if event_type == "response.reasoning_text.delta" { + let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); + let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("reasoning-{item_id}"); + if !active_reasoning.contains_key(item_id) { + active_reasoning.insert(item_id.to_string(), ()); + controller.enqueue(StreamPart::ReasoningStart { + id: block_id.clone(), + provider_metadata: Some(crate::xai::options::xai_metadata( + json!({ "itemId": item_id }), + )), + }); + } + controller.enqueue(StreamPart::ReasoningDelta { + id: block_id, + delta: delta.to_string(), + provider_metadata: Some(crate::xai::options::xai_metadata( + json!({ "itemId": item_id }), + )), + }); + return; + } + + // ── reasoning_text.done ── + if event_type == "response.reasoning_text.done" { + return; + } + + // ── output_text.delta ── + if event_type == "response.output_text.delta" { + let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); + let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("text-{item_id}"); + if !content_blocks.contains(&block_id) { + content_blocks.push(block_id.clone()); + controller.text_start(block_id.clone()); + } + controller.text_delta(block_id, delta.to_string()); + return; + } + + // ── output_text.done ── + if event_type == "response.output_text.done" { + if let Some(annotations) = parsed.get("annotations").and_then(|v| v.as_array()) + { + for ann in annotations { + if ann.get("type").and_then(|v| v.as_str()) == Some("url_citation") { + let url = ann.get("url").and_then(|v| v.as_str()).unwrap_or(""); + let title = ann + .get("title") + .and_then(|v| v.as_str()) + .unwrap_or(url) + .to_string(); + controller.enqueue(StreamPart::Source(Source::Url { + id: generate_source_id(), + url: url.to_string(), + title: Some(title), + provider_metadata: None, + })); + } + } + } + return; + } + + // ── output_text.annotation.added ── + if event_type == "response.output_text.annotation.added" { + let annotation = parsed.get("annotation").cloned().unwrap_or(Value::Null); + if annotation.get("type").and_then(|v| v.as_str()) == Some("url_citation") { + let url = annotation.get("url").and_then(|v| v.as_str()).unwrap_or(""); + let title = annotation + .get("title") + .and_then(|v| v.as_str()) + .unwrap_or(url) + .to_string(); + controller.enqueue(StreamPart::Source(Source::Url { + id: generate_source_id(), + url: url.to_string(), + title: Some(title), + provider_metadata: None, + })); + } + return; + } + + // ── response.done / response.completed / response.incomplete ── + if event_type == "response.done" + || event_type == "response.completed" + || event_type == "response.incomplete" + { + let response = parsed.get("response").cloned().unwrap_or(Value::Null); + if let Some(usage) = response.get("usage") + && let Ok(u) = + serde_json::from_value::(usage.clone()) + { + let mut converted = convert_xai_responses_usage(&u); + converted.raw = usage.as_object().cloned(); + *final_usage = Some(converted); + *cost_in_usd_ticks = u.cost_in_usd_ticks; + } + + *service_tier = response + .get("service_tier") + .and_then(Value::as_str) + .map(str::to_string); + *prompt_cache_key = response + .get("prompt_cache_key") + .and_then(Value::as_str) + .map(str::to_string); + *safety_identifier = response + .get("safety_identifier") + .and_then(Value::as_str) + .map(str::to_string); + + if event_type == "response.incomplete" { + let reason = response + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|v| v.as_str()); + *final_finish_reason = Some(FinishReason { + unified: reason + .map(map_xai_responses_finish_reason) + .unwrap_or(FinishReasonUnified::Other), + raw: reason + .map(std::string::ToString::to_string) + .or(Some("incomplete".to_string())), + }); + } else if let Some(status) = response.get("status").and_then(|v| v.as_str()) { + *final_finish_reason = Some(FinishReason { + unified: if *has_function_call { + FinishReasonUnified::ToolCalls + } else { + map_xai_responses_finish_reason(status) + }, + raw: Some(status.to_string()), + }); + } + return; + } + + // ── response.failed ── + if event_type == "response.failed" { + let response = parsed.get("response").cloned().unwrap_or(Value::Null); + let reason = response + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|v| v.as_str()); + *final_finish_reason = Some(FinishReason { + unified: reason + .map(map_xai_responses_finish_reason) + .unwrap_or(FinishReasonUnified::Error), + raw: reason + .map(std::string::ToString::to_string) + .or(Some("error".to_string())), + }); + if let Some(usage) = response.get("usage") + && let Ok(u) = + serde_json::from_value::(usage.clone()) + { + let mut converted = convert_xai_responses_usage(&u); + converted.raw = usage.as_object().cloned(); + *final_usage = Some(converted); + } + if response.get("error").is_some_and(|error| !error.is_null()) { + controller.enqueue(StreamPart::Error { + error: crate::xai::error::xai_stream_error( + &parsed, + stream_error_url, + stream_request_body.clone(), + stream_response_headers.clone(), + ), + }); + } + return; + } + + // ── error event ── + if event_type == "error" { + controller.enqueue(StreamPart::Error { + error: crate::xai::error::xai_stream_error( + &parsed, + stream_error_url, + stream_request_body.clone(), + stream_response_headers.clone(), + ), + }); + *final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::Error, + raw: Some("error".to_string()), + }); + *finished = true; + return; + } + + // ── custom_tool_call_input.delta / .done ── + if event_type == "response.custom_tool_call_input.delta" + || event_type == "response.custom_tool_call_input.done" + { + return; + } + + // ── function_call_arguments.delta ── + if event_type == "response.function_call_arguments.delta" { + let output_index = parsed + .get("output_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); + if let Some((tool_call_id, _)) = ongoing_tool_calls.get(&output_index) { + controller.tool_input_delta(tool_call_id.clone(), delta.to_string()); + } + return; + } + + // ── function_call_arguments.done ── + if event_type == "response.function_call_arguments.done" { + return; + } + + // ── output_item.added / output_item.done ── + if event_type == "response.output_item.added" + || event_type == "response.output_item.done" + || matches!( + event_type, + "response.image_generation_call.in_progress" + | "response.image_generation_call.generating" + | "response.image_generation_call.completed" + ) + { + let item = if event_type.starts_with("response.image_generation_call.") { + json!({ "type": "image_generation_call", "id": parsed.get("item_id") }) + } else { + parsed.get("item").cloned().unwrap_or(Value::Null) + }; + let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); + let output_index = parsed + .get("output_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + + // ── reasoning item ── + if item_type == "reasoning" { + if event_type == "response.output_item.done" { + let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("reasoning-{part_id}"); + let encrypted = item + .get("encrypted_content") + .and_then(|v| v.as_str()) + .filter(|v| !v.is_empty()); + + if !active_reasoning.contains_key(part_id) { + active_reasoning.insert(part_id.to_string(), ()); + controller.enqueue(StreamPart::ReasoningStart { + id: block_id.clone(), + provider_metadata: Some(crate::xai::options::xai_metadata( + if part_id.is_empty() { + json!({}) + } else { + json!({ "itemId": part_id }) + }, + )), + }); + } + + let mut meta = json!({}); + if !part_id.is_empty() { + meta["itemId"] = json!(part_id); + } + if let Some(ec) = encrypted { + meta["reasoningEncryptedContent"] = json!(ec); + } + controller.enqueue(StreamPart::ReasoningEnd { + id: block_id, + provider_metadata: Some(crate::xai::options::xai_metadata(json!( + meta + ))), + }); + active_reasoning.remove(part_id); + } + return; + } + + // ── file_search_call item ── + if item_type == "file_search_call" || item_type == "image_generation_call" { + let tool_name = resolve_tool_name(item_type, None, provider_tool_names); + let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); + let input = if item_type == "image_generation_call" { + "{}" + } else { + "" + }; + + if !seen_tool_calls.contains(part_id) { + seen_tool_calls.insert(part_id.to_string()); + controller.tool_input_start(part_id.to_string(), tool_name.clone()); + controller.tool_input_delta(part_id.to_string(), input.to_string()); + controller.tool_input_end(part_id.to_string()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: part_id.to_string(), + tool_name: tool_name.clone(), + input: input.to_string(), + provider_executed: Some(true), + dynamic: None, + provider_metadata: None, + })); + } + + if event_type == "response.output_item.done" { + if item_type == "image_generation_call" { + controller.enqueue(StreamPart::ToolResult( + image_generation_result(&item, tool_name), + )); + return; + } + let queries = item + .get("queries") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + let results = item.get("results").and_then(|v| v.as_array()).map(|arr| { + json!(arr.iter().map(|r| json!({ + "fileId": r.get("file_id").cloned().unwrap_or(Value::Null), + "filename": r.get("filename").cloned().unwrap_or(Value::Null), + "score": r.get("score").cloned().unwrap_or(Value::Null), + "text": r.get("text").cloned().unwrap_or(Value::Null), + })).collect::>()) + }).unwrap_or(Value::Null); + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: part_id.to_string(), + tool_name, + result: json!({ "queries": queries, "results": results }), + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: None, + })); + } + return; + } + + // ── Server-side tool calls ── + if matches!( + item_type, + "web_search_call" + | "x_search_call" + | "code_interpreter_call" + | "code_execution_call" + | "view_image_call" + | "view_x_video_call" + | "custom_tool_call" + | "mcp_call" + ) { + let item_name = item.get("name").and_then(|v| v.as_str()); + let tool_name = + resolve_tool_name(item_type, item_name, provider_tool_names); + let tool_input = get_tool_input(item_type, &item); + let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); + + let should_emit = if item_type == "custom_tool_call" { + event_type == "response.output_item.done" + } else { + !seen_tool_calls.contains(part_id) + }; + + if should_emit && !seen_tool_calls.contains(part_id) { + seen_tool_calls.insert(part_id.to_string()); + controller.tool_input_start(part_id.to_string(), tool_name.clone()); + controller.tool_input_delta(part_id.to_string(), tool_input.clone()); + controller.tool_input_end(part_id.to_string()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: part_id.to_string(), + tool_name: tool_name.clone(), + input: tool_input, + provider_executed: Some(true), + dynamic: None, + provider_metadata: None, + })); + } + + if event_type == "response.output_item.done" { + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: part_id.to_string(), + tool_name, + result: if item_type == "web_search_call" { + map_web_search_action( + item.get("action").unwrap_or(&Value::Null), + ) + } else { + json!({}) + }, + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: None, + })); + } + return; + } + + // ── message item ── + if item_type == "message" { + let content_parts = item + .get("content") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + for cp in &content_parts { + if let Some(text) = cp.get("text").and_then(|v| v.as_str()) + && !text.is_empty() + { + let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("text-{part_id}"); + if !content_blocks.contains(&block_id) { + content_blocks.push(block_id.clone()); + controller.text_start(block_id.clone()); + controller.text_delta(block_id, text.to_string()); + } + } + if let Some(annotations) = + cp.get("annotations").and_then(|v| v.as_array()) + { + for ann in annotations { + if ann.get("type").and_then(|v| v.as_str()) + == Some("url_citation") + { + let url = + ann.get("url").and_then(|v| v.as_str()).unwrap_or(""); + let title = ann + .get("title") + .and_then(|v| v.as_str()) + .unwrap_or(url) + .to_string(); + controller.enqueue(StreamPart::Source(Source::Url { + id: generate_source_id(), + url: url.to_string(), + title: Some(title), + provider_metadata: None, + })); + } + } + } + } + return; + } + + // ── function_call item ── + if item_type == "function_call" { + let call_id = item.get("call_id").and_then(|v| v.as_str()).unwrap_or(""); + let name = item.get("name").and_then(|v| v.as_str()).unwrap_or(""); + + if event_type == "response.output_item.added" { + ongoing_tool_calls + .insert(output_index, (call_id.to_string(), name.to_string())); + controller.tool_input_start(call_id.to_string(), name.to_string()); + } else { + // output_item.done + *has_function_call = true; + ongoing_tool_calls.remove(&output_index); + let arguments = + item.get("arguments").and_then(|v| v.as_str()).unwrap_or(""); + controller.tool_input_end(call_id.to_string()); + let input = arguments.to_string(); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: call_id.to_string(), + tool_name: name.to_string(), + input, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); + } + } + } + + // All other event types (web_search_call.in_progress, etc.) are ignored. + } + Err(error) => { + if error.is_recoverable_stream_error() { + if include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: Value::Null, + }); + } + controller.enqueue(StreamPart::Error { error }); + } else { + controller.error(error); + } + } + } + } + + /// Close the open text blocks and emit `Finish`: at the end of the + /// body, or right after an error event. + fn finish(&mut self, controller: &mut Controller) { + // Close any remaining open text blocks. + for block_id in std::mem::take(&mut self.content_blocks) { + controller.text_end(block_id); + } + + // Final part: Finish. + let provider_meta = response_provider_metadata( + self.cost_in_usd_ticks, + self.service_tier.as_deref(), + self.prompt_cache_key.as_deref(), + self.safety_identifier.as_deref(), + ); + + controller.enqueue(StreamPart::Finish { + finish_reason: self.final_finish_reason.clone().unwrap_or(FINISH_OTHER), + usage: self.final_usage.take().unwrap_or_else(zero_usage), + provider_metadata: provider_meta, + }); + } +} + +// ── Request body builder ───────────────────────────────────────────────────── + +#[derive(Debug, Clone)] +pub struct ResponsesRequestBodyResult { + pub body: Value, + pub warnings: Vec, + pub provider_tool_names: std::collections::HashMap, +} + +/// Build the xAI Responses request body, returning warnings and provider tool +/// names. +/// +/// # Errors +/// +/// Propagates conversion errors from `convert_to_xai_responses_input`. +pub fn build_responses_request_body( + model_id: &str, + options: &CallOptions, + stream: bool, +) -> Result { + let mut warnings: Vec = Vec::new(); + let xai_opts = &options.provider_options; + validate_response_options(xai_opts)?; + + if options.frequency_penalty.is_some() { + warnings.push(Warning::Unsupported { + feature: "frequencyPenalty".to_string(), + details: None, + }); + } + + if options.presence_penalty.is_some() { + warnings.push(Warning::Unsupported { + feature: "presencePenalty".to_string(), + details: None, + }); + } + + if options.stop_sequences.is_some() { + warnings.push(Warning::Unsupported { + feature: "stopSequences".to_string(), + details: None, + }); + } + + let (input, input_warnings) = convert_to_xai_responses_input(&options.prompt)?; + warnings.extend(input_warnings); + + let prepared = prepare_responses_tools(&options.tools, options.tool_choice.as_ref())?; + for tw in &prepared.tool_warnings { + warnings.push(tw.clone()); + } + + // Reasoning effort. + let mut reasoning_effort: Option = xai_option(xai_opts, "reasoningEffort").map(|v| { + v.as_str() + .map(std::string::ToString::to_string) + .unwrap_or(v.to_string()) + }); + + if reasoning_effort.is_none() && options.reasoning.is_some_and(ReasoningEffort::is_custom) { + let reasoning = options.reasoning.unwrap(); + if !supports_reasoning_effort(model_id) { + warnings.push(Warning::Unsupported { + feature: "reasoning".to_string(), + details: Some(format!( + "reasoning \"{reasoning}\" is not supported by this model." + )), + }); + } else if reasoning == ReasoningEffort::None { + reasoning_effort = Some("none".to_string()); + } else { + reasoning_effort = Some(match reasoning { + ReasoningEffort::Minimal => "low".to_string(), + ReasoningEffort::Low => "low".to_string(), + ReasoningEffort::Medium => "medium".to_string(), + ReasoningEffort::High => "high".to_string(), + ReasoningEffort::Xhigh if model_id == "grok-4.6" => "xhigh".to_string(), + ReasoningEffort::Xhigh => "high".to_string(), + _ => reasoning.to_string(), + }); + if reasoning_effort.as_deref() != Some(reasoning.to_string().as_str()) { + warnings.push(Warning::Compatibility { + feature: "reasoning".to_string(), + details: Some(format!( + "reasoning \"{reasoning}\" is not directly supported by this model. mapped to effort \"{}\".", + reasoning_effort.as_deref().unwrap_or_default() + )), + }); + } + } + } + + // Store option. + let store = xai_option(xai_opts, "store").and_then(|v| v.as_bool()); + + // Include array. + let mut include: Option> = xai_option(xai_opts, "include").and_then(|v| { + v.as_array().map(|arr| { + arr.iter() + .filter_map(|v| v.as_str().map(std::string::ToString::to_string)) + .collect() + }) + }); + + if store == Some(false) { + include = match include { + None => Some(vec!["reasoning.encrypted_content".to_string()]), + Some(mut v) => { + v.push("reasoning.encrypted_content".to_string()); + Some(v) + } + }; + } + + let previous_response_id = xai_option(xai_opts, "previousResponseId") + .and_then(|v| v.as_str().map(std::string::ToString::to_string)); + + // Logprobs. + let logprobs = xai_option(xai_opts, "logprobs") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let top_logprobs = xai_option(xai_opts, "topLogprobs"); + + // Build body. + let mut body = json!({ "model": model_id, "input": input }); + + if stream { + body["stream"] = json!(true); + } + if let Some(max_tokens) = options.max_output_tokens { + body["max_output_tokens"] = json!(max_tokens); + } + if let Some(temp) = options.temperature { + body["temperature"] = json!(temp); + } + if let Some(tp) = options.top_p { + body["top_p"] = json!(tp); + } + if let Some(tk) = options.top_k { + body["top_k"] = json!(tk); + } + if let Some(seed) = options.seed { + body["seed"] = json!(seed); + } + if logprobs || top_logprobs.is_some() { + body["logprobs"] = json!(true); + } + if let Some(tlp) = top_logprobs { + body["top_logprobs"] = tlp; + } + + if let Some(effort) = &reasoning_effort { + body["reasoning"] = json!({ "effort": effort }); + } + for (option, parameter) in [ + ("minP", "min_p"), + ("maxTurns", "max_turns"), + ("parallelToolCalls", "parallel_tool_calls"), + ("promptCacheKey", "prompt_cache_key"), + ("safetyIdentifier", "safety_identifier"), + ("serviceTier", "service_tier"), + ("user", "user"), + ] { + if let Some(value) = xai_option(xai_opts, option) { + body[parameter] = value; + } + } + + // Store. + if store == Some(false) { + body["store"] = json!(false); + } + + // Include. + if let Some(inc) = &include { + body["include"] = json!(inc); + } + + // Previous response ID. + if let Some(prev_id) = &previous_response_id { + body["previous_response_id"] = json!(prev_id); + } + + // Response format. + if let Some(ref rf) = options.response_format + && let ResponseFormat::Json { + schema, + name, + description, + } = rf + { + let format = if schema.is_some() { + let mut fmt = json!({ "type": "json_schema", "strict": true }); + fmt["name"] = json!(name.clone().unwrap_or_else(|| "response".to_string())); + if let Some(desc) = description { + fmt["description"] = json!(desc); + } + if let Some(s) = schema { + fmt["schema"] = s.clone(); + } + fmt + } else { + json!({ "type": "json_object" }) + }; + body["text"] = json!({ "format": format }); + } + + // Tools. + if let Some(tools) = &prepared.tools { + body["tools"] = json!(tools); + } + if let Some(tc) = prepared.tool_choice { + body["tool_choice"] = tc; + } + + Ok(ResponsesRequestBodyResult { + body, + warnings, + provider_tool_names: prepared.provider_tool_names, + }) +} + +// ── Tool name resolution helpers ───────────────────────────────────────────── + +/// Sub-tool names that map to web_search. +pub const WEB_SEARCH_SUB_TOOLS: &[&str] = + &["web_search", "web_search_with_snippets", "browse_page"]; +/// Sub-tool names that map to x_search. +pub const X_SEARCH_SUB_TOOLS: &[&str] = &[ + "x_user_search", + "x_keyword_search", + "x_semantic_search", + "x_thread_fetch", +]; + +/// Resolve the tool name for a server-side tool call, using the user-provided +/// provider tool name if available, falling back to a default. +#[must_use] +pub fn resolve_tool_name( + part_type: &str, + part_name: Option<&str>, + provider_tool_names: &std::collections::HashMap, +) -> String { + let name = part_name.unwrap_or(""); + let default_web = provider_tool_names + .get("xai.web_search") + .cloned() + .unwrap_or_else(|| "web_search".to_string()); + let default_x = provider_tool_names + .get("xai.x_search") + .cloned() + .unwrap_or_else(|| "x_search".to_string()); + let default_code = provider_tool_names + .get("xai.code_execution") + .cloned() + .unwrap_or_else(|| "code_execution".to_string()); + let default_mcp = provider_tool_names + .get("xai.mcp") + .cloned() + .unwrap_or_else(|| part_name.unwrap_or("mcp").to_string()); + let default_file_search = provider_tool_names + .get("xai.file_search") + .cloned() + .unwrap_or_else(|| "file_search".to_string()); + + if WEB_SEARCH_SUB_TOOLS.contains(&name) || part_type == "web_search_call" { + default_web + } else if X_SEARCH_SUB_TOOLS.contains(&name) || part_type == "x_search_call" { + default_x + } else if name == "code_execution" + || part_type == "code_interpreter_call" + || part_type == "code_execution_call" + { + default_code + } else if part_type == "mcp_call" { + default_mcp + } else if part_type == "file_search_call" { + default_file_search + } else if part_type == "image_generation_call" { + provider_tool_names + .get("xai.image_generation") + .cloned() + .unwrap_or_else(|| "image_generation".to_string()) + } else if !name.is_empty() { + name.to_string() + } else { + String::new() + } +} + +/// Get the tool input for a server-side tool call. +#[must_use] +pub fn get_tool_input(part_type: &str, part: &Value) -> String { + match part_type { + "custom_tool_call" => part + .get("input") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + "mcp_call" => part + .get("arguments") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + _ => part + .get("arguments") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + } +} diff --git a/aimux-providers/src/xai/responses/responses_language_model_options.rs b/aimux-providers/src/xai/responses/responses_language_model_options.rs new file mode 100644 index 000000000..ef4cb51a1 --- /dev/null +++ b/aimux-providers/src/xai/responses/responses_language_model_options.rs @@ -0,0 +1,86 @@ +//! The `providerOptions.xai` checks of `xai-responses-language-model-options.ts`. + +use aimux_core::error::AiMuxError; +use aimux_core::shared::SharedProviderOptions; +use serde_json::Value; + +pub(super) fn xai_option(options: &Option, key: &str) -> Option { + options + .as_ref() + .and_then(|m| crate::xai::options::xai_options(Some(m))) + .and_then(|o| o.get(key)) + .cloned() +} + +pub(super) fn validate_value( + key: &str, + value: &Value, + valid: impl FnOnce(&Value) -> bool, +) -> Result<(), AiMuxError> { + if valid(value) { + Ok(()) + } else { + Err(AiMuxError::InvalidArgument(format!( + "Invalid xai option {key}: {value}" + ))) + } +} + +pub(super) fn string_array(value: &Value, maximum: usize) -> bool { + value + .as_array() + .is_some_and(|items| items.len() <= maximum && items.iter().all(Value::is_string)) +} + +pub(super) fn validate_response_options( + options: &Option, +) -> Result<(), AiMuxError> { + for (key, value) in options + .as_ref() + .and_then(|options| crate::xai::options::xai_options(Some(options))) + .into_iter() + .flatten() + { + validate_value(key, value, |value| match key.as_str() { + "reasoningEffort" => value + .as_str() + .is_some_and(|v| ["none", "low", "medium", "high", "xhigh"].contains(&v)), + "reasoningSummary" => value + .as_str() + .is_some_and(|v| ["auto", "concise", "detailed"].contains(&v)), + "serviceTier" => value + .as_str() + .is_some_and(|v| ["default", "priority"].contains(&v)), + "topLogprobs" => value + .as_f64() + .is_some_and(|v| v.fract() == 0.0 && (0.0..=8.0).contains(&v)), + "minP" => value.as_f64().is_some_and(|v| (0.0..=1.0).contains(&v)), + "maxTurns" => value + .as_f64() + .is_some_and(|v| v.fract() == 0.0 && v.abs() <= 9_007_199_254_740_991.0), + "logprobs" | "store" | "parallelToolCalls" => value.is_boolean(), + "promptCacheKey" | "safetyIdentifier" | "user" | "previousResponseId" => { + value.is_string() + } + "include" => { + value.is_null() + || value.as_array().is_some_and(|items| { + items.iter().all(|item| { + item.as_str().is_some_and(|v| { + [ + "file_search_call.results", + "web_search_call.action.sources", + "code_interpreter_call.outputs", + "reasoning.encrypted_content", + "no_inline_citations", + ] + .contains(&v) + }) + }) + }) + } + _ => true, + })?; + } + Ok(()) +} diff --git a/aimux-providers/src/xai/responses/responses_prepare_tools.rs b/aimux-providers/src/xai/responses/responses_prepare_tools.rs new file mode 100644 index 000000000..09ddb6597 --- /dev/null +++ b/aimux-providers/src/xai/responses/responses_prepare_tools.rs @@ -0,0 +1,257 @@ +//! Tools and tool choice for the Responses API (`xai-responses-prepare-tools.ts`). + +use aimux_core::error::AiMuxError; +use aimux_core::options::LanguageModelV4ToolChoice; +use aimux_core::tool::Tool; +use aimux_core::types::Warning; +use serde_json::{Value, json}; + +use super::responses_language_model_options::{string_array, validate_value}; + +#[derive(Debug, Clone)] +pub struct PreparedResponsesTools { + pub tools: Option>, + pub tool_choice: Option, + pub tool_warnings: Vec, + /// Names of provider-executed tools, keyed by tool id (e.g. "xai.web_search"). + pub provider_tool_names: std::collections::HashMap, +} + +/// # Errors +/// Rejects provider-tool arguments that do not match the upstream schemas. +pub fn prepare_responses_tools( + tools: &Option>, + tool_choice: Option<&LanguageModelV4ToolChoice>, +) -> Result { + let non_empty = tools.as_ref().filter(|t| !t.is_empty()); + for tool in non_empty.into_iter().flatten() { + if let Tool::Provider(tool) = tool { + validate_provider_tool(tool)?; + } + } + let mut tool_warnings: Vec = Vec::new(); + let mut provider_tool_names: std::collections::HashMap = + std::collections::HashMap::new(); + + let tools_opt = match non_empty { + None => None, + Some(tools) => { + let xai_tools: Vec = tools + .iter() + .filter_map(|t| match t { + Tool::Function(ft) => { + let mut func = json!({ + "type": "function", + "name": ft.name, + "parameters": ft.input_schema, + }); + if let Some(ref desc) = ft.description { + func["description"] = json!(desc); + } + if let Some(strict) = ft.strict { + func["strict"] = json!(strict); + } + Some(func) + } + Tool::Provider(pt) => { + provider_tool_names.insert(pt.id.to_string(), pt.name.clone()); + prepare_provider_tool(pt, &mut tool_warnings) + } + }) + .collect(); + Some(xai_tools) + } + }; + + let tool_choice_opt = match (&tools_opt, tool_choice) { + (None, _) => None, + (Some(_), None) => None, + (Some(_), Some(tc)) => { + // Build a map of tool name -> is_provider_tool from the original tools. + let non_empty_tools = non_empty.unwrap(); + match tc { + LanguageModelV4ToolChoice::Auto => Some(json!("auto")), + LanguageModelV4ToolChoice::None => Some(json!("none")), + LanguageModelV4ToolChoice::Required => Some(json!("required")), + LanguageModelV4ToolChoice::Tool { tool_name } => { + let selected = non_empty_tools.iter().rev().find(|t| match t { + Tool::Function(ft) => ft.name == *tool_name, + Tool::Provider(pt) => pt.name == *tool_name, + }); + + match selected { + None => None, + Some(Tool::Function(_)) => { + Some(json!({ "type": "function", "name": tool_name })) + } + Some(Tool::Provider(pt)) => { + // Server-side tools cannot be forced via toolChoice. + tool_warnings.push(Warning::Unsupported { + feature: format!("toolChoice for server-side tool \"{}\"", pt.name), + details: None, + }); + None + } + } + } + } + } + }; + + Ok(PreparedResponsesTools { + tools: tools_opt, + tool_choice: tool_choice_opt, + tool_warnings, + provider_tool_names, + }) +} + +fn prepare_provider_tool( + pt: &aimux_core::tool::ProviderTool, + warnings: &mut Vec, +) -> Option { + let args = &pt.args; + match pt.id.as_str() { + "xai.web_search" => { + let mut tool = json!({ "type": "web_search" }); + if let Some(v) = args.get("allowedDomains") { + tool["allowed_domains"] = v.clone(); + } + if let Some(v) = args.get("excludedDomains") { + tool["excluded_domains"] = v.clone(); + } + if let Some(v) = args.get("enableImageSearch") { + tool["enable_image_search"] = v.clone(); + } + if let Some(v) = args.get("enableImageUnderstanding") { + tool["enable_image_understanding"] = v.clone(); + } + Some(tool) + } + "xai.x_search" => { + let mut tool = json!({ "type": "x_search" }); + if let Some(v) = args.get("allowedXHandles") { + tool["allowed_x_handles"] = v.clone(); + } + if let Some(v) = args.get("excludedXHandles") { + tool["excluded_x_handles"] = v.clone(); + } + if let Some(v) = args.get("fromDate") { + tool["from_date"] = v.clone(); + } + if let Some(v) = args.get("toDate") { + tool["to_date"] = v.clone(); + } + if let Some(v) = args.get("enableImageUnderstanding") { + tool["enable_image_understanding"] = v.clone(); + } + if let Some(v) = args.get("enableVideoUnderstanding") { + tool["enable_video_understanding"] = v.clone(); + } + Some(tool) + } + "xai.code_execution" => Some(json!({ "type": "code_interpreter" })), + "xai.view_image" => Some(json!({ "type": "view_image" })), + "xai.view_x_video" => Some(json!({ "type": "view_x_video" })), + "xai.image_generation" => { + let mut tool = json!({ "type": "image_generation" }); + if let Some(action) = args.get("action") { + tool["action"] = action.clone(); + } + Some(tool) + } + "xai.file_search" => { + let mut tool = json!({ "type": "file_search" }); + if let Some(v) = args.get("vectorStoreIds") { + tool["vector_store_ids"] = v.clone(); + } + if let Some(v) = args.get("maxNumResults") { + tool["max_num_results"] = v.clone(); + } + Some(tool) + } + "xai.mcp" => { + let mut tool = json!({ "type": "mcp" }); + if let Some(v) = args.get("serverUrl") { + tool["server_url"] = v.clone(); + } + if let Some(v) = args.get("serverLabel") { + tool["server_label"] = v.clone(); + } + if let Some(v) = args.get("serverDescription") { + tool["server_description"] = v.clone(); + } + if let Some(v) = args.get("allowedTools") { + tool["allowed_tools"] = v.clone(); + } + if let Some(v) = args.get("headers") { + tool["headers"] = v.clone(); + } + if let Some(v) = args.get("authorization") { + tool["authorization"] = v.clone(); + } + Some(tool) + } + _ => { + warnings.push(Warning::Unsupported { + feature: format!("provider-defined tool {}", pt.name), + details: None, + }); + None + } + } +} + +fn validate_provider_tool(tool: &aimux_core::tool::ProviderTool) -> Result<(), AiMuxError> { + let id = tool.id.as_str(); + if ![ + "xai.web_search", + "xai.x_search", + "xai.image_generation", + "xai.file_search", + "xai.mcp", + ] + .contains(&id) + { + return Ok(()); + } + let args = &tool.args; + for (key, value) in args { + validate_value(key, value, |value| match (id, key.as_str()) { + ("xai.web_search", "allowedDomains" | "excludedDomains") => string_array(value, 5), + ("xai.x_search", "allowedXHandles" | "excludedXHandles") => string_array(value, 10), + ("xai.web_search", "enableImageSearch" | "enableImageUnderstanding") + | ("xai.x_search", "enableImageUnderstanding" | "enableVideoUnderstanding") => { + value.is_boolean() + } + ("xai.x_search", "fromDate" | "toDate") => value.is_string(), + ("xai.image_generation", "action") => value + .as_str() + .is_some_and(|v| ["auto", "generate", "edit"].contains(&v)), + ("xai.file_search", "vectorStoreIds") | ("xai.mcp", "allowedTools") => { + string_array(value, usize::MAX) + } + ("xai.file_search", "maxNumResults") => value.is_number(), + ("xai.mcp", "serverUrl" | "serverLabel" | "serverDescription" | "authorization") => { + value.is_string() + } + ("xai.mcp", "headers") => value + .as_object() + .is_some_and(|headers| headers.values().all(Value::is_string)), + _ => true, + })?; + } + for key in match id { + "xai.file_search" => &["vectorStoreIds"][..], + "xai.mcp" => &["serverUrl"][..], + _ => &[], + } { + if !args.contains_key(*key) { + return Err(AiMuxError::InvalidArgument(format!( + "Missing {key} for {}", + tool.name + ))); + } + } + Ok(()) +} diff --git a/aimux-providers/src/xai/speech.rs b/aimux-providers/src/xai/speech_model.rs similarity index 98% rename from aimux-providers/src/xai/speech.rs rename to aimux-providers/src/xai/speech_model.rs index 6e73a308d..1cb0807a3 100644 --- a/aimux-providers/src/xai/speech.rs +++ b/aimux-providers/src/xai/speech_model.rs @@ -166,7 +166,7 @@ impl SpeechModel for XaiSpeechModel { exchange.request(url, options), body.clone(), aimux_provider_utils::create_json_response_handler::(), - super::xai_failed_response_handler(), + super::error::xai_failed_response_handler(), ) .await?; let envelope = resp.value; @@ -185,7 +185,7 @@ impl SpeechModel for XaiSpeechModel { exchange.request(url, options), body.clone(), aimux_provider_utils::create_binary_response_handler(), - super::xai_failed_response_handler(), + super::error::xai_failed_response_handler(), ) .await?; (resp.value.to_vec(), resp.response_headers, None, None) diff --git a/aimux-providers/src/xai/supports_reasoning_effort.rs b/aimux-providers/src/xai/supports_reasoning_effort.rs new file mode 100644 index 000000000..9002cac00 --- /dev/null +++ b/aimux-providers/src/xai/supports_reasoning_effort.rs @@ -0,0 +1,25 @@ +//! `supports-reasoning-effort.ts`: which Grok models accept `reasoning_effort`. + +/// Models that reject the `reasoning_effort` parameter. +/// Matches `^grok-4\.20(-\d{4})?-(non-)?reasoning$`. +fn is_model_without_reasoning_effort(model_id: &str) -> bool { + let rest = match model_id.strip_prefix("grok-4.20") { + Some(r) => r, + None => return false, + }; + let rest = if rest.len() >= 5 + && rest.starts_with('-') + && rest[1..5].chars().all(|c| c.is_ascii_digit()) + { + &rest[5..] + } else { + rest + }; + rest == "-reasoning" || rest == "-non-reasoning" +} + +/// Whether the model accepts the `reasoning_effort` parameter. +#[must_use] +pub fn supports_reasoning_effort(model_id: &str) -> bool { + !is_model_without_reasoning_effort(model_id) +} diff --git a/aimux-providers/src/xai/tests/error.rs b/aimux-providers/src/xai/tests/error.rs index 1af3d8306..cf17a61ea 100644 --- a/aimux-providers/src/xai/tests/error.rs +++ b/aimux-providers/src/xai/tests/error.rs @@ -7,7 +7,7 @@ use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue}; use reqwest::{StatusCode, Url}; use serde_json::{Value, json}; -use crate::xai::xai_failed_response_handler; +use crate::xai::error::xai_failed_response_handler; /// The upstream `makeResponse` helper: a 400 with a JSON body. fn make_response(url: &str, body: &Value) -> FetchResponse { diff --git a/aimux-providers/src/xai/tests/responses/convert_responses_usage.rs b/aimux-providers/src/xai/tests/responses/convert_responses_usage.rs index 6697c3606..766a83823 100644 --- a/aimux-providers/src/xai/tests/responses/convert_responses_usage.rs +++ b/aimux-providers/src/xai/tests/responses/convert_responses_usage.rs @@ -6,8 +6,8 @@ use aimux_core::types::Usage; use serde_json::{Value, json}; -use crate::xai::responses::convert::convert_xai_responses_usage; -use crate::xai::responses::types::XaiResponsesUsage; +use crate::xai::responses::XaiResponsesUsage; +use crate::xai::responses::convert_xai_responses_usage; /// `convertXaiResponsesUsage(usage)` for a usage object as the API returns it. fn convert(usage: Value) -> Usage { diff --git a/aimux-providers/src/xai/tests/responses/convert_to_responses_input.rs b/aimux-providers/src/xai/tests/responses/convert_to_responses_input.rs index 85d628940..266412704 100644 --- a/aimux-providers/src/xai/tests/responses/convert_to_responses_input.rs +++ b/aimux-providers/src/xai/tests/responses/convert_to_responses_input.rs @@ -19,7 +19,7 @@ use aimux_core::types::Warning; use serde_json::{Value, json}; use crate::xai::options::NAMESPACE; -use crate::xai::responses::convert::convert_to_xai_responses_input; +use crate::xai::responses::convert_to_xai_responses_input; /// The prompt in its JSON form, as the Rust prompt types deserialize it. fn parse_prompt(prompt: Value) -> LanguageModelPrompt { diff --git a/aimux-providers/src/xai/tests/responses/responses_language_model.rs b/aimux-providers/src/xai/tests/responses/responses_language_model.rs index 558f427b9..36bc877a8 100644 --- a/aimux-providers/src/xai/tests/responses/responses_language_model.rs +++ b/aimux-providers/src/xai/tests/responses/responses_language_model.rs @@ -13,20 +13,25 @@ //! - `undefined` properties are absent from the expected JSON. //! - `toMatchSnapshot()` values are `snapshots/.json`, converted //! from the vitest `.snap` file (their text matches the fixtures). +//! +//! The last top-level cases are the older Responses subset (formerly +//! `tests/xai_upstream_e2e_test.rs`), on the same [`MockFetch`]. use std::sync::Arc; +use aimux_core::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{ AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, }; use aimux_core::options::{CallOptions, ResponseFormat}; -use aimux_core::result::{GenerateResult, StreamResult}; +use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions, provider_namespace}; +use aimux_core::stream_part::StreamPart; use aimux_core::tool::{FunctionTool, ProviderTool, Tool}; -use aimux_core::types::ReasoningEffort; -use futures::StreamExt; +use aimux_core::types::{FinishReasonUnified, ReasoningEffort}; +use futures::{StreamExt, TryStreamExt}; use serde::Serialize; use serde_json::{Value, json}; @@ -352,6 +357,59 @@ fn empty_response(usage: Value) -> Value { // ───────────────────────────────────────────────────────────────────────────── +// ── The harness of the older Responses subset ───────────────────────────────── + +fn model(fetch: &std::sync::Arc) -> XaiResponsesModel { + create_xai(XAIProviderSettings { + base_url: Some("https://example.test/v1/".into()), + api_key: Some("test-key".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() + .responses("grok-4-fast-non-reasoning") +} + +fn options() -> CallOptions { + CallOptions { + prompt: vec![LanguageModelMessage::user_text("hello")], + ..Default::default() + } +} + +fn assert_request(fetch: &MockFetch, stream: bool, tools: bool) { + let requests = fetch.seen(); + assert_eq!(requests.len(), 1); + let request = &requests[0]; + assert_eq!(request.method, "POST"); + assert_eq!(request.url, "https://example.test/v1/responses"); + assert_eq!(request.headers["authorization"], "Bearer test-key"); + let mut expected = json!({ + "model": "grok-4-fast-non-reasoning", + "input": [{"role": "user", "content": [{"type": "input_text", "text": "hello"}]}] + }); + if stream { + expected["stream"] = json!(true); + } + if tools { + expected["tools"] = json!([{"type": "web_search"}]); + } + assert_eq!(request.json_body(), expected); +} + +fn sse(events: &[Value]) -> Canned { + Canned { + status: 200, + headers: vec![("content-type".into(), "text/event-stream".into())], + body: (events + .iter() + .map(|event| format!("data: {event}\n\n")) + .collect::() + + "data: [DONE]\n\n") + .into_bytes(), + } +} + mod xai_responses_language_model { use super::*; @@ -5198,3 +5256,220 @@ mod xai_responses_language_model { } } } + +/// partial: `should extract usage correctly`; full port in +/// src/xai/tests/responses/responses_language_model.rs. +#[tokio::test] +async fn generate_converts_usage() { + // The port used an unrelated search fixture for this case; restore the TS response. + let usage = json!({ + "input_tokens": 345, "output_tokens": 538, "total_tokens": 883, + "output_tokens_details": {"reasoning_tokens": 123} + }); + let body = json!({ + "id": "resp_123", "object": "response", "status": "completed", + "model": "grok-4-fast-non-reasoning", "output": [], "usage": usage + }); + let fetch = MockFetch::new(vec![Canned::json(&body)]); + let result = model(&fetch).do_generate(&options()).await.unwrap(); + assert_request(&fetch, false, false); + assert!(result.content.is_empty()); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(result.finish_reason.raw.as_deref(), Some("completed")); + assert_eq!( + serde_json::to_value(result.usage).unwrap(), + json!({ + "inputTokens": {"total": 345, "noCache": 345, "cacheRead": 0}, + "outputTokens": {"total": 538, "text": 415, "reasoning": 123}, + "raw": usage + }) + ); + assert!(result.provider_metadata.is_none()); + assert_eq!(result.response.unwrap().id.as_deref(), Some("resp_123")); +} + +/// partial: `should expose cost_in_usd_ticks in finish providerMetadata`; full port in +/// src/xai/tests/responses/responses_language_model.rs. +#[tokio::test] +async fn stream_reports_cost_and_text() { + let events = vec![ + json!({"type": "response.created", "response": {"id": "resp_123", "object": "response", "model": "grok-4-fast-non-reasoning", "output": []}}), + json!({"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "Hello"}), + json!({"type": "response.completed", "response": {"id": "resp_123", "object": "response", "model": "grok-4-fast-non-reasoning", "status": "completed", "output": [{"type": "message", "id": "msg_001", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": "Hello"}]}], "usage": {"input_tokens": 10, "output_tokens": 5, "cost_in_usd_ticks": 113500}}}), + ]; + let fetch = MockFetch::new(vec![sse(&events)]); + let parts: Vec<_> = model(&fetch) + .do_stream(&options()) + .await + .unwrap() + .stream + .try_collect() + .await + .unwrap(); + assert_request(&fetch, true, false); + let text: String = parts + .iter() + .filter_map(|part| match part { + StreamPart::TextDelta { delta, .. } => Some(delta.as_str()), + _ => None, + }) + .collect(); + assert_eq!(text, "Hello"); + let finishes: Vec<_> = parts + .iter() + .filter_map(|part| match part { + StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } => Some((finish_reason, usage, provider_metadata)), + _ => None, + }) + .collect(); + assert_eq!(finishes.len(), 1); + let (reason, usage, metadata) = finishes[0]; + assert_eq!(reason.unified, FinishReasonUnified::Stop); + assert_eq!(reason.raw.as_deref(), Some("completed")); + assert_eq!(usage.input_tokens.total, Some(10)); + assert_eq!(usage.output_tokens.total, Some(5)); + assert_eq!( + metadata.as_ref().unwrap(), + &provider_namespace("xai", json!({"costInUsdTicks": 113500})).unwrap() + ); + assert!( + !parts + .iter() + .any(|part| matches!(part, StreamPart::Error { .. })) + ); +} + +/// partial: `should forward web_search_call action and sources to tool-result`; full port in +/// src/xai/tests/responses/responses_language_model.rs. +#[tokio::test] +async fn tool_results_preserve_actions_and_sources() { + let body = json!({ + "id": "resp_123", + "object": "response", + "status": "completed", + "model": "grok-4-fast-non-reasoning", + "output": [ + {"type": "web_search_call", "id": "ws_action_1", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "search", "query": "latest AI news", "sources": [{"type": "url", "url": "https://example.com/a"}, {"type": "url", "url": "https://example.com/b"}]}}, + {"type": "web_search_call", "id": "ws_action_2", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "open_page", "url": "https://example.com/a"}}, + {"type": "web_search_call", "id": "ws_action_3", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "find_in_page", "url": "https://example.com/a", "pattern": "climate"}}, + {"type": "web_search_call", "id": "ws_action_4", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "search", "query": null, "sources": null}}, + {"type": "web_search_call", "id": "ws_action_5", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "open_page", "url": null}}, + ], + "usage": {"input_tokens": 10, "output_tokens": 5}, + }); + let fetch = MockFetch::new(vec![Canned::json(&body)]); + let mut options = options(); + options.tools = Some(vec![Tool::Provider(ProviderTool { + id: "xai.web_search".parse().unwrap(), + name: "web_search".into(), + args: serde_json::Map::new(), + supports_deferred_results: None, + })]); + let result = model(&fetch).do_generate(&options).await.unwrap(); + assert_request(&fetch, false, true); + let expected = vec![ + json!({"action": {"type": "search", "query": "latest AI news"}, "sources": [{"type": "url", "url": "https://example.com/a"}, {"type": "url", "url": "https://example.com/b"}]}), + json!({"action": {"type": "openPage", "url": "https://example.com/a"}}), + json!({"action": {"type": "findInPage", "url": "https://example.com/a", "pattern": "climate"}}), + json!({"action": {"type": "search"}}), + json!({"action": {"type": "openPage", "url": null}}), + ]; + assert_eq!(result.content.len(), expected.len() * 2); + for (index, (pair, expected_result)) in result + .content + .as_chunks::<2>() + .0 + .iter() + .zip(expected) + .enumerate() + { + let id = format!("ws_action_{}", index + 1); + let GenerateContent::ToolCall(call) = &pair[0] else { + panic!("missing tool call") + }; + assert_eq!(call.tool_call_id, id); + assert_eq!(call.tool_name, "web_search"); + assert_eq!(call.input, ""); + assert_eq!(call.provider_executed, Some(true)); + let GenerateContent::ToolResult(output) = &pair[1] else { + panic!("missing tool result") + }; + assert_eq!(output.tool_call_id, id); + assert_eq!(output.tool_name, "web_search"); + assert_eq!(output.result, expected_result); + assert_eq!(output.is_error, None); + } + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); +} + +/// partial: `should set finish reason to error`; full port in +/// src/xai/tests/responses/responses_language_model.rs. +#[tokio::test] +async fn failed_response_emits_error_and_finish() { + let events = vec![ + json!({"type": "response.created", "response": {"id": "resp_123", "object": "response", "model": "grok-4-fast-non-reasoning", "output": []}}), + json!({"type": "response.failed", "response": {"error": {"code": "server_error", "message": "Internal server error"}, "usage": {"input_tokens": 50, "output_tokens": 0}}}), + ]; + let fetch = MockFetch::new(vec![sse(&events)]); + let parts: Vec<_> = model(&fetch) + .do_stream(&options()) + .await + .unwrap() + .stream + .try_collect() + .await + .unwrap(); + assert_request(&fetch, true, false); + let errors: Vec<_> = parts + .iter() + .filter_map(|part| match part { + StreamPart::Error { error } => Some(error), + _ => None, + }) + .collect(); + assert_eq!(errors.len(), 1); + let AiMuxError::ApiCall(error) = errors[0] else { + panic!("expected API error") + }; + assert_eq!(error.message, "Internal server error"); + assert_eq!(error.provider_code.as_deref(), Some("server_error")); + assert_eq!(error.status_code, Some(500)); + assert!(error.is_retryable); + assert_eq!(error.data.as_ref(), Some(&events[1])); + let finishes: Vec<_> = parts + .iter() + .filter_map(|part| match part { + StreamPart::Finish { + finish_reason, + usage, + .. + } => Some((finish_reason, usage)), + _ => None, + }) + .collect(); + assert_eq!(finishes.len(), 1); + let (reason, usage) = finishes[0]; + assert_eq!(reason.unified, FinishReasonUnified::Error); + assert_eq!(reason.raw.as_deref(), Some("error")); + assert_eq!(usage.input_tokens.total, Some(50)); + assert_eq!(usage.output_tokens.total, Some(0)); +} + +/// `xaiResponsesSupportedUrls` (xai-responses-language-model.ts:142-149): +/// image, PDF and text URLs go to the API as they are. +#[tokio::test] +async fn responses_support_http_image_pdf_and_text_urls() { + let fetch = MockFetch::new(vec![]); + let urls = model(&fetch).supported_urls().await; + let mut kinds: Vec<_> = urls.0.keys().map(String::as_str).collect(); + kinds.sort_unstable(); + assert_eq!(kinds, ["application/pdf", "image/*", "text/*"]); + for patterns in urls.0.values() { + assert!(patterns[0].is_match("https://example.com/file")); + assert!(!patterns[0].is_match("ftp://example.com/file")); + } +} diff --git a/aimux-providers/src/xai/tests/responses/responses_prepare_tools.rs b/aimux-providers/src/xai/tests/responses/responses_prepare_tools.rs index b9f2666bb..d51e5fac2 100644 --- a/aimux-providers/src/xai/tests/responses/responses_prepare_tools.rs +++ b/aimux-providers/src/xai/tests/responses/responses_prepare_tools.rs @@ -8,7 +8,7 @@ use aimux_core::tool::{FunctionTool, ProviderTool, Tool}; use aimux_core::types::Warning; use serde_json::{Value, json}; -use crate::xai::responses::convert::{PreparedResponsesTools, prepare_responses_tools}; +use crate::xai::responses::{PreparedResponsesTools, prepare_responses_tools}; /// A provider-defined tool as the caller passes it: `{ type: 'provider', id, name, args }`. fn provider_tool(id: &str, name: &str, args: Value) -> Tool { diff --git a/aimux-providers/src/xai/tests/speech_model.rs b/aimux-providers/src/xai/tests/speech_model.rs index 4fde61eef..9cbce7e2f 100644 --- a/aimux-providers/src/xai/tests/speech_model.rs +++ b/aimux-providers/src/xai/tests/speech_model.rs @@ -9,14 +9,22 @@ //! - The injected `_internal.currentDate` has no counterpart: aimux stamps the //! response with `Timestamp::now()` and has no clock seam. The stamp is //! checked to fall within the call instead of equalling the injected date. +//! +//! The older ports (formerly `tests/xai_speech_test.rs`) follow the wiremock +//! harness helpers: they run against a loopback wiremock `MockServer`. +use std::collections::HashMap; use std::sync::Arc; -use aimux_core::speech_model::{SpeechCallOptions, SpeechModel}; +use aimux_core::error::AiMuxError; +use aimux_core::shared::{SharedProviderOptions, Warning}; +use aimux_core::speech_model::{AudioData, SpeechCallOptions, SpeechModel, SpeechResult}; +use serde_json::{Value, json}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; use crate::test_support::mock_fetch::{Canned, MockFetch}; -use crate::xai::XaiSpeechModel; -use crate::{XAIProviderSettings, create_xai}; +use crate::{XAIProviderSettings, XaiSpeechModel, create_xai}; /// The audio the upstream server answers with (`new Uint8Array([1, 2, 3, 4])`). const AUDIO: [u8; 4] = [1, 2, 3, 4]; @@ -48,6 +56,88 @@ fn audio_answer(headers: &[(&str, &str)]) -> Canned { } } +// ── The wiremock harness of the older ports ────────────────────────────────── + +fn wiremock_speech_model( + server: &MockServer, + headers: Option>>, +) -> XaiSpeechModel { + create_xai(XAIProviderSettings { + api_key: Some("test-key".into()), + base_url: Some(server.uri()), + headers, + ..Default::default() + }) + .expect("valid settings") + .speech() +} + +/// Start a server answering `POST /tts` with `response`. +async fn server_with(response: ResponseTemplate) -> MockServer { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/tts")) + .respond_with(response) + .mount(&server) + .await; + server +} + +fn audio_response(content_type: &str, headers: &[(&str, &str)]) -> ResponseTemplate { + headers.iter().fold( + ResponseTemplate::new(200) + .insert_header("content-type", content_type) + .set_body_bytes(AUDIO.to_vec()), + |response, (name, value)| response.insert_header(*name, *value), + ) +} + +/// The `with_timestamps` JSON envelope; 'AQIDBA==' is `AUDIO` in base64. +fn timestamps_response(headers: &[(&str, &str)]) -> ResponseTemplate { + headers.iter().fold( + ResponseTemplate::new(200).set_body_json(json!({ + "audio": "AQIDBA==", + "content_type": "audio/mpeg", + "duration": 1.19, + "audio_timestamps": { + "graph_chars": ["H", "i"], + "graph_times": [[0.04, 0.06], [0.06, 0.1]], + }, + })), + |response, (name, value)| response.insert_header(*name, *value), + ) +} + +fn options(text: &str) -> SpeechCallOptions { + SpeechCallOptions::new(text) +} + +/// `providerOptions: { xai: {...} }`. +fn with_xai(mut options: SpeechCallOptions, xai: Value) -> SpeechCallOptions { + options.provider_options = Some(SharedProviderOptions::from([( + "xai".to_string(), + xai.as_object().expect("an object").clone(), + )])); + options +} + +/// The JSON body of the only request the server received. +async fn request_body(server: &MockServer) -> Value { + let requests = server.received_requests().await.expect("requests recorded"); + assert_eq!(requests.len(), 1); + serde_json::from_slice(&requests[0].body).unwrap() +} + +fn xai_metadata(result: &SpeechResult) -> Value { + Value::Object(result.provider_metadata.as_ref().expect("metadata")["xai"].clone()) +} + +fn is_unsupported(warnings: &[Warning], expected: &str) -> bool { + warnings + .iter() + .any(|w| matches!(w, Warning::Unsupported { feature, .. } if feature == expected)) +} + mod xai_speech_model { use super::*; @@ -99,4 +189,332 @@ mod do_generate { Some("test-request-id") ); } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should send text with xAI defaults + #[tokio::test] + async fn should_send_text_with_xai_defaults() { + let server = server_with(audio_response("audio/mpeg", &[])).await; + let result = wiremock_speech_model(&server, None) + .do_generate(&options("Hello from the AI SDK!")) + .await + .unwrap(); + + let expected = json!({ + "text": "Hello from the AI SDK!", + "voice_id": "eve", + "language": "auto", + "output_format": { "codec": "mp3" }, + }); + assert_eq!(request_body(&server).await, expected); + assert_eq!(result.request.unwrap().body, Some(expected)); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests[0].method.as_str(), "POST"); + assert_eq!(requests[0].url.path(), "/tts"); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should pass standard speech options + #[tokio::test] + async fn should_pass_standard_speech_options() { + let server = server_with(audio_response("audio/wav", &[])).await; + let mut opts = options("Hello from the AI SDK!"); + opts.voice = Some("ara".into()); + opts.language = Some("en".into()); + opts.output_format = Some("wav".into()); + opts.speed = Some(1.2); + wiremock_speech_model(&server, None) + .do_generate(&opts) + .await + .unwrap(); + + assert_eq!( + request_body(&server).await, + json!({ + "text": "Hello from the AI SDK!", + "voice_id": "ara", + "language": "en", + "output_format": { "codec": "wav" }, + "speed": 1.2, + }) + ); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should accept the %s output format + #[tokio::test] + async fn should_accept_each_output_format() { + for format in ["mp3", "wav", "pcm", "mulaw", "alaw"] { + let server = server_with(audio_response("audio/mpeg", &[])).await; + let mut opts = options("Hello from the AI SDK!"); + opts.output_format = Some(format.into()); + let result = wiremock_speech_model(&server, None) + .do_generate(&opts) + .await + .unwrap(); + + assert_eq!( + request_body(&server).await["output_format"]["codec"], + format + ); + assert!(result.warnings.is_empty(), "{format}"); + } + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should map provider options onto xAI request fields + #[tokio::test] + async fn should_map_provider_options_onto_xai_request_fields() { + let server = server_with(audio_response("audio/mpeg", &[])).await; + let opts = with_xai( + options("Hello from the AI SDK!"), + json!({ + "sampleRate": 44100, + "bitRate": 192000, + "optimizeStreamingLatency": 1, + "textNormalization": true, + }), + ); + wiremock_speech_model(&server, None) + .do_generate(&opts) + .await + .unwrap(); + + let body = request_body(&server).await; + assert_eq!( + body["output_format"], + json!({ "codec": "mp3", "sample_rate": 44100, "bit_rate": 192000 }) + ); + assert_eq!(body["optimize_streaming_latency"], 1); + assert_eq!(body["text_normalization"], true); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should pass withTimestamps and replace provider options + #[tokio::test] + async fn should_pass_with_timestamps_and_replace_provider_options() { + let server = server_with(timestamps_response(&[])).await; + let opts = with_xai( + options("Hello from the AI SDK!"), + json!({ "withTimestamps": true, "replace": { "nginx": "/ˈɛndʒɪn ˈɛks/" } }), + ); + wiremock_speech_model(&server, None) + .do_generate(&opts) + .await + .unwrap(); + + let body = request_body(&server).await; + assert_eq!(body["with_timestamps"], true); + assert_eq!(body["replace"], json!({ "nginx": "/ˈɛndʒɪn ˈɛks/" })); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should decode the with_timestamps envelope and extract provider metadata + #[tokio::test] + async fn should_decode_the_with_timestamps_envelope_and_extract_provider_metadata() { + let trace = "993675dc-8ea6-4f54-b4ad-a59ac2615026"; + let server = server_with(timestamps_response(&[("x-trace-id", trace)])).await; + let opts = with_xai(options("Hi"), json!({ "withTimestamps": true })); + let result = wiremock_speech_model(&server, None) + .do_generate(&opts) + .await + .unwrap(); + + assert!(matches!(&result.audio, AudioData::Binary(audio) if audio == &AUDIO)); + assert_eq!( + xai_metadata(&result), + json!({ + "traceId": trace, + "duration": 1.19, + "contentType": "audio/mpeg", + "audioTimestamps": { + "graphChars": ["H", "i"], + "graphTimes": [[0.04, 0.06], [0.06, 0.1]], + }, + }) + ); + assert_eq!(result.response.body.unwrap()["audio"], "AQIDBA=="); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should extract the trace id from binary responses + #[tokio::test] + async fn should_extract_the_trace_id_from_binary_responses() { + let trace = "06e3dab5-e3ba-4c6b-83a6-1e9ea11d78af"; + let server = server_with(audio_response("audio/mpeg", &[("x-trace-id", trace)])).await; + let result = wiremock_speech_model(&server, None) + .do_generate(&options("Hello from the AI SDK!")) + .await + .unwrap(); + + assert_eq!(xai_metadata(&result), json!({ "traceId": trace })); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should return empty provider metadata when xAI headers are absent + #[tokio::test] + async fn should_return_empty_provider_metadata_when_xai_headers_are_absent() { + let server = server_with(audio_response("audio/mpeg", &[])).await; + let result = wiremock_speech_model(&server, None) + .do_generate(&options("Hello from the AI SDK!")) + .await + .unwrap(); + + assert_eq!(xai_metadata(&result), json!({})); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should warn and use mp3 for unsupported output formats + #[tokio::test] + async fn should_warn_and_use_mp3_for_unsupported_output_formats() { + let server = server_with(audio_response("audio/mpeg", &[])).await; + let mut opts = options("Hello from the AI SDK!"); + opts.output_format = Some("flac".into()); + let result = wiremock_speech_model(&server, None) + .do_generate(&opts) + .await + .unwrap(); + + assert_eq!(request_body(&server).await["output_format"]["codec"], "mp3"); + assert!(is_unsupported(&result.warnings, "outputFormat")); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should warn and ignore bitRate for non-mp3 output + #[tokio::test] + async fn should_warn_and_ignore_bit_rate_for_non_mp3_output() { + let server = server_with(audio_response("audio/wav", &[])).await; + let mut opts = with_xai( + options("Hello from the AI SDK!"), + json!({ "bitRate": 192000 }), + ); + opts.output_format = Some("wav".into()); + let result = wiremock_speech_model(&server, None) + .do_generate(&opts) + .await + .unwrap(); + + assert_eq!( + request_body(&server).await["output_format"], + json!({ "codec": "wav" }) + ); + assert!(is_unsupported(&result.warnings, "providerOptions")); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should warn when instructions are provided + #[tokio::test] + async fn should_warn_when_instructions_are_provided() { + let server = server_with(audio_response("audio/mpeg", &[])).await; + let mut opts = options("Hello from the AI SDK!"); + opts.instructions = Some("Speak cheerfully".into()); + let result = wiremock_speech_model(&server, None) + .do_generate(&opts) + .await + .unwrap(); + + assert!(is_unsupported(&result.warnings, "instructions")); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should pass headers and the xAI user agent + #[tokio::test] + async fn should_pass_headers_and_the_xai_user_agent() { + let server = server_with(audio_response("audio/mpeg", &[])).await; + let provider_headers = HashMap::from([( + "Custom-Provider-Header".to_string(), + Some("provider-header-value".to_string()), + )]); + let mut opts = options("Hello from the AI SDK!"); + opts.headers = Some(HashMap::from([( + "Custom-Request-Header".to_string(), + "request-header-value".to_string(), + )])); + wiremock_speech_model(&server, Some(provider_headers)) + .do_generate(&opts) + .await + .unwrap(); + + let requests = server.received_requests().await.unwrap(); + let header = |name: &str| requests[0].headers.get(name).unwrap().to_str().unwrap(); + assert_eq!(header("authorization"), "Bearer test-key"); + assert_eq!(header("content-type"), "application/json"); + assert_eq!(header("custom-provider-header"), "provider-header-value"); + assert_eq!(header("custom-request-header"), "request-header-value"); + assert!(header("user-agent").contains("ai-sdk-xai/5.0.12")); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should return binary audio data + #[tokio::test] + async fn should_return_binary_audio_data() { + let server = server_with(audio_response("audio/mpeg", &[])).await; + let result = wiremock_speech_model(&server, None) + .do_generate(&options("Hello from the AI SDK!")) + .await + .unwrap(); + + assert!(matches!(&result.audio, AudioData::Binary(audio) if audio == &AUDIO)); + assert!(result.warnings.is_empty()); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should handle API errors + #[tokio::test] + async fn should_handle_api_errors() { + let server = server_with( + ResponseTemplate::new(400).set_body_string( + json!({ "error": { "message": "Invalid text", "type": "invalid_request_error" } }) + .to_string(), + ), + ) + .await; + let error = wiremock_speech_model(&server, None) + .do_generate(&options("Hello from the AI SDK!")) + .await + .unwrap_err(); + + let AiMuxError::ApiCall(error) = error else { + panic!("expected an ApiCall error, got {error:?}"); + }; + assert_eq!(error.message, "Invalid text"); + assert_eq!(error.status_code, Some(400)); + } + + /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should use the real date when no custom date provider is specified + #[tokio::test] + async fn should_use_the_real_date_when_no_custom_date_provider_is_specified() { + let server = server_with(audio_response("audio/mpeg", &[])).await; + let before = chrono::Utc::now().timestamp_millis(); + let result = wiremock_speech_model(&server, None) + .do_generate(&options("Hello from the AI SDK!")) + .await + .unwrap(); + let after = chrono::Utc::now().timestamp_millis(); + + // A `Date`: millisecond precision. + let timestamp = result.response.timestamp.unwrap().unix_millis(); + assert!(before <= timestamp && timestamp <= after); + assert_eq!(result.response.model_id.as_deref(), Some("")); + } +} + +/// Checks the provider id and the empty model id; specificationVersion is not checked. partial: should expose correct provider and model information; full port in src/xai/tests/ +#[tokio::test] +async fn should_expose_correct_provider_and_model_information() { + let model = wiremock_speech_model(&MockServer::start().await, None); + assert_eq!(model.provider(), "xai.speech"); + assert_eq!(model.model_id(), ""); +} + +/// Checks the response model id and the x-request-id header; the timestamp is only checked as present. partial: should include response timestamp, model id, and headers; full port in src/xai/tests/ +#[tokio::test] +async fn should_include_response_timestamp_model_id_and_headers() { + let server = server_with(audio_response( + "audio/mpeg", + &[("x-request-id", "test-request-id")], + )) + .await; + let result = wiremock_speech_model(&server, None) + .do_generate(&options("Hello from the AI SDK!")) + .await + .unwrap(); + + assert_eq!(result.response.model_id.as_deref(), Some("")); + assert_eq!( + result + .response + .headers + .unwrap() + .get("x-request-id") + .map(String::as_str), + Some("test-request-id") + ); + assert!(result.response.timestamp.is_some()); } diff --git a/aimux-providers/src/xai/tests/supports_reasoning_effort.rs b/aimux-providers/src/xai/tests/supports_reasoning_effort.rs index e7123aa70..f844b43fb 100644 --- a/aimux-providers/src/xai/tests/supports_reasoning_effort.rs +++ b/aimux-providers/src/xai/tests/supports_reasoning_effort.rs @@ -1,6 +1,6 @@ //! Upstream `supports-reasoning-effort.test.ts`. -use crate::xai::convert::supports_reasoning_effort; +use crate::xai::supports_reasoning_effort::supports_reasoning_effort; #[allow( clippy::module_inception, diff --git a/aimux-providers/src/xai/tests/transcription_model.rs b/aimux-providers/src/xai/tests/transcription_model.rs index 246e65642..a4794031e 100644 --- a/aimux-providers/src/xai/tests/transcription_model.rs +++ b/aimux-providers/src/xai/tests/transcription_model.rs @@ -24,18 +24,32 @@ //! header. //! - `argument: 'providerOptions'`: aimux's `InvalidArgument` carries only a //! message, so the message is checked. +//! +//! The older ports (formerly `tests/xai_transcription_test.rs`) follow the +//! wiremock harness helpers: they run against a loopback wiremock `MockServer`. use std::collections::HashMap; -use std::sync::Arc; - -use aimux_core::shared::SharedProviderOptions; -use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use aimux_core::error::AiMuxError; +use aimux_core::shared::{SharedProviderOptions, Warning}; +use aimux_core::transcription_model::{ + AudioChunk, AudioInput, InputAudioFormat, TranscriptionCallOptions, TranscriptionModel, + TranscriptionResponse, TranscriptionResult, TranscriptionStreamOptions, + TranscriptionStreamPart, +}; +use futures::{SinkExt, StreamExt}; use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::handshake::server::{Request, Response}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; use crate::test_support::mock_fetch::{Canned, MockFetch, Seen}; -use crate::xai::XaiTranscriptionModel; use crate::xai::options::NAMESPACE; -use crate::{XAIProviderSettings, create_xai}; +use crate::{XAIProviderSettings, XaiTranscriptionModel, create_xai}; /// `createXai({ apiKey: 'test-api-key', headers }).transcription()`, sending /// through `fetch`. @@ -115,6 +129,303 @@ fn assert_user_agent_names_the_package(seen: &Seen) { ); } +// ── The wiremock harness of the older ports ────────────────────────────────── + +type Headers = HashMap>; + +fn model(base_url: &str, headers: Option) -> XaiTranscriptionModel { + create_xai(XAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(base_url.into()), + headers, + ..Default::default() + }) + .expect("valid settings") + .transcription() +} + +fn wiremock_xai_options(value: Value) -> SharedProviderOptions { + SharedProviderOptions::from([("xai".to_string(), value.as_object().unwrap().clone())]) +} + +// ── doGenerate helpers ────────────────────────────────────────────────────── + +fn call(provider_options: Option, media_type: &str) -> TranscriptionCallOptions { + let mut options = + TranscriptionCallOptions::new(AudioInput::Binary(vec![1, 2, 3, 4]), media_type); + options.provider_options = provider_options.map(wiremock_xai_options); + options +} + +fn json_response(body: Value, headers: &[(&str, &str)]) -> ResponseTemplate { + headers.iter().fold( + ResponseTemplate::new(200).set_body_json(body), + |t, (k, v)| t.insert_header(*k, *v), + ) +} + +async fn mock_stt(server: &MockServer, template: ResponseTemplate) { + Mock::given(method("POST")) + .and(path("/stt")) + .respond_with(template) + .mount(server) + .await; +} + +fn wiremock_default_body() -> Value { + json!({ + "text": "Hello from the AI SDK!", + "language": "en", + "duration": 2.5, + "words": [ + {"text": "Hello", "start": 0, "end": 1}, + {"text": "from the AI SDK!", "start": 1, "end": 2.5} + ] + }) +} + +async fn generate( + server: &MockServer, + options: &TranscriptionCallOptions, + headers: Option, +) -> TranscriptionResult { + model(&server.uri(), headers) + .do_generate(options) + .await + .expect("do_generate") +} + +struct Part { + name: String, + filename: Option, + content_type: Option, + value: String, +} + +fn quoted(line: &str, key: &str) -> Option { + let start = line.find(&format!(" {key}=\""))? + key.len() + 3; + let end = line[start..].find('"')?; + Some(line[start..start + end].to_string()) +} + +/// The multipart fields of a captured request, in wire order. +fn parse_multipart(request: &wiremock::Request) -> Vec { + let content_type = request + .headers + .get("content-type") + .unwrap() + .to_str() + .unwrap(); + let boundary = content_type.split("boundary=").nth(1).unwrap(); + let body = String::from_utf8_lossy(&request.body).into_owned(); + body.split(&format!("--{boundary}")) + .skip(1) + .map(|chunk| chunk.trim_start_matches("\r\n")) + .filter(|chunk| !chunk.starts_with("--")) + .map(|chunk| { + let (head, value) = chunk.split_once("\r\n\r\n").unwrap(); + let line = |prefix: &str| head.lines().find(|l| l.starts_with(prefix)); + Part { + name: quoted(line("Content-Disposition").unwrap(), "name").unwrap(), + filename: quoted(line("Content-Disposition").unwrap(), "filename"), + content_type: line("Content-Type").map(|l| l["Content-Type: ".len()..].to_string()), + value: value.strip_suffix("\r\n").unwrap_or(value).to_string(), + } + }) + .collect() +} + +fn full_options() -> Value { + json!({ + "audioFormat": "pcm", "sampleRate": 16000, "language": "en", "format": true, + "multichannel": true, "channels": 2, "diarize": true, + "keyterm": ["AI SDK", "Grok"], "fillerWords": true + }) +} + +// ── doStream helpers ──────────────────────────────────────────────────────── + +#[derive(Default, Clone)] +struct Capture { + uri: String, + headers: HashMap, + /// Frames that arrived before the server sent `transcript.created`. + early_frames: usize, + binary: Vec>, + texts: Vec, +} + +async fn serve(listener: TcpListener, cap: Arc>, events: Vec) { + let (stream, _) = listener.accept().await.unwrap(); + let handshake_cap = Arc::clone(&cap); + #[allow(clippy::result_large_err)] + let handshake = move |req: &Request, resp: Response| { + let mut cap = handshake_cap.lock().unwrap(); + cap.uri = req.uri().to_string(); + cap.headers = req + .headers() + .iter() + .filter_map(|(k, v)| Some((k.as_str().to_string(), v.to_str().ok()?.to_string()))) + .collect(); + Ok(resp) + }; + let mut ws = tokio_tungstenite::accept_hdr_async(stream, handshake) + .await + .expect("ws handshake"); + + // The client holds audio until `transcript.created`. + if let Ok(Some(Ok(_))) = tokio::time::timeout(Duration::from_millis(100), ws.next()).await { + cap.lock().unwrap().early_frames += 1; + } + ws.send(Message::Text( + json!({"type": "transcript.created"}).to_string(), + )) + .await + .unwrap(); + loop { + match ws.next().await { + Some(Ok(Message::Binary(bytes))) => cap.lock().unwrap().binary.push(bytes.to_vec()), + Some(Ok(Message::Text(text))) => { + let text = text.as_str().to_string(); + cap.lock().unwrap().texts.push(text.clone()); + if text.contains("audio.done") { + break; + } + } + Some(Ok(_)) => {} + _ => return, + } + } + for event in events { + ws.send(Message::Text(event.to_string())).await.unwrap(); + } + // Wait for the client to close. + let _ = tokio::time::timeout(Duration::from_secs(2), async { + while let Some(Ok(message)) = ws.next().await { + if message.is_close() { + break; + } + } + }) + .await; +} + +async fn start(events: Vec) -> (String, Arc>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); + let cap = Arc::new(Mutex::new(Capture::default())); + tokio::spawn(serve(listener, Arc::clone(&cap), events)); + (base_url, cap) +} + +fn wiremock_stream_options( + provider_options: Option, + format_type: &str, + rate: Option, +) -> TranscriptionStreamOptions { + TranscriptionStreamOptions { + audio: Box::pin(futures::stream::iter(vec![AudioChunk::Binary(vec![ + 1, 2, 3, + ])])), + input_audio_format: InputAudioFormat { + format_type: format_type.to_string(), + rate, + }, + provider_options: provider_options.map(wiremock_xai_options), + abort_signal: None, + headers: None, + include_raw_chunks: false, + timeout: None, + } +} + +struct Run { + cap: Capture, + parts: Vec>, + response: Option, +} + +impl Run { + /// The parts as JSON; panics on a stream error. + fn values(&self) -> Vec { + self.parts + .iter() + .map(|p| serde_json::to_value(p.as_ref().expect("stream part")).unwrap()) + .collect() + } +} + +async fn run_with( + options: TranscriptionStreamOptions, + headers: Option, + events: Vec, +) -> Run { + let (base_url, cap) = start(events).await; + let mut result = model(&base_url, headers) + .do_stream(options) + .expect("the model has do_stream") + .await + .expect("do_stream"); + let response = result.response.take(); + let parts = tokio::time::timeout(Duration::from_secs(10), result.stream.collect::>()) + .await + .expect("stream timed out"); + let cap = cap.lock().unwrap().clone(); + Run { + cap, + parts, + response, + } +} + +async fn run(provider_options: Option, events: Vec) -> Run { + run_with( + wiremock_stream_options(provider_options, "audio/pcm", Some(16000)), + None, + events, + ) + .await +} + +fn partial(text: &str, is_final: bool, speech_final: bool, span: Option<(f64, f64)>) -> Value { + let mut event = json!({ + "type": "transcript.partial", "text": text, + "is_final": is_final, "speech_final": speech_final + }); + if let Some((start, duration)) = span { + event["start"] = json!(start); + event["duration"] = json!(duration); + } + event +} + +fn done(text: &str, duration: f64) -> Value { + json!({"type": "transcript.done", "text": text, "duration": duration}) +} + +fn kinds(parts: &[Value]) -> Vec<&str> { + parts.iter().map(|p| p["type"].as_str().unwrap()).collect() +} + +fn of_kind(parts: &[Value], kind: &str) -> Vec { + parts + .iter() + .filter(|p| p["type"] == kind) + .cloned() + .collect() +} + +fn finish(text: &str, language: Value, duration: Option) -> Value { + let mut part = json!({"type": "finish", "text": text, "segments": []}); + if !language.is_null() { + part["language"] = language; + } + if let Some(duration) = duration { + part["durationInSeconds"] = json!(duration); + } + part +} + mod xai_transcription_model { use super::*; @@ -204,6 +515,131 @@ mod do_generate { assert_eq!(headers["x-request-id"], "test-request-id"); assert_eq!(headers["x-ratelimit-remaining"], "123"); } + + /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doGenerate › should send a multipart request with the audio file + #[tokio::test] + async fn should_send_a_multipart_request_with_the_audio_file() { + let server = MockServer::start().await; + mock_stt(&server, json_response(wiremock_default_body(), &[])).await; + + generate(&server, &call(None, "audio/wav"), None).await; + + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests[0].method.as_str(), "POST"); + assert_eq!(requests[0].url.path(), "/stt"); + let parts = parse_multipart(&requests[0]); + let file = parts.iter().find(|p| p.name == "file").expect("file field"); + assert_eq!(file.filename.as_deref(), Some("audio.wav")); + assert_eq!(file.content_type.as_deref(), Some("audio/wav")); + } + + /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doGenerate › should map provider options onto xAI request fields + #[tokio::test] + async fn should_map_provider_options_onto_xai_request_fields() { + let server = MockServer::start().await; + mock_stt(&server, json_response(wiremock_default_body(), &[])).await; + + generate(&server, &call(Some(full_options()), "audio/pcm"), None).await; + + let parts = parse_multipart(&server.received_requests().await.unwrap()[0]); + let values = |name: &str| -> Vec<&str> { + parts + .iter() + .filter(|p| p.name == name) + .map(|p| p.value.as_str()) + .collect() + }; + for (name, expected) in [ + ("audio_format", "pcm"), + ("sample_rate", "16000"), + ("language", "en"), + ("format", "true"), + ("multichannel", "true"), + ("channels", "2"), + ("diarize", "true"), + ("filler_words", "true"), + ] { + assert_eq!(values(name), vec![expected], "field {name}"); + } + assert_eq!(values("keyterm"), vec!["AI SDK", "Grok"]); + } + + /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doGenerate › should append file after all other multipart fields + #[tokio::test] + async fn should_append_file_after_all_other_multipart_fields() { + let server = MockServer::start().await; + mock_stt( + &server, + json_response(json!({"text": "Hello from the AI SDK!"}), &[]), + ) + .await; + + generate(&server, &call(Some(full_options()), "audio/pcm"), None).await; + + let parts = parse_multipart(&server.received_requests().await.unwrap()[0]); + let names: Vec<&str> = parts.iter().map(|p| p.name.as_str()).collect(); + assert_eq!( + names, + [ + "audio_format", + "sample_rate", + "language", + "format", + "multichannel", + "channels", + "diarize", + "filler_words", + "keyterm", + "keyterm", + "file" + ] + ); + } + + /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doGenerate › should extract text, segments, language, and duration + #[tokio::test] + async fn should_extract_text_segments_language_and_duration() { + let server = MockServer::start().await; + mock_stt(&server, json_response(wiremock_default_body(), &[])).await; + + let result = generate(&server, &call(None, "audio/wav"), None).await; + + assert_eq!(result.text, "Hello from the AI SDK!"); + assert_eq!(result.language.as_deref(), Some("en")); + assert_eq!(result.duration_in_seconds, Some(2.5)); + let segments: Vec<_> = result + .segments + .iter() + .map(|s| (s.text.as_str(), s.start_second, s.end_second)) + .collect(); + assert_eq!( + segments, + [("Hello", 0.0, 1.0), ("from the AI SDK!", 1.0, 2.5)] + ); + assert!(result.warnings.is_empty()); + } + + /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doGenerate › should handle missing words, duration, and empty language + #[tokio::test] + async fn should_handle_missing_words_duration_and_empty_language() { + let server = MockServer::start().await; + mock_stt( + &server, + json_response( + json!({"text": "Hello from the AI SDK!", "language": ""}), + &[], + ), + ) + .await; + + let result = generate(&server, &call(None, "audio/wav"), None).await; + + assert_eq!(result.text, "Hello from the AI SDK!"); + assert_eq!(result.language, None); + assert_eq!(result.duration_in_seconds, None); + assert!(result.segments.is_empty()); + assert!(result.warnings.is_empty()); + } } #[cfg(feature = "realtime")] @@ -855,4 +1291,372 @@ mod do_stream { "the audio stream is cancelled" ); } + + /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doStream › should emit one transcript-final per utterance and reconstruct the finish text when transcript.done is empty + #[tokio::test] + async fn should_emit_one_final_per_utterance_and_reconstruct_finish_text() { + let events = vec![ + partial("Hello wor", false, false, Some((0.0, 0.8))), + partial("Hello world.", true, false, Some((0.0, 1.0))), + partial("Hello world.", true, true, Some((0.0, 1.0))), + done("", 1.0), + ]; + let parts = run(None, events).await.values(); + + assert_eq!( + of_kind(&parts, "transcript-final"), + vec![json!({ + "type": "transcript-final", "text": "Hello world.", "startSecond": 0.0, + "endSecond": 1.0 + })] + ); + // the speech_final:false re-send surfaces as a partial + assert_eq!(of_kind(&parts, "transcript-partial").len(), 2); + assert_eq!( + parts.last().unwrap(), + &finish("Hello world.", Value::Null, Some(1.0)) + ); + } + + /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doStream › should treat is_final fragments as partials and use the speech_final text for finish + #[tokio::test] + async fn should_treat_is_final_fragments_as_partials_and_use_speech_final_text() { + let events = vec![ + partial("No", true, false, Some((0.0, 0.5))), + partial(", I'm not", true, false, Some((0.5, 0.7))), + partial("No, I'm not.", true, true, Some((0.0, 1.2))), + done("", 1.2), + ]; + let parts = run(None, events).await.values(); + + assert_eq!( + kinds(&parts), + [ + "stream-start", + "transcript-partial", + "transcript-partial", + "transcript-final", + "finish" + ] + ); + assert_eq!(parts.last().unwrap()["text"], "No, I'm not."); + } + + /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doStream › should fall back to the latest pending text when no speech_final arrived before transcript.done + #[tokio::test] + async fn should_fall_back_to_latest_pending_text_without_speech_final() { + let events = vec![ + partial("Hello wor", false, false, None), + partial("Hello world", true, false, None), + done("", 1.0), + ]; + let parts = run(None, events).await.values(); + + let last = parts.last().unwrap(); + assert_eq!(kinds(&parts).last(), Some(&"finish")); + assert_eq!(last["text"], "Hello world"); + } + + /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doStream › should warn on unrecognized inputAudioFormat types + #[tokio::test] + async fn should_warn_on_unrecognized_input_audio_format_types() { + let options = wiremock_stream_options(None, "audio/wav", None); + let run = run_with(options, None, vec![done("Hello", 1.0)]).await; + + match &run.parts[0] { + Ok(TranscriptionStreamPart::StreamStart { warnings }) => match warnings.as_slice() { + [Warning::Other { message }] => { + assert!( + message.contains("Unrecognized inputAudioFormat.type \"audio/wav\""), + "{message}" + ); + } + other => panic!("expected one Other warning, got {other:?}"), + }, + other => panic!("expected StreamStart, got {other:?}"), + } + } +} + +/// Checks the provider id and the empty model id; specificationVersion is not checked. partial: should expose correct provider and model information; full port in src/xai/tests/ +#[test] +fn should_expose_correct_provider_and_model_information() { + let model = model("https://api.x.ai/v1", None); + assert_eq!(model.provider(), "xai.transcription"); + assert_eq!(model.model_id(), ""); +} + +/// Checks authorization, custom headers and the multipart content type; the user agent is checked without its version and the boundary only by prefix. partial: should pass headers and the xAI user agent; full port in src/xai/tests/ +#[tokio::test] +async fn should_pass_headers_and_the_xai_user_agent() { + let server = MockServer::start().await; + mock_stt(&server, json_response(wiremock_default_body(), &[])).await; + + let mut options = call(None, "audio/wav"); + options.headers = Some(HashMap::from([( + "Custom-Request-Header".to_string(), + "request-header-value".to_string(), + )])); + let provider_headers = HashMap::from([( + "Custom-Provider-Header".to_string(), + Some("provider-header-value".to_string()), + )]); + generate(&server, &options, Some(provider_headers)).await; + + let requests = server.received_requests().await.unwrap(); + let header = |name: &str| { + requests[0] + .headers + .get(name) + .unwrap() + .to_str() + .unwrap() + .to_string() + }; + assert_eq!(header("authorization"), "Bearer test-api-key"); + assert!(header("content-type").starts_with("multipart/form-data; boundary=")); + assert_eq!(header("custom-provider-header"), "provider-header-value"); + assert_eq!(header("custom-request-header"), "request-header-value"); + assert!(header("user-agent").contains("ai-sdk-xai/")); +} + +/// Checks the response model id, headers and body; the timestamp is only checked as present. partial: should include response timestamp, model id, and headers; full port in src/xai/tests/ +#[tokio::test] +async fn should_include_response_timestamp_model_id_and_headers() { + let server = MockServer::start().await; + let headers = [ + ("x-request-id", "test-request-id"), + ("x-ratelimit-remaining", "123"), + ]; + mock_stt(&server, json_response(wiremock_default_body(), &headers)).await; + + let result = generate(&server, &call(None, "audio/wav"), None).await; + + assert!(result.response.timestamp.is_some()); + assert_eq!(result.response.model_id.as_deref(), Some("")); + let response_headers = result.response.headers.unwrap(); + assert_eq!(response_headers["content-type"], "application/json"); + assert_eq!(response_headers["x-request-id"], "test-request-id"); + assert_eq!(response_headers["x-ratelimit-remaining"], "123"); + assert_eq!(result.response.body, Some(wiremock_default_body())); +} + +/// Checks the multichannel-without-channels error and that no socket is opened; the error's argument field is not checked. partial: should require channels when streaming multichannel audio; full port in src/xai/tests/ +#[tokio::test] +async fn should_require_channels_when_streaming_multichannel_audio() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); + + let options = wiremock_stream_options( + Some(json!({"multichannel": true})), + "audio/pcm", + Some(16000), + ); + let error = model(&base_url, None) + .do_stream(options) + .expect("the model has do_stream") + .await + .unwrap_err(); + + assert!( + matches!(&error, AiMuxError::InvalidArgument(message) if message.contains( + "providerOptions.xai.channels is required when providerOptions.xai.multichannel is true" + )), + "got {error:?}" + ); + // Nothing connected. + assert!( + tokio::time::timeout(Duration::from_millis(100), listener.accept()) + .await + .is_err() + ); +} + +/// Checks the streaming URL, handshake, audio frames and parts; the response timestamp is only checked as present. partial: should stream xAI STT over WebSocket; full port in src/xai/tests/ +#[tokio::test] +async fn should_stream_xai_stt_over_websocket() { + let options = json!({ + "language": "en", "diarize": true, "keyterm": ["AI SDK", "Grok"], + "streaming": { + "interimResults": true, "endpointing": 500, + "smartTurn": 0.7, "smartTurnTimeout": 3000 + } + }); + let events = vec![ + partial("Hel", false, false, Some((0.0, 0.5))), + partial("Hello", true, true, Some((0.0, 1.0))), + done("Hello", 1.0), + ]; + let run = run(Some(options), events).await; + + assert_eq!( + run.cap.uri, + "/stt?sample_rate=16000&encoding=pcm&language=en&diarize=true&interim_results=true&endpointing=500&smart_turn=0.7&smart_turn_timeout=3000&keyterm=AI+SDK&keyterm=Grok" + ); + assert_eq!(run.cap.headers["authorization"], "Bearer test-api-key"); + // Audio waits for `transcript.created`, then the frames and `audio.done` follow. + assert_eq!(run.cap.early_frames, 0); + assert_eq!(run.cap.binary, vec![vec![1u8, 2, 3]]); + assert_eq!(run.cap.texts.len(), 1); + assert_eq!( + serde_json::from_str::(&run.cap.texts[0]).unwrap(), + json!({"type": "audio.done"}) + ); + + assert_eq!( + run.values(), + vec![ + json!({"type": "stream-start", "warnings": []}), + json!({ + "type": "transcript-partial", "text": "Hel", "startSecond": 0.0, + "durationInSeconds": 0.5 + }), + json!({ + "type": "transcript-final", "text": "Hello", "startSecond": 0.0, + "endSecond": 1.0 + }), + finish("Hello", json!("en"), Some(1.0)), + ] + ); + let response = run.response.unwrap(); + assert_eq!(response.model_id.as_deref(), Some("")); + assert!(response.timestamp.is_some()); +} + +/// Checks that a provider header set to None is left off the WebSocket handshake (analogue: per-call headers cannot hold None). partial: should strip undefined header values before the WebSocket constructor; full port in src/xai/tests/ +#[tokio::test] +async fn should_not_send_removed_header_values_on_the_handshake() { + let headers = HashMap::from([ + ( + "Custom-Header".to_string(), + Some("custom-value".to_string()), + ), + ("X-Unset".to_string(), None), + ]); + let options = wiremock_stream_options(None, "audio/pcm", Some(16000)); + let run = run_with(options, Some(headers), vec![done("Hello", 1.0)]).await; + + assert_eq!(run.cap.headers["authorization"], "Bearer test-api-key"); + assert_eq!(run.cap.headers["custom-header"], "custom-value"); + assert!(!run.cap.headers.contains_key("x-unset")); +} + +/// Checks that finalized utterances are joined into the finish text; the finish part type is not checked. partial: should join finalized utterances per channel when transcript.done is empty; full port in src/xai/tests/ +#[tokio::test] +async fn should_join_finalized_utterances_when_transcript_done_is_empty() { + let events = vec![ + partial("First utterance.", true, true, Some((0.0, 1.0))), + partial("Second utterance.", true, true, Some((1.0, 1.0))), + done("", 2.0), + ]; + let parts = run(None, events).await.values(); + + assert_eq!( + parts.last().unwrap()["text"], + "First utterance. Second utterance." + ); +} + +/// Checks that an error event fails the stream with the server message; the WebSocket close is not checked. partial: should error the stream with the server message on error events; full port in src/xai/tests/ +#[tokio::test] +async fn should_error_the_stream_with_the_server_message_on_error_events() { + let events = vec![json!({"type": "error", "message": "invalid sample_rate"})]; + let run = run(None, events).await; + + assert!(matches!( + run.parts[0], + Ok(TranscriptionStreamPart::StreamStart { .. }) + )); + assert_eq!(run.parts.len(), 2, "parts: {:?}", run.parts); + assert!( + matches!(&run.parts[1], Err(AiMuxError::Other(message)) if message == "invalid sample_rate"), + "got {:?}", + run.parts[1] + ); +} + +/// Only `transcribe` downloads an audio URL (layer 5's `AudioInput::Url`); +/// the model itself refuses one before any request, as every transcription +/// model does. +#[tokio::test] +async fn url_audio_is_refused_before_any_request() { + let server = MockServer::start().await; + let mut options = call(None, "audio/wav"); + options.audio = AudioInput::Url("https://cdn.test/a.wav".into()); + let error = model(&server.uri(), None) + .do_generate(&options) + .await + .unwrap_err(); + assert_eq!(error.to_string(), AudioInput::url_error().to_string()); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +/// A refused connection errors the stream: upstream connects inside the +/// stream, so `doStream` resolves and the stream carries the failure. +#[tokio::test] +async fn a_connect_failure_errors_the_stream() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); + drop(listener); + + let result = model(&base_url, None) + .do_stream(wiremock_stream_options(None, "audio/pcm", Some(16000))) + .expect("the model has do_stream") + .await + .expect("do_stream resolves; the connect failure arrives on the stream"); + let parts = tokio::time::timeout(Duration::from_secs(10), result.stream.collect::>()) + .await + .expect("stream timed out"); + + assert_eq!(parts.len(), 1, "{parts:?}"); + assert!(parts[0].is_err(), "{parts:?}"); +} + +/// Dropping the stream closes the socket: upstream's `cancel` calls +/// `connection.close()`. +#[tokio::test] +async fn dropping_the_stream_closes_the_socket() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); + let (closed_tx, closed_rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap(); + ws.send(Message::Text( + json!({"type": "transcript.created"}).to_string(), + )) + .await + .unwrap(); + let mut closed = false; + while let Some(Ok(message)) = ws.next().await { + if message.is_close() { + closed = true; + break; + } + } + let _ = closed_tx.send(closed); + }); + + let mut options = wiremock_stream_options(None, "audio/pcm", Some(16000)); + options.audio = Box::pin(futures::stream::pending()); + let mut result = model(&base_url, None) + .do_stream(options) + .expect("the model has do_stream") + .await + .expect("do_stream"); + let first = result.stream.next().await; + assert!( + matches!(first, Some(Ok(TranscriptionStreamPart::StreamStart { .. }))), + "{first:?}" + ); + drop(result); + + let closed = tokio::time::timeout(Duration::from_secs(5), closed_rx) + .await + .expect("the server saw the socket end") + .unwrap(); + assert!( + closed, + "no close frame was sent when the stream was dropped" + ); } diff --git a/aimux-providers/src/xai/tests/video_model.rs b/aimux-providers/src/xai/tests/video_model.rs index 38a7f8ee3..efd19ea8a 100644 --- a/aimux-providers/src/xai/tests/video_model.rs +++ b/aimux-providers/src/xai/tests/video_model.rs @@ -1,5 +1,7 @@ -//! Ports of `src/xai-video-model.test.ts` from `@ai-sdk/xai@5.0.12`: the -//! `doStart` and `doStatus` cases that `tests/xai_video_test.rs` does not cover. +//! Ports of `src/xai-video-model.test.ts` from `@ai-sdk/xai@5.0.12`, in two +//! harnesses: the `doStart` and `doStatus` cases on the in-crate [`MockFetch`], +//! then the older ports (formerly `tests/xai_video_test.rs`) on a loopback +//! wiremock `MockServer`, after the wiremock harness helpers. //! //! Harness differences from the TypeScript (expectations are unchanged): //! - `createTestServer` is the in-crate [`MockFetch`]: it answers requests in @@ -20,17 +22,19 @@ //! - `createModel`'s default `api-key` header is not used: no case here checks a //! header, and the Rust provider sends `Authorization: Bearer`. +use std::collections::HashMap; use std::sync::Arc; -use serde_json::{Value, json}; - use aimux_core::Timestamp; use aimux_core::error::AiMuxError; use aimux_core::shared::{AspectRatio, Size, Warning, provider_namespace}; use aimux_core::video_model::{ - VideoCallOptions, VideoFile, VideoFileData, VideoModel, VideoOperationStart, - VideoOperationStatus, VideoResult, + VideoCallOptions, VideoData, VideoFile, VideoFileData, VideoFrameImage, VideoFrameType, + VideoModel, VideoOperationStart, VideoOperationStatus, VideoResult, }; +use serde_json::{Map, Value, json}; +use wiremock::matchers::method; +use wiremock::{Mock, MockServer, Request, ResponseTemplate}; use crate::test_support::mock_fetch::{Canned, MockFetch, Seen}; use crate::xai::options::NAMESPACE; @@ -237,6 +241,161 @@ fn names_feature(warnings: &[Warning], feature: &str) -> bool { }) } +// ── The wiremock harness of the older ports ────────────────────────────────── + +// ---------------------------------------------------------------- helpers + +fn done_status() -> Value { + json!({ + "status": "done", + "video": {"url": VIDEO_URL, "duration": 5, "respect_moderation": true}, + "model": V1, + "progress": 100 + }) +} + +fn json_response(body: Value) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(body) +} + +async fn server_with(post: Value, status: ResponseTemplate) -> MockServer { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(json_response(post)) + .mount(&server) + .await; + Mock::given(method("GET")) + .respond_with(status) + .mount(&server) + .await; + server +} + +fn model_with(server: &MockServer, model_id: &str, headers: &[(&str, &str)]) -> XaiVideoModel { + create_xai(XAIProviderSettings { + api_key: Some("test-key".into()), + base_url: Some(server.uri()), + headers: (!headers.is_empty()).then(|| { + headers + .iter() + .map(|(k, v)| ((*k).to_string(), Some((*v).to_string()))) + .collect() + }), + ..Default::default() + }) + .unwrap() + .video(model_id) +} + +fn base() -> VideoCallOptions { + VideoCallOptions::new(PROMPT) +} + +/// Sets `providerOptions.xai`. +fn xai(mut options: VideoCallOptions, value: Value) -> VideoCallOptions { + options + .provider_options + .insert("xai".into(), value.as_object().unwrap().clone()); + options +} + +fn ar(s: &str) -> Option { + Some(AspectRatio::parse(s).unwrap()) +} + +fn last_frame(file: VideoFile) -> Option> { + Some(vec![VideoFrameImage { + image: file, + frame_type: VideoFrameType::LastFrame, + }]) +} + +fn header(request: &Request, name: &str) -> String { + request + .headers + .get(name) + .map(|v| v.to_str().unwrap().to_owned()) + .unwrap_or_default() +} + +struct Started { + start: VideoOperationStart, + method: String, + path: String, + body: Value, +} + +async fn try_start(model_id: &str, options: &VideoCallOptions) -> Result { + let server = server_with( + json!({"request_id": "req-123"}), + json_response(done_status()), + ) + .await; + let start = model_with(&server, model_id, &[]).do_start(options).await?; + let request = server.received_requests().await.unwrap().remove(0); + Ok(Started { + start, + method: request.method.to_string(), + path: request.url.path().to_owned(), + body: serde_json::from_slice(&request.body).unwrap(), + }) +} + +async fn wiremock_start(model_id: &str, options: VideoCallOptions) -> Started { + try_start(model_id, &options).await.unwrap() +} + +/// The `details` of every `Unsupported` warning for `feature`. +fn details<'a>(warnings: &'a [Warning], feature: &str) -> Vec<&'a str> { + warnings + .iter() + .filter_map(|w| match w { + Warning::Unsupported { + feature: f, + details, + } if f == feature => Some(details.as_deref().unwrap_or("")), + _ => None, + }) + .collect() +} + +fn warned(warnings: &[Warning], feature: &str) -> bool { + !details(warnings, feature).is_empty() +} + +async fn poll_once( + status: ResponseTemplate, + request_id: &str, +) -> (Result, Vec) { + let server = server_with(json!({"request_id": "req-123"}), status).await; + let result = model_with(&server, V1, &[]) + .do_status(&json!({ "requestId": request_id }), &base()) + .await; + (result, server.received_requests().await.unwrap()) +} + +fn wiremock_completed(result: Result) -> VideoResult { + match result { + Ok(VideoOperationStatus::Completed(done)) => done, + other => panic!("expected completed, got {other:?}"), + } +} + +fn xai_meta(result: &VideoResult) -> Map { + result.provider_metadata.as_ref().unwrap()["xai"].clone() +} + +fn assert_terminal(result: Result, message: &str) { + match result { + Err(AiMuxError::ApiCall(e)) => { + assert_eq!(e.message, message); + assert!(!e.is_retryable); + assert_eq!(e.status_code, Some(200)); + } + other => panic!("expected terminal api error, got {other:?}"), + } +} + mod xai_video_model { use super::*; @@ -253,6 +412,13 @@ mod xai_video_model { assert_eq!(video.specification_version(), "v4"); assert_eq!(video.max_videos_per_call(), Some(1)); } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › constructor › should send the grok-imagine-video-1.5 model id in the request body + #[tokio::test] + async fn sends_the_1_5_model_id_in_the_body() { + let s = wiremock_start(V15, base()).await; + assert_eq!(s.body["model"], V15); + } } mod do_start { @@ -995,6 +1161,158 @@ mod xai_video_model { let (_, request) = start(V15, &options).await; assert!(request.json_body().get("referenceVoiceIds").is_none()); } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should map current xAI generation options + #[tokio::test] + async fn maps_current_generation_options() { + let mut o = xai( + base(), + json!({ + "storageOptions": { + "filename": "result.mp4", + "expiresAfter": 86_400, + "publicUrl": {"expiresAfter": 3_600} + }, + "keyframes": [{"imageUrl": "https://example.com/middle.png", "timestampSeconds": 2.5}] + }), + ); + o.generate_audio = Some(false); + o.frame_images = last_frame(url_file("https://example.com/end.png", None)); + let s = wiremock_start(V15, o).await; + assert_eq!(s.body["generate_audio"], false); + assert_eq!( + s.body["last_frame"], + json!({"url": "https://example.com/end.png"}) + ); + assert_eq!( + s.body["storage_options"], + json!({ + "filename": "result.mp4", + "expires_after": 86_400, + "public_url": {"expires_after": 3_600} + }) + ); + assert_eq!( + s.body["keyframes"], + json!([{"image": {"url": "https://example.com/middle.png"}, "timestamp_s": 2.5}]) + ); + assert_eq!(s.start.operation, json!({"requestId": "req-123"})); + } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should warn and omit last_frame for grok-imagine-video + #[tokio::test] + async fn warns_and_omits_last_frame_for_grok_imagine_video() { + let mut o = base(); + o.frame_images = last_frame(url_file("https://example.com/end.png", None)); + let s = wiremock_start(V1, o).await; + assert!(s.body.get("last_frame").is_none()); + assert_eq!( + details(&s.start.warnings, "frameImages"), + [ + "xAI only supports last_frame with \"grok-imagine-video-1.5\". \ + The last frame was ignored." + ] + ); + } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should pass headers + #[tokio::test] + async fn do_start_passes_headers() { + let server = server_with( + json!({"request_id": "req-123"}), + json_response(done_status()), + ) + .await; + let model = model_with(&server, V1, &[("X-Custom", "value")]); + let mut o = base(); + o.headers = Some(HashMap::from([( + "X-Request-Header".to_string(), + "request-value".to_string(), + )])); + model.do_start(&o).await.unwrap(); + let request = server.received_requests().await.unwrap().remove(0); + assert_eq!(header(&request, "authorization"), "Bearer test-key"); + assert_eq!(header(&request, "x-custom"), "value"); + assert_eq!(header(&request, "x-request-header"), "request-value"); + } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should throw when no request_id returned + #[tokio::test] + async fn do_start_fails_without_request_id() { + let server = server_with(json!({}), json_response(done_status())).await; + let err = model_with(&server, V1, &[]) + .do_start(&base()) + .await + .unwrap_err(); + match err { + AiMuxError::Other(message) => { + assert_eq!(message, "No request_id returned from xAI API.") + } + other => panic!("unexpected error: {other:?}"), + } + } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should return warnings for unsupported features + #[tokio::test] + async fn warns_for_fps_and_seed() { + let mut o = base(); + o.fps = Some(30); + o.seed = Some(42); + let s = wiremock_start(V1, o).await; + assert!(warned(&s.start.warnings, "fps")); + assert!(warned(&s.start.warnings, "seed")); + } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should warn when explicit R2V has no references at all + #[tokio::test] + async fn warns_when_reference_to_video_has_no_references() { + let s = wiremock_start(V15, xai(base(), json!({"mode": "reference-to-video"}))).await; + assert!(s.body.get("reference_images").is_none()); + assert!( + details(&s.start.warnings, "referenceImages")[0] + .contains("without reference images") + ); + } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should allow duration and aspectRatio with reference images + #[tokio::test] + async fn allows_duration_and_aspect_ratio_with_reference_images() { + let mut o = xai( + base(), + json!({"referenceImageUrls": ["https://example.com/ref.jpg"]}), + ); + o.duration = Some(8); + o.aspect_ratio = ar("16:9"); + let s = wiremock_start(V1, o).await; + assert_eq!(s.body["duration"], 8); + assert_eq!(s.body["aspect_ratio"], "16:9"); + assert!(!warned(&s.start.warnings, "duration")); + assert!(!warned(&s.start.warnings, "aspectRatio")); + } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should omit reference_audios for an empty referenceVoiceIds array + #[tokio::test] + async fn omits_reference_audios_for_empty_voice_ids() { + let o = xai( + base(), + json!({ + "mode": "reference-to-video", + "referenceImageUrls": ["https://example.com/ref1.jpg"], + "referenceVoiceIds": [] + }), + ); + let s = wiremock_start(V15, o).await; + assert!(s.body.get("reference_audios").is_none()); + assert!(!warned(&s.start.warnings, "referenceVoiceIds")); + } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should warn and omit referenceVoiceIds outside reference-to-video + #[tokio::test] + async fn warns_and_omits_voice_ids_outside_reference_to_video() { + let s = wiremock_start(V15, xai(base(), json!({"referenceVoiceIds": ["eve"]}))).await; + assert!(s.body.get("reference_audios").is_none()); + assert!(warned(&s.start.warnings, "referenceVoiceIds")); + } } mod do_status { @@ -1169,5 +1487,509 @@ mod xai_video_model { let done = completed(result); assert!(metadata_json(&done)["xai"].get("progress").is_none()); } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStatus › should return pending when status is pending + #[tokio::test] + async fn do_status_returns_pending() { + for template in [ + json_response(json!({"status": "pending"})), + // xAI answers 202 while running, sometimes with an empty body. + ResponseTemplate::new(202), + ResponseTemplate::new(202) + .set_body_json(json!({"status": "pending", "progress": 40})), + ] { + let (result, _) = poll_once(template, "req-123").await; + assert!( + matches!(result, Ok(VideoOperationStatus::Pending)), + "{result:?}" + ); + } + } + + /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStatus › should pass headers to request + #[tokio::test] + async fn do_status_passes_headers() { + let server = server_with(json!({}), json_response(done_status())).await; + let model = model_with(&server, V1, &[("X-Custom", "value")]); + let mut o = base(); + o.headers = Some(HashMap::from([( + "X-Request-Header".to_string(), + "request-value".to_string(), + )])); + model + .do_status(&json!({"requestId": "req-123"}), &o) + .await + .unwrap(); + let request = server.received_requests().await.unwrap().remove(0); + assert_eq!(header(&request, "authorization"), "Bearer test-key"); + assert_eq!(header(&request, "x-custom"), "value"); + assert_eq!(header(&request, "x-request-header"), "request-value"); + } } } + +/// Checks the provider id, model id and maxVideosPerCall; specificationVersion is not checked. partial: should expose correct provider and model information; full port in src/xai/tests/ +#[tokio::test] +async fn constructor_exposes_provider_and_model_information() { + let server = MockServer::start().await; + let model = model_with(&server, V1, &[]); + assert_eq!(model.provider(), "xai.video"); + assert_eq!(model.model_id(), V1); + assert_eq!(model.max_videos_per_call(), Some(1)); +} + +/// Checks the requestId operation and the response model id, timestamp and headers. partial: should return operation with requestId; should include response metadata; full port in src/xai/tests/ +#[tokio::test] +async fn do_start_returns_request_id_and_response_metadata() { + let s = wiremock_start(V1, base()).await; + assert_eq!(s.start.operation, json!({"requestId": "req-123"})); + let response = &s.start.response; + assert_eq!(response.model_id.as_deref(), Some(V1)); + assert!(response.timestamp.is_some()); + assert!(response.headers.is_some()); +} + +/// Checks the POST to /videos/generations with the minimal body and no warnings. partial: should pass correct request body; should return empty warnings for supported features; full port in src/xai/tests/ +#[tokio::test] +async fn do_start_posts_minimal_generation_body() { + let s = wiremock_start(V1, base()).await; + assert_eq!(s.method, "POST"); + assert_eq!(s.path, "/videos/generations"); + assert_eq!(s.body, json!({"model": V1, "prompt": PROMPT})); + assert!(s.start.warnings.is_empty()); +} + +/// Checks the POST path and video field for legacy, edit and extend modes. partial: should use edits endpoint for video editing (legacy videoUrl without mode); should use edits endpoint for video editing with explicit mode; should use extensions endpoint for extend-video mode; full port in src/xai/tests/ +#[tokio::test] +async fn selects_endpoint_by_mode() { + let cases = [ + (json!({"videoUrl": SRC}), "/videos/edits"), + ( + json!({"mode": "edit-video", "videoUrl": SRC}), + "/videos/edits", + ), + ( + json!({"mode": "extend-video", "videoUrl": SRC}), + "/videos/extensions", + ), + ]; + for (options, path) in cases { + let s = wiremock_start(V1, xai(base(), options)).await; + assert_eq!(s.method, "POST"); + assert_eq!(s.path, path); + assert_eq!(s.body["video"], json!({"url": SRC})); + } +} + +/// Checks the n warning for n=3 and its absence for n=1. partial: should warn when n > 1; should not warn when n is 1; full port in src/xai/tests/ +#[tokio::test] +async fn warns_only_when_n_exceeds_one() { + for (n, expected) in [(3, true), (1, false)] { + let mut o = base(); + o.n = n; + let s = wiremock_start(V1, o).await; + assert_eq!(warned(&s.start.warnings, "n"), expected, "n = {n}"); + } +} + +/// Checks that duration and aspect_ratio are sent in the body. partial: should send duration in request body; should send aspect_ratio in request body; full port in src/xai/tests/ +#[tokio::test] +async fn sends_duration_and_aspect_ratio() { + let mut o = base(); + o.duration = Some(10); + o.aspect_ratio = ar("9:16"); + let s = wiremock_start(V1, o).await; + assert_eq!(s.body["duration"], 10); + assert_eq!(s.body["aspect_ratio"], "9:16"); +} + +/// Checks that SDK sizes map to 720p, 480p and 1080p. partial: should map SDK resolution 1280x720 to 720p; should map SDK resolution 854x480 to 480p; should map SDK resolution 640x480 to 480p; should map SDK resolution 1920x1080 to 1080p; full port in src/xai/tests/ +#[tokio::test] +async fn maps_sdk_resolution_to_xai_resolution() { + let cases = [ + (V1, "1280x720", "720p"), + (V1, "854x480", "480p"), + (V1, "640x480", "480p"), + (V15, "1920x1080", "1080p"), + ]; + for (model, resolution, expected) in cases { + let mut o = base(); + o.resolution = size(resolution); + let s = wiremock_start(model, o).await; + assert_eq!(s.body["resolution"], expected, "{resolution}"); + } +} + +/// Checks that a provider resolution is sent and wins over the SDK size. partial: should pass through provider option resolution 1080p; should prefer provider option resolution over SDK resolution; full port in src/xai/tests/ +#[tokio::test] +async fn provider_option_resolution_wins() { + let s = wiremock_start(V15, xai(base(), json!({"resolution": "1080p"}))).await; + assert_eq!(s.body["resolution"], "1080p"); + + let mut o = xai(base(), json!({"resolution": "480p"})); + o.resolution = size("1280x720"); + let s = wiremock_start(V1, o).await; + assert_eq!(s.body["resolution"], "480p"); +} + +/// Checks the 1080p warning on grok-imagine-video (SDK and provider sources) and none on 1.5. partial: should warn when SDK resolution 1920x1080 is used with grok-imagine-video; should warn when provider resolution 1080p is used with grok-imagine-video; should not warn about 1080p with grok-imagine-video-1.5; full port in src/xai/tests/ +#[tokio::test] +async fn warns_about_1080p_only_on_grok_imagine_video() { + let mut sdk = base(); + sdk.resolution = size("1920x1080"); + let provider = xai(base(), json!({"resolution": "1080p"})); + for options in [sdk, provider] { + let s = wiremock_start(V1, options).await; + assert_eq!(s.body["resolution"], "1080p"); + assert!( + details(&s.start.warnings, "resolution")[0].contains("does not support 1080p"), + "{:?}", + s.start.warnings + ); + } + let mut o = base(); + o.resolution = size("1920x1080"); + assert!(wiremock_start(V15, o).await.start.warnings.is_empty()); +} + +/// Checks that unrecognized resolutions are left out of the body and warned about. partial: should warn for unrecognized resolution format; should warn and omit body resolution for completely unknown format; full port in src/xai/tests/ +#[tokio::test] +async fn warns_and_omits_unrecognized_resolution() { + for resolution in ["2560x1440", "3840x2160"] { + let mut o = base(); + o.resolution = size(resolution); + let s = wiremock_start(V1, o).await; + assert!(s.body.get("resolution").is_none()); + assert!(warned(&s.start.warnings, "resolution"), "{resolution}"); + } +} + +/// Checks that URL, binary and base64 image inputs are sent as image.url. partial: should send image object from URL-based image input; should send image object with data URI from file data; should send image object with data URI from base64 string; full port in src/xai/tests/ +#[tokio::test] +async fn sends_image_object_from_url_and_file_inputs() { + let cases = [ + ( + url_file("https://example.com/image.png", None), + "https://example.com/image.png", + ), + ( + VideoFile::File { + media_type: "image/png".into(), + data: VideoFileData::Binary(vec![137, 80, 78, 71]), + }, + "data:image/png;base64,iVBORw==", + ), + ( + VideoFile::File { + media_type: "image/jpeg".into(), + data: VideoFileData::Base64("aGVsbG8=".into()), + }, + "data:image/jpeg;base64,aGVsbG8=", + ), + ]; + for (image, expected) in cases { + let mut o = base(); + o.image = Some(image); + let s = wiremock_start(V1, o).await; + assert_eq!(s.body["image"], json!({"url": expected})); + } +} + +/// Checks which options warn in default, edit and extend modes. partial: should warn about duration in edit mode; should warn about aspectRatio in edit mode; should warn about resolution in edit mode; should not warn about duration outside edit mode; should not warn about aspectRatio outside edit mode; should not warn about resolution outside edit mode; should warn about aspectRatio in extension mode; should warn about resolution in extension mode; should not warn about duration in extension mode; full port in src/xai/tests/ +#[tokio::test] +async fn warns_about_unsupported_options_per_mode() { + type Setter = fn(&mut VideoCallOptions); + let setters: [(&str, Setter); 3] = [ + ("duration", |o| o.duration = Some(10)), + ("aspectRatio", |o| o.aspect_ratio = ar("16:9")), + ("resolution", |o| o.resolution = size("1280x720")), + ]; + for mode in [None, Some("edit-video"), Some("extend-video")] { + for (feature, set) in setters { + let mut o = base(); + set(&mut o); + if let Some(mode) = mode { + o = xai(o, json!({"mode": mode, "videoUrl": SRC})); + } + let expected = matches!( + (mode, feature), + (Some("edit-video"), _) | (Some("extend-video"), "aspectRatio" | "resolution") + ); + let s = wiremock_start(V1, o).await; + assert_eq!( + warned(&s.start.warnings, feature), + expected, + "{mode:?} {feature}" + ); + } + } +} + +/// Checks that unsupported fields are left out of edit and extend bodies. partial: should omit duration, aspect_ratio, and resolution from body in edit mode; should allow duration in extension mode; should omit aspect_ratio and resolution from body in extension mode; should warn about provider-level resolution in extension mode; full port in src/xai/tests/ +#[tokio::test] +async fn omits_unsupported_fields_from_edit_and_extension_bodies() { + let mut edit = xai(base(), json!({"mode": "edit-video", "videoUrl": SRC})); + edit.duration = Some(10); + edit.aspect_ratio = ar("16:9"); + edit.resolution = size("1280x720"); + let s = wiremock_start(V1, edit).await; + for key in ["duration", "aspect_ratio", "resolution"] { + assert!(s.body.get(key).is_none(), "edit body has {key}"); + } + + let mut extend = xai(base(), json!({"mode": "extend-video", "videoUrl": SRC})); + extend.duration = Some(6); + extend.aspect_ratio = ar("16:9"); + extend.resolution = size("1280x720"); + let s = wiremock_start(V1, extend).await; + assert_eq!(s.body["duration"], 6); + assert!(s.body.get("aspect_ratio").is_none()); + assert!(s.body.get("resolution").is_none()); + + let provider_resolution = xai( + base(), + json!({"mode": "extend-video", "videoUrl": SRC, "resolution": "720p"}), + ); + let s = wiremock_start(V1, provider_resolution).await; + assert!(s.body.get("resolution").is_none()); + assert!(warned(&s.start.warnings, "resolution")); +} + +/// Checks reference_images for explicit and implicit reference-to-video. partial: should send reference_images array to /videos/generations with explicit mode; should fallback to reference-to-video mode when referenceImageUrls is set without mode; full port in src/xai/tests/ +#[tokio::test] +async fn sends_reference_images_for_reference_to_video() { + let urls = [ + "https://example.com/ref1.jpg", + "https://example.com/ref2.jpg", + ]; + let explicit = xai( + base(), + json!({"mode": "reference-to-video", "referenceImageUrls": urls}), + ); + let implicit = xai(base(), json!({"referenceImageUrls": urls})); + for options in [explicit, implicit] { + let s = wiremock_start(V1, options).await; + assert_eq!(s.method, "POST"); + assert_eq!(s.path, "/videos/generations"); + assert_eq!( + s.body["reference_images"], + json!([{"url": urls[0]}, {"url": urls[1]}]) + ); + } +} + +/// Checks that 1080p is downgraded to 720p in reference-to-video with a warning. partial: should downgrade R2V 1080p to 720p with a warning; should downgrade R2V 1920x1080 from the SDK resolution to 720p; full port in src/xai/tests/ +#[tokio::test] +async fn downgrades_reference_to_video_1080p() { + let refs = json!(["https://example.com/ref1.jpg"]); + let provider = xai( + base(), + json!({"mode": "reference-to-video", "referenceImageUrls": refs, "resolution": "1080p"}), + ); + let mut sdk = xai( + base(), + json!({"mode": "reference-to-video", "referenceImageUrls": refs}), + ); + sdk.resolution = size("1920x1080"); + for options in [provider, sdk] { + let s = wiremock_start(V15, options).await; + assert_eq!(s.body["resolution"], "720p"); + assert!(warned(&s.start.warnings, "resolution")); + } +} + +/// Checks that inputReferences split into reference_images and reference_audios by media type. partial: should separate image and audio inputReferences; should support audio-only reference-to-video; should not send an empty reference_images array for video-only inputReferences; should combine a pinned first frame with an audio reference; should support explicit audio-only R2V without reference_images; full port in src/xai/tests/ +#[tokio::test] +async fn splits_input_references_by_media_type() { + let image = url_file("https://example.com/ref1.jpg", None); + let audio = url_file("https://example.com/voice.mp3", Some("audio/mpeg")); + let video = url_file("https://example.com/clip.mp4", Some("video/mp4")); + let start_image = url_file("https://example.com/start.jpg", Some("image/jpeg")); + let voice = json!([{"url": "https://example.com/voice.mp3"}]); + + let mut o = base(); + o.input_references = Some(vec![image, audio.clone()]); + let s = wiremock_start(V15, o).await; + assert_eq!( + s.body["reference_images"], + json!([{"url": "https://example.com/ref1.jpg"}]) + ); + assert_eq!(s.body["reference_audios"], voice); + assert!(s.start.warnings.is_empty()); + + let mut o = base(); + o.input_references = Some(vec![audio.clone()]); + let s = wiremock_start(V15, o).await; + assert_eq!(s.body["reference_audios"], voice); + assert!(s.body.get("reference_images").is_none()); + assert!(s.start.warnings.is_empty()); + + let mut o = base(); + o.input_references = Some(vec![video]); + let s = wiremock_start(V15, o).await; + assert!(s.body.get("reference_images").is_none()); + + let mut o = base(); + o.image = Some(start_image); + o.input_references = Some(vec![audio.clone()]); + let s = wiremock_start(V15, o).await; + assert_eq!( + s.body["image"], + json!({"url": "https://example.com/start.jpg"}) + ); + assert_eq!(s.body["reference_audios"], voice); + assert!(s.body.get("reference_images").is_none()); + assert!(s.start.warnings.is_empty()); + + let mut o = xai(base(), json!({"mode": "reference-to-video"})); + o.input_references = Some(vec![audio]); + let s = wiremock_start(V15, o).await; + assert_eq!(s.body["reference_audios"], voice); + assert!(s.body.get("reference_images").is_none()); + assert!(s.start.warnings.is_empty()); +} + +/// Checks that referenceVoiceIds become reference_audios in order and are never sent raw. partial: should send reference_audios for preset reference voices; should send up to 3 preset reference voices in order; should never forward referenceVoiceIds as a raw body key; full port in src/xai/tests/ +#[tokio::test] +async fn sends_reference_voice_ids_as_reference_audios() { + let cases = [ + (json!(["eve"]), json!([{"voice_id": "eve"}])), + ( + json!(["eve", "leo", "rex"]), + json!([{"voice_id": "eve"}, {"voice_id": "leo"}, {"voice_id": "rex"}]), + ), + ]; + for (ids, expected) in cases { + let o = xai( + base(), + json!({ + "mode": "reference-to-video", + "referenceImageUrls": ["https://example.com/ref1.jpg"], + "referenceVoiceIds": ids + }), + ); + let s = wiremock_start(V15, o).await; + assert_eq!(s.path, "/videos/generations"); + assert_eq!( + s.body["reference_images"], + json!([{"url": "https://example.com/ref1.jpg"}]) + ); + assert_eq!(s.body["reference_audios"], expected); + assert!(s.body.get("referenceVoiceIds").is_none()); + } +} + +/// Checks that more than 3 voices, or an empty voice id, gives InvalidArgument. partial: should reject more than 3 preset reference voices; should reject empty-string preset reference voice ids; full port in src/xai/tests/ +#[tokio::test] +async fn rejects_invalid_reference_voice_ids() { + for ids in [json!(["ara", "eve", "leo", "rex"]), json!([""])] { + let o = xai( + base(), + json!({ + "mode": "reference-to-video", + "referenceImageUrls": ["https://example.com/ref1.jpg"], + "referenceVoiceIds": ids + }), + ); + match try_start(V15, &o).await { + Err(AiMuxError::InvalidArgument(_)) => {} + Err(other) => panic!("unexpected error: {other:?}"), + Ok(_) => panic!("expected InvalidArgument for {ids}"), + } + } +} + +/// Checks that the requestId is encoded as one URL path segment, dot segments included. partial: should encode the request ID as a single URL path segment; should preserve the $requestId request ID as a URL path segment; full port in src/xai/tests/ +#[tokio::test] +async fn do_status_encodes_request_id_as_one_path_segment() { + let cases = [ + ("abc/../../internal", "/videos/abc%2F..%2F..%2Finternal"), + (".", "/videos/%252E"), + ("..", "/videos/%252E%252E"), + ]; + for (request_id, expected) in cases { + let (_, requests) = poll_once(json_response(done_status()), request_id).await; + assert_eq!(requests[0].method.as_str(), "GET"); + assert_eq!(requests[0].url.path(), expected); + } +} + +/// Checks the completed video, its xai metadata and the response. partial: should return completed with video data when done; should include response metadata; full port in src/xai/tests/ +#[tokio::test] +async fn do_status_returns_completed_video_and_metadata() { + let (result, _) = poll_once(json_response(done_status()), "req-123").await; + let done = wiremock_completed(result); + assert_eq!(done.videos.len(), 1); + match &done.videos[0] { + VideoData::Url { url, media_type } => { + assert_eq!(url, VIDEO_URL); + assert_eq!(media_type, "video/mp4"); + } + other => panic!("expected url, got {other:?}"), + } + let meta = xai_meta(&done); + assert_eq!(meta.len(), 4); + assert_eq!(meta["requestId"], "req-123"); + assert_eq!(meta["videoUrl"], VIDEO_URL); + assert_eq!(meta["duration"].as_f64(), Some(5.0)); + assert_eq!(meta["progress"].as_f64(), Some(100.0)); + assert_eq!(done.response.model_id.as_deref(), Some(V1)); + assert!(done.response.timestamp.is_some()); + assert!(done.response.headers.is_some()); +} + +/// Checks that expired, failed, missing-URL and blocked statuses give non-retryable errors. partial: should return error status on expired; should return error status on failed; should report an error status when video URL missing on done; should report an error status when respect_moderation is false; full port in src/xai/tests/ +#[tokio::test] +async fn do_status_reports_terminal_failures_as_non_retryable_errors() { + let cases = [ + ( + json!({"status": "expired", "model": V1}), + "Video generation request expired.", + ), + ( + json!({"status": "failed", "model": V1, "progress": 0, + "error": {"message": "Content policy violation"}}), + "Video generation failed: Content policy violation", + ), + ( + json!({"status": "done", "video": null, "model": V1}), + "Video generation completed but no video URL was returned.", + ), + ( + json!({"status": "done", "video": {"url": "", "respect_moderation": false}, "model": V1}), + "Video generation was blocked due to a content policy violation.", + ), + ]; + for (body, message) in cases { + let (result, _) = poll_once(json_response(body), "req-123").await; + assert_terminal(result, message); + } +} + +/// Checks that costInUsdTicks is added and that duration and progress are omitted when absent. partial: should include costInUsdTicks when returned in usage; should omit duration from metadata when absent in response; should omit progress from metadata when absent in response; full port in src/xai/tests/ +#[tokio::test] +async fn do_status_metadata_includes_optional_fields_only_when_present() { + let mut with_usage = done_status(); + with_usage["usage"] = json!({"cost_in_usd_ticks": 4_000_000_000_u64}); + let mut no_duration = done_status(); + no_duration["video"] + .as_object_mut() + .unwrap() + .remove("duration"); + let mut no_progress = done_status(); + no_progress.as_object_mut().unwrap().remove("progress"); + + let meta = xai_meta(&wiremock_completed( + poll_once(json_response(with_usage), "req-123").await.0, + )); + assert_eq!(meta.len(), 5); + assert_eq!(meta["costInUsdTicks"].as_f64(), Some(4_000_000_000.0)); + let meta = xai_meta(&wiremock_completed( + poll_once(json_response(no_duration), "req-123").await.0, + )); + assert!(!meta.contains_key("duration")); + let meta = xai_meta(&wiremock_completed( + poll_once(json_response(no_progress), "req-123").await.0, + )); + assert!(!meta.contains_key("progress")); +} diff --git a/aimux-providers/src/xai/transcription.rs b/aimux-providers/src/xai/transcription_model.rs similarity index 99% rename from aimux-providers/src/xai/transcription.rs rename to aimux-providers/src/xai/transcription_model.rs index 40263e443..e734338b7 100644 --- a/aimux-providers/src/xai/transcription.rs +++ b/aimux-providers/src/xai/transcription_model.rs @@ -175,7 +175,7 @@ impl TranscriptionModel for XaiTranscriptionModel { exchange.request(exchange.url("/stt"), options), HttpBody::Bytes(body, content_type), aimux_provider_utils::create_json_response_handler::(), - super::xai_failed_response_handler(), + super::error::xai_failed_response_handler(), ) .await?; let response = resp.value; diff --git a/aimux-providers/src/xai/video.rs b/aimux-providers/src/xai/video_model.rs similarity index 99% rename from aimux-providers/src/xai/video.rs rename to aimux-providers/src/xai/video_model.rs index b62feaddb..2fc56f6f3 100644 --- a/aimux-providers/src/xai/video.rs +++ b/aimux-providers/src/xai/video_model.rs @@ -706,7 +706,7 @@ impl VideoModel for XaiVideoModel { exchange.request(exchange.url(path), options), body, aimux_provider_utils::create_json_response_handler::(), - super::xai_failed_response_handler(), + super::error::xai_failed_response_handler(), ) .await?; let request_id = resp @@ -742,7 +742,7 @@ impl VideoModel for XaiVideoModel { let resp = aimux_provider_utils::get_from_api( exchange.request(url.clone(), options), status_response_handler(), - super::xai_failed_response_handler(), + super::error::xai_failed_response_handler(), ) .await?; let headers = resp.response_headers; diff --git a/aimux-providers/tests/xai_speech_test.rs b/aimux-providers/tests/xai_speech_test.rs deleted file mode 100644 index dc1b18a55..000000000 --- a/aimux-providers/tests/xai_speech_test.rs +++ /dev/null @@ -1,435 +0,0 @@ -//! Rust translation of the xAI speech (TTS) model tests. -//! -//! Source: `reference/aisdk-pinned/xai/src/xai-speech-model.test.ts` -//! -//! Each test starts a `wiremock` server, mounts a binary audio response (or -//! the `with_timestamps` JSON envelope) on `POST /tts`, calls `do_generate` -//! and asserts on the request the server saw and on the result. -//! -//! The upstream cases this file does not cover, or covers only in part, are -//! ported in full in `src/xai/tests/speech_model.rs`. - -use std::collections::HashMap; - -use serde_json::{Value, json}; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::error::AiMuxError; -use aimux_core::shared::{SharedProviderOptions, Warning}; -use aimux_core::speech_model::{AudioData, SpeechCallOptions, SpeechModel, SpeechResult}; -use aimux_providers::{XAIProviderSettings, XaiSpeechModel, create_xai}; - -const AUDIO: [u8; 4] = [1, 2, 3, 4]; - -fn speech_model( - server: &MockServer, - headers: Option>>, -) -> XaiSpeechModel { - create_xai(XAIProviderSettings { - api_key: Some("test-key".into()), - base_url: Some(server.uri()), - headers, - ..Default::default() - }) - .expect("valid settings") - .speech() -} - -/// Start a server answering `POST /tts` with `response`. -async fn server_with(response: ResponseTemplate) -> MockServer { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/tts")) - .respond_with(response) - .mount(&server) - .await; - server -} - -fn audio_response(content_type: &str, headers: &[(&str, &str)]) -> ResponseTemplate { - headers.iter().fold( - ResponseTemplate::new(200) - .insert_header("content-type", content_type) - .set_body_bytes(AUDIO.to_vec()), - |response, (name, value)| response.insert_header(*name, *value), - ) -} - -/// The `with_timestamps` JSON envelope; 'AQIDBA==' is `AUDIO` in base64. -fn timestamps_response(headers: &[(&str, &str)]) -> ResponseTemplate { - headers.iter().fold( - ResponseTemplate::new(200).set_body_json(json!({ - "audio": "AQIDBA==", - "content_type": "audio/mpeg", - "duration": 1.19, - "audio_timestamps": { - "graph_chars": ["H", "i"], - "graph_times": [[0.04, 0.06], [0.06, 0.1]], - }, - })), - |response, (name, value)| response.insert_header(*name, *value), - ) -} - -fn options(text: &str) -> SpeechCallOptions { - SpeechCallOptions::new(text) -} - -/// `providerOptions: { xai: {...} }`. -fn with_xai(mut options: SpeechCallOptions, xai: Value) -> SpeechCallOptions { - options.provider_options = Some(SharedProviderOptions::from([( - "xai".to_string(), - xai.as_object().expect("an object").clone(), - )])); - options -} - -/// The JSON body of the only request the server received. -async fn request_body(server: &MockServer) -> Value { - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - serde_json::from_slice(&requests[0].body).unwrap() -} - -fn xai_metadata(result: &SpeechResult) -> Value { - Value::Object(result.provider_metadata.as_ref().expect("metadata")["xai"].clone()) -} - -fn is_unsupported(warnings: &[Warning], expected: &str) -> bool { - warnings - .iter() - .any(|w| matches!(w, Warning::Unsupported { feature, .. } if feature == expected)) -} - -mod do_generate { - use super::*; - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should send text with xAI defaults - #[tokio::test] - async fn should_send_text_with_xai_defaults() { - let server = server_with(audio_response("audio/mpeg", &[])).await; - let result = speech_model(&server, None) - .do_generate(&options("Hello from the AI SDK!")) - .await - .unwrap(); - - let expected = json!({ - "text": "Hello from the AI SDK!", - "voice_id": "eve", - "language": "auto", - "output_format": { "codec": "mp3" }, - }); - assert_eq!(request_body(&server).await, expected); - assert_eq!(result.request.unwrap().body, Some(expected)); - let requests = server.received_requests().await.unwrap(); - assert_eq!(requests[0].method.as_str(), "POST"); - assert_eq!(requests[0].url.path(), "/tts"); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should pass standard speech options - #[tokio::test] - async fn should_pass_standard_speech_options() { - let server = server_with(audio_response("audio/wav", &[])).await; - let mut opts = options("Hello from the AI SDK!"); - opts.voice = Some("ara".into()); - opts.language = Some("en".into()); - opts.output_format = Some("wav".into()); - opts.speed = Some(1.2); - speech_model(&server, None) - .do_generate(&opts) - .await - .unwrap(); - - assert_eq!( - request_body(&server).await, - json!({ - "text": "Hello from the AI SDK!", - "voice_id": "ara", - "language": "en", - "output_format": { "codec": "wav" }, - "speed": 1.2, - }) - ); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should accept the %s output format - #[tokio::test] - async fn should_accept_each_output_format() { - for format in ["mp3", "wav", "pcm", "mulaw", "alaw"] { - let server = server_with(audio_response("audio/mpeg", &[])).await; - let mut opts = options("Hello from the AI SDK!"); - opts.output_format = Some(format.into()); - let result = speech_model(&server, None) - .do_generate(&opts) - .await - .unwrap(); - - assert_eq!( - request_body(&server).await["output_format"]["codec"], - format - ); - assert!(result.warnings.is_empty(), "{format}"); - } - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should map provider options onto xAI request fields - #[tokio::test] - async fn should_map_provider_options_onto_xai_request_fields() { - let server = server_with(audio_response("audio/mpeg", &[])).await; - let opts = with_xai( - options("Hello from the AI SDK!"), - json!({ - "sampleRate": 44100, - "bitRate": 192000, - "optimizeStreamingLatency": 1, - "textNormalization": true, - }), - ); - speech_model(&server, None) - .do_generate(&opts) - .await - .unwrap(); - - let body = request_body(&server).await; - assert_eq!( - body["output_format"], - json!({ "codec": "mp3", "sample_rate": 44100, "bit_rate": 192000 }) - ); - assert_eq!(body["optimize_streaming_latency"], 1); - assert_eq!(body["text_normalization"], true); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should pass withTimestamps and replace provider options - #[tokio::test] - async fn should_pass_with_timestamps_and_replace_provider_options() { - let server = server_with(timestamps_response(&[])).await; - let opts = with_xai( - options("Hello from the AI SDK!"), - json!({ "withTimestamps": true, "replace": { "nginx": "/ˈɛndʒɪn ˈɛks/" } }), - ); - speech_model(&server, None) - .do_generate(&opts) - .await - .unwrap(); - - let body = request_body(&server).await; - assert_eq!(body["with_timestamps"], true); - assert_eq!(body["replace"], json!({ "nginx": "/ˈɛndʒɪn ˈɛks/" })); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should decode the with_timestamps envelope and extract provider metadata - #[tokio::test] - async fn should_decode_the_with_timestamps_envelope_and_extract_provider_metadata() { - let trace = "993675dc-8ea6-4f54-b4ad-a59ac2615026"; - let server = server_with(timestamps_response(&[("x-trace-id", trace)])).await; - let opts = with_xai(options("Hi"), json!({ "withTimestamps": true })); - let result = speech_model(&server, None) - .do_generate(&opts) - .await - .unwrap(); - - assert!(matches!(&result.audio, AudioData::Binary(audio) if audio == &AUDIO)); - assert_eq!( - xai_metadata(&result), - json!({ - "traceId": trace, - "duration": 1.19, - "contentType": "audio/mpeg", - "audioTimestamps": { - "graphChars": ["H", "i"], - "graphTimes": [[0.04, 0.06], [0.06, 0.1]], - }, - }) - ); - assert_eq!(result.response.body.unwrap()["audio"], "AQIDBA=="); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should extract the trace id from binary responses - #[tokio::test] - async fn should_extract_the_trace_id_from_binary_responses() { - let trace = "06e3dab5-e3ba-4c6b-83a6-1e9ea11d78af"; - let server = server_with(audio_response("audio/mpeg", &[("x-trace-id", trace)])).await; - let result = speech_model(&server, None) - .do_generate(&options("Hello from the AI SDK!")) - .await - .unwrap(); - - assert_eq!(xai_metadata(&result), json!({ "traceId": trace })); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should return empty provider metadata when xAI headers are absent - #[tokio::test] - async fn should_return_empty_provider_metadata_when_xai_headers_are_absent() { - let server = server_with(audio_response("audio/mpeg", &[])).await; - let result = speech_model(&server, None) - .do_generate(&options("Hello from the AI SDK!")) - .await - .unwrap(); - - assert_eq!(xai_metadata(&result), json!({})); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should warn and use mp3 for unsupported output formats - #[tokio::test] - async fn should_warn_and_use_mp3_for_unsupported_output_formats() { - let server = server_with(audio_response("audio/mpeg", &[])).await; - let mut opts = options("Hello from the AI SDK!"); - opts.output_format = Some("flac".into()); - let result = speech_model(&server, None) - .do_generate(&opts) - .await - .unwrap(); - - assert_eq!(request_body(&server).await["output_format"]["codec"], "mp3"); - assert!(is_unsupported(&result.warnings, "outputFormat")); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should warn and ignore bitRate for non-mp3 output - #[tokio::test] - async fn should_warn_and_ignore_bit_rate_for_non_mp3_output() { - let server = server_with(audio_response("audio/wav", &[])).await; - let mut opts = with_xai( - options("Hello from the AI SDK!"), - json!({ "bitRate": 192000 }), - ); - opts.output_format = Some("wav".into()); - let result = speech_model(&server, None) - .do_generate(&opts) - .await - .unwrap(); - - assert_eq!( - request_body(&server).await["output_format"], - json!({ "codec": "wav" }) - ); - assert!(is_unsupported(&result.warnings, "providerOptions")); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should warn when instructions are provided - #[tokio::test] - async fn should_warn_when_instructions_are_provided() { - let server = server_with(audio_response("audio/mpeg", &[])).await; - let mut opts = options("Hello from the AI SDK!"); - opts.instructions = Some("Speak cheerfully".into()); - let result = speech_model(&server, None) - .do_generate(&opts) - .await - .unwrap(); - - assert!(is_unsupported(&result.warnings, "instructions")); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should pass headers and the xAI user agent - #[tokio::test] - async fn should_pass_headers_and_the_xai_user_agent() { - let server = server_with(audio_response("audio/mpeg", &[])).await; - let provider_headers = HashMap::from([( - "Custom-Provider-Header".to_string(), - Some("provider-header-value".to_string()), - )]); - let mut opts = options("Hello from the AI SDK!"); - opts.headers = Some(HashMap::from([( - "Custom-Request-Header".to_string(), - "request-header-value".to_string(), - )])); - speech_model(&server, Some(provider_headers)) - .do_generate(&opts) - .await - .unwrap(); - - let requests = server.received_requests().await.unwrap(); - let header = |name: &str| requests[0].headers.get(name).unwrap().to_str().unwrap(); - assert_eq!(header("authorization"), "Bearer test-key"); - assert_eq!(header("content-type"), "application/json"); - assert_eq!(header("custom-provider-header"), "provider-header-value"); - assert_eq!(header("custom-request-header"), "request-header-value"); - assert!(header("user-agent").contains("ai-sdk-xai/5.0.12")); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should return binary audio data - #[tokio::test] - async fn should_return_binary_audio_data() { - let server = server_with(audio_response("audio/mpeg", &[])).await; - let result = speech_model(&server, None) - .do_generate(&options("Hello from the AI SDK!")) - .await - .unwrap(); - - assert!(matches!(&result.audio, AudioData::Binary(audio) if audio == &AUDIO)); - assert!(result.warnings.is_empty()); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should handle API errors - #[tokio::test] - async fn should_handle_api_errors() { - let server = server_with( - ResponseTemplate::new(400).set_body_string( - json!({ "error": { "message": "Invalid text", "type": "invalid_request_error" } }) - .to_string(), - ), - ) - .await; - let error = speech_model(&server, None) - .do_generate(&options("Hello from the AI SDK!")) - .await - .unwrap_err(); - - let AiMuxError::ApiCall(error) = error else { - panic!("expected an ApiCall error, got {error:?}"); - }; - assert_eq!(error.message, "Invalid text"); - assert_eq!(error.status_code, Some(400)); - } - - /// upstream: @ai-sdk/xai src/xai-speech-model.test.ts › doGenerate › should use the real date when no custom date provider is specified - #[tokio::test] - async fn should_use_the_real_date_when_no_custom_date_provider_is_specified() { - let server = server_with(audio_response("audio/mpeg", &[])).await; - let before = chrono::Utc::now().timestamp_millis(); - let result = speech_model(&server, None) - .do_generate(&options("Hello from the AI SDK!")) - .await - .unwrap(); - let after = chrono::Utc::now().timestamp_millis(); - - // A `Date`: millisecond precision. - let timestamp = result.response.timestamp.unwrap().unix_millis(); - assert!(before <= timestamp && timestamp <= after); - assert_eq!(result.response.model_id.as_deref(), Some("")); - } -} - -/// Checks the provider id and the empty model id; specificationVersion is not checked. partial: should expose correct provider and model information; full port in src/xai/tests/ -#[tokio::test] -async fn should_expose_correct_provider_and_model_information() { - let model = speech_model(&MockServer::start().await, None); - assert_eq!(model.provider(), "xai.speech"); - assert_eq!(model.model_id(), ""); -} - -/// Checks the response model id and the x-request-id header; the timestamp is only checked as present. partial: should include response timestamp, model id, and headers; full port in src/xai/tests/ -#[tokio::test] -async fn should_include_response_timestamp_model_id_and_headers() { - let server = server_with(audio_response( - "audio/mpeg", - &[("x-request-id", "test-request-id")], - )) - .await; - let result = speech_model(&server, None) - .do_generate(&options("Hello from the AI SDK!")) - .await - .unwrap(); - - assert_eq!(result.response.model_id.as_deref(), Some("")); - assert_eq!( - result - .response - .headers - .unwrap() - .get("x-request-id") - .map(String::as_str), - Some("test-request-id") - ); - assert!(result.response.timestamp.is_some()); -} diff --git a/aimux-providers/tests/xai_transcription_test.rs b/aimux-providers/tests/xai_transcription_test.rs deleted file mode 100644 index 9e9ef2934..000000000 --- a/aimux-providers/tests/xai_transcription_test.rs +++ /dev/null @@ -1,828 +0,0 @@ -//! Rust translation of the xAI transcription model tests. -//! Source: `reference/aisdk-pinned/xai/src/xai-transcription-model.test.ts` -//! -//! `do_generate` runs against a wiremock server (`POST {base}/stt`). -//! `do_stream` runs against a local WebSocket server that plays the xAI STT -//! role: it captures the handshake (URI + headers), sends `transcript.created`, -//! records the binary audio frames and the final `audio.done` text frame, then -//! sends the scripted events. -//! -//! The upstream cases this file does not cover, or covers only in part, are -//! ported in full in `src/xai/tests/transcription_model.rs`. - -use std::collections::HashMap; -use std::sync::{Arc, Mutex}; -use std::time::Duration; - -use futures::{SinkExt, StreamExt}; -use serde_json::{Value, json}; -use tokio::net::TcpListener; -use tokio_tungstenite::tungstenite::Message; -use tokio_tungstenite::tungstenite::handshake::server::{Request, Response}; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::error::AiMuxError; -use aimux_core::shared::{SharedProviderOptions, Warning}; -use aimux_core::transcription_model::{ - AudioChunk, AudioInput, InputAudioFormat, TranscriptionCallOptions, TranscriptionModel, - TranscriptionResponse, TranscriptionResult, TranscriptionStreamOptions, - TranscriptionStreamPart, -}; -use aimux_providers::{XAIProviderSettings, XaiTranscriptionModel, create_xai}; - -type Headers = HashMap>; - -fn model(base_url: &str, headers: Option) -> XaiTranscriptionModel { - create_xai(XAIProviderSettings { - api_key: Some("test-api-key".into()), - base_url: Some(base_url.into()), - headers, - ..Default::default() - }) - .expect("valid settings") - .transcription() -} - -fn xai_options(value: Value) -> SharedProviderOptions { - SharedProviderOptions::from([("xai".to_string(), value.as_object().unwrap().clone())]) -} - -// ── doGenerate helpers ────────────────────────────────────────────────────── - -fn call(provider_options: Option, media_type: &str) -> TranscriptionCallOptions { - let mut options = - TranscriptionCallOptions::new(AudioInput::Binary(vec![1, 2, 3, 4]), media_type); - options.provider_options = provider_options.map(xai_options); - options -} - -fn json_response(body: Value, headers: &[(&str, &str)]) -> ResponseTemplate { - headers.iter().fold( - ResponseTemplate::new(200).set_body_json(body), - |t, (k, v)| t.insert_header(*k, *v), - ) -} - -async fn mock_stt(server: &MockServer, template: ResponseTemplate) { - Mock::given(method("POST")) - .and(path("/stt")) - .respond_with(template) - .mount(server) - .await; -} - -fn default_body() -> Value { - json!({ - "text": "Hello from the AI SDK!", - "language": "en", - "duration": 2.5, - "words": [ - {"text": "Hello", "start": 0, "end": 1}, - {"text": "from the AI SDK!", "start": 1, "end": 2.5} - ] - }) -} - -async fn generate( - server: &MockServer, - options: &TranscriptionCallOptions, - headers: Option, -) -> TranscriptionResult { - model(&server.uri(), headers) - .do_generate(options) - .await - .expect("do_generate") -} - -struct Part { - name: String, - filename: Option, - content_type: Option, - value: String, -} - -fn quoted(line: &str, key: &str) -> Option { - let start = line.find(&format!(" {key}=\""))? + key.len() + 3; - let end = line[start..].find('"')?; - Some(line[start..start + end].to_string()) -} - -/// The multipart fields of a captured request, in wire order. -fn parse_multipart(request: &wiremock::Request) -> Vec { - let content_type = request - .headers - .get("content-type") - .unwrap() - .to_str() - .unwrap(); - let boundary = content_type.split("boundary=").nth(1).unwrap(); - let body = String::from_utf8_lossy(&request.body).into_owned(); - body.split(&format!("--{boundary}")) - .skip(1) - .map(|chunk| chunk.trim_start_matches("\r\n")) - .filter(|chunk| !chunk.starts_with("--")) - .map(|chunk| { - let (head, value) = chunk.split_once("\r\n\r\n").unwrap(); - let line = |prefix: &str| head.lines().find(|l| l.starts_with(prefix)); - Part { - name: quoted(line("Content-Disposition").unwrap(), "name").unwrap(), - filename: quoted(line("Content-Disposition").unwrap(), "filename"), - content_type: line("Content-Type").map(|l| l["Content-Type: ".len()..].to_string()), - value: value.strip_suffix("\r\n").unwrap_or(value).to_string(), - } - }) - .collect() -} - -fn full_options() -> Value { - json!({ - "audioFormat": "pcm", "sampleRate": 16000, "language": "en", "format": true, - "multichannel": true, "channels": 2, "diarize": true, - "keyterm": ["AI SDK", "Grok"], "fillerWords": true - }) -} - -// ── doStream helpers ──────────────────────────────────────────────────────── - -#[derive(Default, Clone)] -struct Capture { - uri: String, - headers: HashMap, - /// Frames that arrived before the server sent `transcript.created`. - early_frames: usize, - binary: Vec>, - texts: Vec, -} - -async fn serve(listener: TcpListener, cap: Arc>, events: Vec) { - let (stream, _) = listener.accept().await.unwrap(); - let handshake_cap = Arc::clone(&cap); - #[allow(clippy::result_large_err)] - let handshake = move |req: &Request, resp: Response| { - let mut cap = handshake_cap.lock().unwrap(); - cap.uri = req.uri().to_string(); - cap.headers = req - .headers() - .iter() - .filter_map(|(k, v)| Some((k.as_str().to_string(), v.to_str().ok()?.to_string()))) - .collect(); - Ok(resp) - }; - let mut ws = tokio_tungstenite::accept_hdr_async(stream, handshake) - .await - .expect("ws handshake"); - - // The client holds audio until `transcript.created`. - if let Ok(Some(Ok(_))) = tokio::time::timeout(Duration::from_millis(100), ws.next()).await { - cap.lock().unwrap().early_frames += 1; - } - ws.send(Message::Text( - json!({"type": "transcript.created"}).to_string(), - )) - .await - .unwrap(); - loop { - match ws.next().await { - Some(Ok(Message::Binary(bytes))) => cap.lock().unwrap().binary.push(bytes.to_vec()), - Some(Ok(Message::Text(text))) => { - let text = text.as_str().to_string(); - cap.lock().unwrap().texts.push(text.clone()); - if text.contains("audio.done") { - break; - } - } - Some(Ok(_)) => {} - _ => return, - } - } - for event in events { - ws.send(Message::Text(event.to_string())).await.unwrap(); - } - // Wait for the client to close. - let _ = tokio::time::timeout(Duration::from_secs(2), async { - while let Some(Ok(message)) = ws.next().await { - if message.is_close() { - break; - } - } - }) - .await; -} - -async fn start(events: Vec) -> (String, Arc>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let cap = Arc::new(Mutex::new(Capture::default())); - tokio::spawn(serve(listener, Arc::clone(&cap), events)); - (base_url, cap) -} - -fn stream_options( - provider_options: Option, - format_type: &str, - rate: Option, -) -> TranscriptionStreamOptions { - TranscriptionStreamOptions { - audio: Box::pin(futures::stream::iter(vec![AudioChunk::Binary(vec![ - 1, 2, 3, - ])])), - input_audio_format: InputAudioFormat { - format_type: format_type.to_string(), - rate, - }, - provider_options: provider_options.map(xai_options), - abort_signal: None, - headers: None, - include_raw_chunks: false, - timeout: None, - } -} - -struct Run { - cap: Capture, - parts: Vec>, - response: Option, -} - -impl Run { - /// The parts as JSON; panics on a stream error. - fn values(&self) -> Vec { - self.parts - .iter() - .map(|p| serde_json::to_value(p.as_ref().expect("stream part")).unwrap()) - .collect() - } -} - -async fn run_with( - options: TranscriptionStreamOptions, - headers: Option, - events: Vec, -) -> Run { - let (base_url, cap) = start(events).await; - let mut result = model(&base_url, headers) - .do_stream(options) - .expect("the model has do_stream") - .await - .expect("do_stream"); - let response = result.response.take(); - let parts = tokio::time::timeout(Duration::from_secs(10), result.stream.collect::>()) - .await - .expect("stream timed out"); - let cap = cap.lock().unwrap().clone(); - Run { - cap, - parts, - response, - } -} - -async fn run(provider_options: Option, events: Vec) -> Run { - run_with( - stream_options(provider_options, "audio/pcm", Some(16000)), - None, - events, - ) - .await -} - -fn partial(text: &str, is_final: bool, speech_final: bool, span: Option<(f64, f64)>) -> Value { - let mut event = json!({ - "type": "transcript.partial", "text": text, - "is_final": is_final, "speech_final": speech_final - }); - if let Some((start, duration)) = span { - event["start"] = json!(start); - event["duration"] = json!(duration); - } - event -} - -fn done(text: &str, duration: f64) -> Value { - json!({"type": "transcript.done", "text": text, "duration": duration}) -} - -fn kinds(parts: &[Value]) -> Vec<&str> { - parts.iter().map(|p| p["type"].as_str().unwrap()).collect() -} - -fn of_kind(parts: &[Value], kind: &str) -> Vec { - parts - .iter() - .filter(|p| p["type"] == kind) - .cloned() - .collect() -} - -fn finish(text: &str, language: Value, duration: Option) -> Value { - let mut part = json!({"type": "finish", "text": text, "segments": []}); - if !language.is_null() { - part["language"] = language; - } - if let Some(duration) = duration { - part["durationInSeconds"] = json!(duration); - } - part -} - -mod do_generate { - use super::*; - - /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doGenerate › should send a multipart request with the audio file - #[tokio::test] - async fn should_send_a_multipart_request_with_the_audio_file() { - let server = MockServer::start().await; - mock_stt(&server, json_response(default_body(), &[])).await; - - generate(&server, &call(None, "audio/wav"), None).await; - - let requests = server.received_requests().await.unwrap(); - assert_eq!(requests[0].method.as_str(), "POST"); - assert_eq!(requests[0].url.path(), "/stt"); - let parts = parse_multipart(&requests[0]); - let file = parts.iter().find(|p| p.name == "file").expect("file field"); - assert_eq!(file.filename.as_deref(), Some("audio.wav")); - assert_eq!(file.content_type.as_deref(), Some("audio/wav")); - } - - /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doGenerate › should map provider options onto xAI request fields - #[tokio::test] - async fn should_map_provider_options_onto_xai_request_fields() { - let server = MockServer::start().await; - mock_stt(&server, json_response(default_body(), &[])).await; - - generate(&server, &call(Some(full_options()), "audio/pcm"), None).await; - - let parts = parse_multipart(&server.received_requests().await.unwrap()[0]); - let values = |name: &str| -> Vec<&str> { - parts - .iter() - .filter(|p| p.name == name) - .map(|p| p.value.as_str()) - .collect() - }; - for (name, expected) in [ - ("audio_format", "pcm"), - ("sample_rate", "16000"), - ("language", "en"), - ("format", "true"), - ("multichannel", "true"), - ("channels", "2"), - ("diarize", "true"), - ("filler_words", "true"), - ] { - assert_eq!(values(name), vec![expected], "field {name}"); - } - assert_eq!(values("keyterm"), vec!["AI SDK", "Grok"]); - } - - /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doGenerate › should append file after all other multipart fields - #[tokio::test] - async fn should_append_file_after_all_other_multipart_fields() { - let server = MockServer::start().await; - mock_stt( - &server, - json_response(json!({"text": "Hello from the AI SDK!"}), &[]), - ) - .await; - - generate(&server, &call(Some(full_options()), "audio/pcm"), None).await; - - let parts = parse_multipart(&server.received_requests().await.unwrap()[0]); - let names: Vec<&str> = parts.iter().map(|p| p.name.as_str()).collect(); - assert_eq!( - names, - [ - "audio_format", - "sample_rate", - "language", - "format", - "multichannel", - "channels", - "diarize", - "filler_words", - "keyterm", - "keyterm", - "file" - ] - ); - } - - /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doGenerate › should extract text, segments, language, and duration - #[tokio::test] - async fn should_extract_text_segments_language_and_duration() { - let server = MockServer::start().await; - mock_stt(&server, json_response(default_body(), &[])).await; - - let result = generate(&server, &call(None, "audio/wav"), None).await; - - assert_eq!(result.text, "Hello from the AI SDK!"); - assert_eq!(result.language.as_deref(), Some("en")); - assert_eq!(result.duration_in_seconds, Some(2.5)); - let segments: Vec<_> = result - .segments - .iter() - .map(|s| (s.text.as_str(), s.start_second, s.end_second)) - .collect(); - assert_eq!( - segments, - [("Hello", 0.0, 1.0), ("from the AI SDK!", 1.0, 2.5)] - ); - assert!(result.warnings.is_empty()); - } - - /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doGenerate › should handle missing words, duration, and empty language - #[tokio::test] - async fn should_handle_missing_words_duration_and_empty_language() { - let server = MockServer::start().await; - mock_stt( - &server, - json_response( - json!({"text": "Hello from the AI SDK!", "language": ""}), - &[], - ), - ) - .await; - - let result = generate(&server, &call(None, "audio/wav"), None).await; - - assert_eq!(result.text, "Hello from the AI SDK!"); - assert_eq!(result.language, None); - assert_eq!(result.duration_in_seconds, None); - assert!(result.segments.is_empty()); - assert!(result.warnings.is_empty()); - } -} - -mod do_stream { - use super::*; - - /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doStream › should emit one transcript-final per utterance and reconstruct the finish text when transcript.done is empty - #[tokio::test] - async fn should_emit_one_final_per_utterance_and_reconstruct_finish_text() { - let events = vec![ - partial("Hello wor", false, false, Some((0.0, 0.8))), - partial("Hello world.", true, false, Some((0.0, 1.0))), - partial("Hello world.", true, true, Some((0.0, 1.0))), - done("", 1.0), - ]; - let parts = run(None, events).await.values(); - - assert_eq!( - of_kind(&parts, "transcript-final"), - vec![json!({ - "type": "transcript-final", "text": "Hello world.", "startSecond": 0.0, - "endSecond": 1.0 - })] - ); - // the speech_final:false re-send surfaces as a partial - assert_eq!(of_kind(&parts, "transcript-partial").len(), 2); - assert_eq!( - parts.last().unwrap(), - &finish("Hello world.", Value::Null, Some(1.0)) - ); - } - - /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doStream › should treat is_final fragments as partials and use the speech_final text for finish - #[tokio::test] - async fn should_treat_is_final_fragments_as_partials_and_use_speech_final_text() { - let events = vec![ - partial("No", true, false, Some((0.0, 0.5))), - partial(", I'm not", true, false, Some((0.5, 0.7))), - partial("No, I'm not.", true, true, Some((0.0, 1.2))), - done("", 1.2), - ]; - let parts = run(None, events).await.values(); - - assert_eq!( - kinds(&parts), - [ - "stream-start", - "transcript-partial", - "transcript-partial", - "transcript-final", - "finish" - ] - ); - assert_eq!(parts.last().unwrap()["text"], "No, I'm not."); - } - - /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doStream › should fall back to the latest pending text when no speech_final arrived before transcript.done - #[tokio::test] - async fn should_fall_back_to_latest_pending_text_without_speech_final() { - let events = vec![ - partial("Hello wor", false, false, None), - partial("Hello world", true, false, None), - done("", 1.0), - ]; - let parts = run(None, events).await.values(); - - let last = parts.last().unwrap(); - assert_eq!(kinds(&parts).last(), Some(&"finish")); - assert_eq!(last["text"], "Hello world"); - } - - /// upstream: @ai-sdk/xai src/xai-transcription-model.test.ts › doStream › should warn on unrecognized inputAudioFormat types - #[tokio::test] - async fn should_warn_on_unrecognized_input_audio_format_types() { - let options = stream_options(None, "audio/wav", None); - let run = run_with(options, None, vec![done("Hello", 1.0)]).await; - - match &run.parts[0] { - Ok(TranscriptionStreamPart::StreamStart { warnings }) => match warnings.as_slice() { - [Warning::Other { message }] => { - assert!( - message.contains("Unrecognized inputAudioFormat.type \"audio/wav\""), - "{message}" - ); - } - other => panic!("expected one Other warning, got {other:?}"), - }, - other => panic!("expected StreamStart, got {other:?}"), - } - } -} - -/// Checks the provider id and the empty model id; specificationVersion is not checked. partial: should expose correct provider and model information; full port in src/xai/tests/ -#[test] -fn should_expose_correct_provider_and_model_information() { - let model = model("https://api.x.ai/v1", None); - assert_eq!(model.provider(), "xai.transcription"); - assert_eq!(model.model_id(), ""); -} - -/// Checks authorization, custom headers and the multipart content type; the user agent is checked without its version and the boundary only by prefix. partial: should pass headers and the xAI user agent; full port in src/xai/tests/ -#[tokio::test] -async fn should_pass_headers_and_the_xai_user_agent() { - let server = MockServer::start().await; - mock_stt(&server, json_response(default_body(), &[])).await; - - let mut options = call(None, "audio/wav"); - options.headers = Some(HashMap::from([( - "Custom-Request-Header".to_string(), - "request-header-value".to_string(), - )])); - let provider_headers = HashMap::from([( - "Custom-Provider-Header".to_string(), - Some("provider-header-value".to_string()), - )]); - generate(&server, &options, Some(provider_headers)).await; - - let requests = server.received_requests().await.unwrap(); - let header = |name: &str| { - requests[0] - .headers - .get(name) - .unwrap() - .to_str() - .unwrap() - .to_string() - }; - assert_eq!(header("authorization"), "Bearer test-api-key"); - assert!(header("content-type").starts_with("multipart/form-data; boundary=")); - assert_eq!(header("custom-provider-header"), "provider-header-value"); - assert_eq!(header("custom-request-header"), "request-header-value"); - assert!(header("user-agent").contains("ai-sdk-xai/")); -} - -/// Checks the response model id, headers and body; the timestamp is only checked as present. partial: should include response timestamp, model id, and headers; full port in src/xai/tests/ -#[tokio::test] -async fn should_include_response_timestamp_model_id_and_headers() { - let server = MockServer::start().await; - let headers = [ - ("x-request-id", "test-request-id"), - ("x-ratelimit-remaining", "123"), - ]; - mock_stt(&server, json_response(default_body(), &headers)).await; - - let result = generate(&server, &call(None, "audio/wav"), None).await; - - assert!(result.response.timestamp.is_some()); - assert_eq!(result.response.model_id.as_deref(), Some("")); - let response_headers = result.response.headers.unwrap(); - assert_eq!(response_headers["content-type"], "application/json"); - assert_eq!(response_headers["x-request-id"], "test-request-id"); - assert_eq!(response_headers["x-ratelimit-remaining"], "123"); - assert_eq!(result.response.body, Some(default_body())); -} - -/// Checks the multichannel-without-channels error and that no socket is opened; the error's argument field is not checked. partial: should require channels when streaming multichannel audio; full port in src/xai/tests/ -#[tokio::test] -async fn should_require_channels_when_streaming_multichannel_audio() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - - let options = stream_options( - Some(json!({"multichannel": true})), - "audio/pcm", - Some(16000), - ); - let error = model(&base_url, None) - .do_stream(options) - .expect("the model has do_stream") - .await - .unwrap_err(); - - assert!( - matches!(&error, AiMuxError::InvalidArgument(message) if message.contains( - "providerOptions.xai.channels is required when providerOptions.xai.multichannel is true" - )), - "got {error:?}" - ); - // Nothing connected. - assert!( - tokio::time::timeout(Duration::from_millis(100), listener.accept()) - .await - .is_err() - ); -} - -/// Checks the streaming URL, handshake, audio frames and parts; the response timestamp is only checked as present. partial: should stream xAI STT over WebSocket; full port in src/xai/tests/ -#[tokio::test] -async fn should_stream_xai_stt_over_websocket() { - let options = json!({ - "language": "en", "diarize": true, "keyterm": ["AI SDK", "Grok"], - "streaming": { - "interimResults": true, "endpointing": 500, - "smartTurn": 0.7, "smartTurnTimeout": 3000 - } - }); - let events = vec![ - partial("Hel", false, false, Some((0.0, 0.5))), - partial("Hello", true, true, Some((0.0, 1.0))), - done("Hello", 1.0), - ]; - let run = run(Some(options), events).await; - - assert_eq!( - run.cap.uri, - "/stt?sample_rate=16000&encoding=pcm&language=en&diarize=true&interim_results=true&endpointing=500&smart_turn=0.7&smart_turn_timeout=3000&keyterm=AI+SDK&keyterm=Grok" - ); - assert_eq!(run.cap.headers["authorization"], "Bearer test-api-key"); - // Audio waits for `transcript.created`, then the frames and `audio.done` follow. - assert_eq!(run.cap.early_frames, 0); - assert_eq!(run.cap.binary, vec![vec![1u8, 2, 3]]); - assert_eq!(run.cap.texts.len(), 1); - assert_eq!( - serde_json::from_str::(&run.cap.texts[0]).unwrap(), - json!({"type": "audio.done"}) - ); - - assert_eq!( - run.values(), - vec![ - json!({"type": "stream-start", "warnings": []}), - json!({ - "type": "transcript-partial", "text": "Hel", "startSecond": 0.0, - "durationInSeconds": 0.5 - }), - json!({ - "type": "transcript-final", "text": "Hello", "startSecond": 0.0, - "endSecond": 1.0 - }), - finish("Hello", json!("en"), Some(1.0)), - ] - ); - let response = run.response.unwrap(); - assert_eq!(response.model_id.as_deref(), Some("")); - assert!(response.timestamp.is_some()); -} - -/// Checks that a provider header set to None is left off the WebSocket handshake (analogue: per-call headers cannot hold None). partial: should strip undefined header values before the WebSocket constructor; full port in src/xai/tests/ -#[tokio::test] -async fn should_not_send_removed_header_values_on_the_handshake() { - let headers = HashMap::from([ - ( - "Custom-Header".to_string(), - Some("custom-value".to_string()), - ), - ("X-Unset".to_string(), None), - ]); - let options = stream_options(None, "audio/pcm", Some(16000)); - let run = run_with(options, Some(headers), vec![done("Hello", 1.0)]).await; - - assert_eq!(run.cap.headers["authorization"], "Bearer test-api-key"); - assert_eq!(run.cap.headers["custom-header"], "custom-value"); - assert!(!run.cap.headers.contains_key("x-unset")); -} - -/// Checks that finalized utterances are joined into the finish text; the finish part type is not checked. partial: should join finalized utterances per channel when transcript.done is empty; full port in src/xai/tests/ -#[tokio::test] -async fn should_join_finalized_utterances_when_transcript_done_is_empty() { - let events = vec![ - partial("First utterance.", true, true, Some((0.0, 1.0))), - partial("Second utterance.", true, true, Some((1.0, 1.0))), - done("", 2.0), - ]; - let parts = run(None, events).await.values(); - - assert_eq!( - parts.last().unwrap()["text"], - "First utterance. Second utterance." - ); -} - -/// Checks that an error event fails the stream with the server message; the WebSocket close is not checked. partial: should error the stream with the server message on error events; full port in src/xai/tests/ -#[tokio::test] -async fn should_error_the_stream_with_the_server_message_on_error_events() { - let events = vec![json!({"type": "error", "message": "invalid sample_rate"})]; - let run = run(None, events).await; - - assert!(matches!( - run.parts[0], - Ok(TranscriptionStreamPart::StreamStart { .. }) - )); - assert_eq!(run.parts.len(), 2, "parts: {:?}", run.parts); - assert!( - matches!(&run.parts[1], Err(AiMuxError::Other(message)) if message == "invalid sample_rate"), - "got {:?}", - run.parts[1] - ); -} - -/// Only `transcribe` downloads an audio URL (layer 5's `AudioInput::Url`); -/// the model itself refuses one before any request, as every transcription -/// model does. -#[tokio::test] -async fn url_audio_is_refused_before_any_request() { - let server = MockServer::start().await; - let mut options = call(None, "audio/wav"); - options.audio = AudioInput::Url("https://cdn.test/a.wav".into()); - let error = model(&server.uri(), None) - .do_generate(&options) - .await - .unwrap_err(); - assert_eq!(error.to_string(), AudioInput::url_error().to_string()); - assert!(server.received_requests().await.unwrap().is_empty()); -} - -/// A refused connection errors the stream: upstream connects inside the -/// stream, so `doStream` resolves and the stream carries the failure. -#[tokio::test] -async fn a_connect_failure_errors_the_stream() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - drop(listener); - - let result = model(&base_url, None) - .do_stream(stream_options(None, "audio/pcm", Some(16000))) - .expect("the model has do_stream") - .await - .expect("do_stream resolves; the connect failure arrives on the stream"); - let parts = tokio::time::timeout(Duration::from_secs(10), result.stream.collect::>()) - .await - .expect("stream timed out"); - - assert_eq!(parts.len(), 1, "{parts:?}"); - assert!(parts[0].is_err(), "{parts:?}"); -} - -/// Dropping the stream closes the socket: upstream's `cancel` calls -/// `connection.close()`. -#[tokio::test] -async fn dropping_the_stream_closes_the_socket() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let (closed_tx, closed_rx) = tokio::sync::oneshot::channel(); - tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap(); - ws.send(Message::Text( - json!({"type": "transcript.created"}).to_string(), - )) - .await - .unwrap(); - let mut closed = false; - while let Some(Ok(message)) = ws.next().await { - if message.is_close() { - closed = true; - break; - } - } - let _ = closed_tx.send(closed); - }); - - let mut options = stream_options(None, "audio/pcm", Some(16000)); - options.audio = Box::pin(futures::stream::pending()); - let mut result = model(&base_url, None) - .do_stream(options) - .expect("the model has do_stream") - .await - .expect("do_stream"); - let first = result.stream.next().await; - assert!( - matches!(first, Some(Ok(TranscriptionStreamPart::StreamStart { .. }))), - "{first:?}" - ); - drop(result); - - let closed = tokio::time::timeout(Duration::from_secs(5), closed_rx) - .await - .expect("the server saw the socket end") - .unwrap(); - assert!( - closed, - "no close frame was sent when the stream was dropped" - ); -} diff --git a/aimux-providers/tests/xai_upstream_e2e_test.rs b/aimux-providers/tests/xai_upstream_e2e_test.rs deleted file mode 100644 index 637ee4cbd..000000000 --- a/aimux-providers/tests/xai_upstream_e2e_test.rs +++ /dev/null @@ -1,286 +0,0 @@ -//! Small Responses API subset from the ported xAI upstream cases. - -#[path = "common/mock_fetch.rs"] -mod mock_fetch; - -use aimux_core::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::LanguageModelMessage; -use aimux_core::options::CallOptions; -use aimux_core::result::GenerateContent; -use aimux_core::shared::provider_namespace; -use aimux_core::stream_part::StreamPart; -use aimux_core::tool::{ProviderTool, Tool}; -use aimux_core::types::FinishReasonUnified; -use aimux_providers::xai::{XAIProviderSettings, XaiResponsesModel, create_xai}; -use futures::TryStreamExt; -use mock_fetch::{Canned, MockFetch}; -use serde_json::{Value, json}; - -fn model(fetch: &std::sync::Arc) -> XaiResponsesModel { - create_xai(XAIProviderSettings { - base_url: Some("https://example.test/v1/".into()), - api_key: Some("test-key".into()), - fetch: Some(fetch.transport()), - ..Default::default() - }) - .unwrap() - .responses("grok-4-fast-non-reasoning") -} - -fn options() -> CallOptions { - CallOptions { - prompt: vec![LanguageModelMessage::user_text("hello")], - ..Default::default() - } -} - -fn assert_request(fetch: &MockFetch, stream: bool, tools: bool) { - let requests = fetch.seen(); - assert_eq!(requests.len(), 1); - let request = &requests[0]; - assert_eq!(request.method, "POST"); - assert_eq!(request.url, "https://example.test/v1/responses"); - assert_eq!(request.headers["authorization"], "Bearer test-key"); - let mut expected = json!({ - "model": "grok-4-fast-non-reasoning", - "input": [{"role": "user", "content": [{"type": "input_text", "text": "hello"}]}] - }); - if stream { - expected["stream"] = json!(true); - } - if tools { - expected["tools"] = json!([{"type": "web_search"}]); - } - assert_eq!(request.json_body(), expected); -} - -fn sse(events: &[Value]) -> Canned { - Canned { - status: 200, - headers: vec![("content-type".into(), "text/event-stream".into())], - body: (events - .iter() - .map(|event| format!("data: {event}\n\n")) - .collect::() - + "data: [DONE]\n\n") - .into_bytes(), - } -} - -/// partial: `should extract usage correctly`; full port in -/// src/xai/tests/responses/responses_language_model.rs. -#[tokio::test] -async fn generate_converts_usage() { - // The port used an unrelated search fixture for this case; restore the TS response. - let usage = json!({ - "input_tokens": 345, "output_tokens": 538, "total_tokens": 883, - "output_tokens_details": {"reasoning_tokens": 123} - }); - let body = json!({ - "id": "resp_123", "object": "response", "status": "completed", - "model": "grok-4-fast-non-reasoning", "output": [], "usage": usage - }); - let fetch = MockFetch::new(vec![Canned::json(&body)]); - let result = model(&fetch).do_generate(&options()).await.unwrap(); - assert_request(&fetch, false, false); - assert!(result.content.is_empty()); - assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); - assert_eq!(result.finish_reason.raw.as_deref(), Some("completed")); - assert_eq!( - serde_json::to_value(result.usage).unwrap(), - json!({ - "inputTokens": {"total": 345, "noCache": 345, "cacheRead": 0}, - "outputTokens": {"total": 538, "text": 415, "reasoning": 123}, - "raw": usage - }) - ); - assert!(result.provider_metadata.is_none()); - assert_eq!(result.response.unwrap().id.as_deref(), Some("resp_123")); -} - -/// partial: `should expose cost_in_usd_ticks in finish providerMetadata`; full port in -/// src/xai/tests/responses/responses_language_model.rs. -#[tokio::test] -async fn stream_reports_cost_and_text() { - let events = vec![ - json!({"type": "response.created", "response": {"id": "resp_123", "object": "response", "model": "grok-4-fast-non-reasoning", "output": []}}), - json!({"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "Hello"}), - json!({"type": "response.completed", "response": {"id": "resp_123", "object": "response", "model": "grok-4-fast-non-reasoning", "status": "completed", "output": [{"type": "message", "id": "msg_001", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": "Hello"}]}], "usage": {"input_tokens": 10, "output_tokens": 5, "cost_in_usd_ticks": 113500}}}), - ]; - let fetch = MockFetch::new(vec![sse(&events)]); - let parts: Vec<_> = model(&fetch) - .do_stream(&options()) - .await - .unwrap() - .stream - .try_collect() - .await - .unwrap(); - assert_request(&fetch, true, false); - let text: String = parts - .iter() - .filter_map(|part| match part { - StreamPart::TextDelta { delta, .. } => Some(delta.as_str()), - _ => None, - }) - .collect(); - assert_eq!(text, "Hello"); - let finishes: Vec<_> = parts - .iter() - .filter_map(|part| match part { - StreamPart::Finish { - finish_reason, - usage, - provider_metadata, - } => Some((finish_reason, usage, provider_metadata)), - _ => None, - }) - .collect(); - assert_eq!(finishes.len(), 1); - let (reason, usage, metadata) = finishes[0]; - assert_eq!(reason.unified, FinishReasonUnified::Stop); - assert_eq!(reason.raw.as_deref(), Some("completed")); - assert_eq!(usage.input_tokens.total, Some(10)); - assert_eq!(usage.output_tokens.total, Some(5)); - assert_eq!( - metadata.as_ref().unwrap(), - &provider_namespace("xai", json!({"costInUsdTicks": 113500})).unwrap() - ); - assert!( - !parts - .iter() - .any(|part| matches!(part, StreamPart::Error { .. })) - ); -} - -/// partial: `should forward web_search_call action and sources to tool-result`; full port in -/// src/xai/tests/responses/responses_language_model.rs. -#[tokio::test] -async fn tool_results_preserve_actions_and_sources() { - let body = json!({ - "id": "resp_123", - "object": "response", - "status": "completed", - "model": "grok-4-fast-non-reasoning", - "output": [ - {"type": "web_search_call", "id": "ws_action_1", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "search", "query": "latest AI news", "sources": [{"type": "url", "url": "https://example.com/a"}, {"type": "url", "url": "https://example.com/b"}]}}, - {"type": "web_search_call", "id": "ws_action_2", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "open_page", "url": "https://example.com/a"}}, - {"type": "web_search_call", "id": "ws_action_3", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "find_in_page", "url": "https://example.com/a", "pattern": "climate"}}, - {"type": "web_search_call", "id": "ws_action_4", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "search", "query": null, "sources": null}}, - {"type": "web_search_call", "id": "ws_action_5", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "open_page", "url": null}}, - ], - "usage": {"input_tokens": 10, "output_tokens": 5}, - }); - let fetch = MockFetch::new(vec![Canned::json(&body)]); - let mut options = options(); - options.tools = Some(vec![Tool::Provider(ProviderTool { - id: "xai.web_search".parse().unwrap(), - name: "web_search".into(), - args: serde_json::Map::new(), - supports_deferred_results: None, - })]); - let result = model(&fetch).do_generate(&options).await.unwrap(); - assert_request(&fetch, false, true); - let expected = vec![ - json!({"action": {"type": "search", "query": "latest AI news"}, "sources": [{"type": "url", "url": "https://example.com/a"}, {"type": "url", "url": "https://example.com/b"}]}), - json!({"action": {"type": "openPage", "url": "https://example.com/a"}}), - json!({"action": {"type": "findInPage", "url": "https://example.com/a", "pattern": "climate"}}), - json!({"action": {"type": "search"}}), - json!({"action": {"type": "openPage", "url": null}}), - ]; - assert_eq!(result.content.len(), expected.len() * 2); - for (index, (pair, expected_result)) in result - .content - .as_chunks::<2>() - .0 - .iter() - .zip(expected) - .enumerate() - { - let id = format!("ws_action_{}", index + 1); - let GenerateContent::ToolCall(call) = &pair[0] else { - panic!("missing tool call") - }; - assert_eq!(call.tool_call_id, id); - assert_eq!(call.tool_name, "web_search"); - assert_eq!(call.input, ""); - assert_eq!(call.provider_executed, Some(true)); - let GenerateContent::ToolResult(output) = &pair[1] else { - panic!("missing tool result") - }; - assert_eq!(output.tool_call_id, id); - assert_eq!(output.tool_name, "web_search"); - assert_eq!(output.result, expected_result); - assert_eq!(output.is_error, None); - } - assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); -} - -/// partial: `should set finish reason to error`; full port in -/// src/xai/tests/responses/responses_language_model.rs. -#[tokio::test] -async fn failed_response_emits_error_and_finish() { - let events = vec![ - json!({"type": "response.created", "response": {"id": "resp_123", "object": "response", "model": "grok-4-fast-non-reasoning", "output": []}}), - json!({"type": "response.failed", "response": {"error": {"code": "server_error", "message": "Internal server error"}, "usage": {"input_tokens": 50, "output_tokens": 0}}}), - ]; - let fetch = MockFetch::new(vec![sse(&events)]); - let parts: Vec<_> = model(&fetch) - .do_stream(&options()) - .await - .unwrap() - .stream - .try_collect() - .await - .unwrap(); - assert_request(&fetch, true, false); - let errors: Vec<_> = parts - .iter() - .filter_map(|part| match part { - StreamPart::Error { error } => Some(error), - _ => None, - }) - .collect(); - assert_eq!(errors.len(), 1); - let AiMuxError::ApiCall(error) = errors[0] else { - panic!("expected API error") - }; - assert_eq!(error.message, "Internal server error"); - assert_eq!(error.provider_code.as_deref(), Some("server_error")); - assert_eq!(error.status_code, Some(500)); - assert!(error.is_retryable); - assert_eq!(error.data.as_ref(), Some(&events[1])); - let finishes: Vec<_> = parts - .iter() - .filter_map(|part| match part { - StreamPart::Finish { - finish_reason, - usage, - .. - } => Some((finish_reason, usage)), - _ => None, - }) - .collect(); - assert_eq!(finishes.len(), 1); - let (reason, usage) = finishes[0]; - assert_eq!(reason.unified, FinishReasonUnified::Error); - assert_eq!(reason.raw.as_deref(), Some("error")); - assert_eq!(usage.input_tokens.total, Some(50)); - assert_eq!(usage.output_tokens.total, Some(0)); -} - -/// `xaiResponsesSupportedUrls` (xai-responses-language-model.ts:142-149): -/// image, PDF and text URLs go to the API as they are. -#[tokio::test] -async fn responses_support_http_image_pdf_and_text_urls() { - let fetch = MockFetch::new(vec![]); - let urls = model(&fetch).supported_urls().await; - let mut kinds: Vec<_> = urls.0.keys().map(String::as_str).collect(); - kinds.sort_unstable(); - assert_eq!(kinds, ["application/pdf", "image/*", "text/*"]); - for patterns in urls.0.values() { - assert!(patterns[0].is_match("https://example.com/file")); - assert!(!patterns[0].is_match("ftp://example.com/file")); - } -} diff --git a/aimux-providers/tests/xai_video_test.rs b/aimux-providers/tests/xai_video_test.rs deleted file mode 100644 index 4c5ec5e25..000000000 --- a/aimux-providers/tests/xai_video_test.rs +++ /dev/null @@ -1,868 +0,0 @@ -//! Rust translation of the xAI video model tests. -//! Source: `reference/aisdk-pinned/xai/src/xai-video-model.test.ts` -//! -//! The HTTP layer is mocked with wiremock. Requests are inspected through -//! `received_requests`; `do_start` / `do_status` are called directly. - -use std::collections::HashMap; - -use serde_json::{Map, Value, json}; -use wiremock::matchers::method; -use wiremock::{Mock, MockServer, Request, ResponseTemplate}; - -use aimux_core::error::AiMuxError; -use aimux_core::shared::{AspectRatio, Size, Warning}; -use aimux_core::video_model::{ - VideoCallOptions, VideoData, VideoFile, VideoFileData, VideoFrameImage, VideoFrameType, - VideoModel, VideoOperationStart, VideoOperationStatus, VideoResult, -}; -use aimux_providers::{XAIProviderSettings, XaiVideoModel, create_xai}; - -const PROMPT: &str = "A chicken flying into the sunset"; -const SRC: &str = "https://example.com/source-video.mp4"; -const VIDEO_URL: &str = "https://vidgen.x.ai/output/video-001.mp4"; -const V15: &str = "grok-imagine-video-1.5"; -const V1: &str = "grok-imagine-video"; - -// ---------------------------------------------------------------- helpers - -fn done_status() -> Value { - json!({ - "status": "done", - "video": {"url": VIDEO_URL, "duration": 5, "respect_moderation": true}, - "model": V1, - "progress": 100 - }) -} - -fn json_response(body: Value) -> ResponseTemplate { - ResponseTemplate::new(200).set_body_json(body) -} - -async fn server_with(post: Value, status: ResponseTemplate) -> MockServer { - let server = MockServer::start().await; - Mock::given(method("POST")) - .respond_with(json_response(post)) - .mount(&server) - .await; - Mock::given(method("GET")) - .respond_with(status) - .mount(&server) - .await; - server -} - -fn model_with(server: &MockServer, model_id: &str, headers: &[(&str, &str)]) -> XaiVideoModel { - create_xai(XAIProviderSettings { - api_key: Some("test-key".into()), - base_url: Some(server.uri()), - headers: (!headers.is_empty()).then(|| { - headers - .iter() - .map(|(k, v)| ((*k).to_string(), Some((*v).to_string()))) - .collect() - }), - ..Default::default() - }) - .unwrap() - .video(model_id) -} - -fn base() -> VideoCallOptions { - VideoCallOptions::new(PROMPT) -} - -/// Sets `providerOptions.xai`. -fn xai(mut options: VideoCallOptions, value: Value) -> VideoCallOptions { - options - .provider_options - .insert("xai".into(), value.as_object().unwrap().clone()); - options -} - -fn ar(s: &str) -> Option { - Some(AspectRatio::parse(s).unwrap()) -} - -fn size(s: &str) -> Option { - Some(Size::parse(s).unwrap()) -} - -fn url_file(url: &str, media_type: Option<&str>) -> VideoFile { - VideoFile::Url { - url: url.into(), - media_type: media_type.map(str::to_owned), - } -} - -fn last_frame(file: VideoFile) -> Option> { - Some(vec![VideoFrameImage { - image: file, - frame_type: VideoFrameType::LastFrame, - }]) -} - -fn header(request: &Request, name: &str) -> String { - request - .headers - .get(name) - .map(|v| v.to_str().unwrap().to_owned()) - .unwrap_or_default() -} - -struct Started { - start: VideoOperationStart, - method: String, - path: String, - body: Value, -} - -async fn try_start(model_id: &str, options: &VideoCallOptions) -> Result { - let server = server_with( - json!({"request_id": "req-123"}), - json_response(done_status()), - ) - .await; - let start = model_with(&server, model_id, &[]).do_start(options).await?; - let request = server.received_requests().await.unwrap().remove(0); - Ok(Started { - start, - method: request.method.to_string(), - path: request.url.path().to_owned(), - body: serde_json::from_slice(&request.body).unwrap(), - }) -} - -async fn start(model_id: &str, options: VideoCallOptions) -> Started { - try_start(model_id, &options).await.unwrap() -} - -/// The `details` of every `Unsupported` warning for `feature`. -fn details<'a>(warnings: &'a [Warning], feature: &str) -> Vec<&'a str> { - warnings - .iter() - .filter_map(|w| match w { - Warning::Unsupported { - feature: f, - details, - } if f == feature => Some(details.as_deref().unwrap_or("")), - _ => None, - }) - .collect() -} - -fn warned(warnings: &[Warning], feature: &str) -> bool { - !details(warnings, feature).is_empty() -} - -async fn poll_once( - status: ResponseTemplate, - request_id: &str, -) -> (Result, Vec) { - let server = server_with(json!({"request_id": "req-123"}), status).await; - let result = model_with(&server, V1, &[]) - .do_status(&json!({ "requestId": request_id }), &base()) - .await; - (result, server.received_requests().await.unwrap()) -} - -fn completed(result: Result) -> VideoResult { - match result { - Ok(VideoOperationStatus::Completed(done)) => done, - other => panic!("expected completed, got {other:?}"), - } -} - -fn xai_meta(result: &VideoResult) -> Map { - result.provider_metadata.as_ref().unwrap()["xai"].clone() -} - -fn assert_terminal(result: Result, message: &str) { - match result { - Err(AiMuxError::ApiCall(e)) => { - assert_eq!(e.message, message); - assert!(!e.is_retryable); - assert_eq!(e.status_code, Some(200)); - } - other => panic!("expected terminal api error, got {other:?}"), - } -} - -mod xai_video_model { - use super::*; - - mod constructor { - use super::*; - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › constructor › should send the grok-imagine-video-1.5 model id in the request body - #[tokio::test] - async fn sends_the_1_5_model_id_in_the_body() { - let s = start(V15, base()).await; - assert_eq!(s.body["model"], V15); - } - } - - mod do_start { - use super::*; - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should map current xAI generation options - #[tokio::test] - async fn maps_current_generation_options() { - let mut o = xai( - base(), - json!({ - "storageOptions": { - "filename": "result.mp4", - "expiresAfter": 86_400, - "publicUrl": {"expiresAfter": 3_600} - }, - "keyframes": [{"imageUrl": "https://example.com/middle.png", "timestampSeconds": 2.5}] - }), - ); - o.generate_audio = Some(false); - o.frame_images = last_frame(url_file("https://example.com/end.png", None)); - let s = start(V15, o).await; - assert_eq!(s.body["generate_audio"], false); - assert_eq!( - s.body["last_frame"], - json!({"url": "https://example.com/end.png"}) - ); - assert_eq!( - s.body["storage_options"], - json!({ - "filename": "result.mp4", - "expires_after": 86_400, - "public_url": {"expires_after": 3_600} - }) - ); - assert_eq!( - s.body["keyframes"], - json!([{"image": {"url": "https://example.com/middle.png"}, "timestamp_s": 2.5}]) - ); - assert_eq!(s.start.operation, json!({"requestId": "req-123"})); - } - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should warn and omit last_frame for grok-imagine-video - #[tokio::test] - async fn warns_and_omits_last_frame_for_grok_imagine_video() { - let mut o = base(); - o.frame_images = last_frame(url_file("https://example.com/end.png", None)); - let s = start(V1, o).await; - assert!(s.body.get("last_frame").is_none()); - assert_eq!( - details(&s.start.warnings, "frameImages"), - [ - "xAI only supports last_frame with \"grok-imagine-video-1.5\". \ - The last frame was ignored." - ] - ); - } - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should pass headers - #[tokio::test] - async fn do_start_passes_headers() { - let server = server_with( - json!({"request_id": "req-123"}), - json_response(done_status()), - ) - .await; - let model = model_with(&server, V1, &[("X-Custom", "value")]); - let mut o = base(); - o.headers = Some(HashMap::from([( - "X-Request-Header".to_string(), - "request-value".to_string(), - )])); - model.do_start(&o).await.unwrap(); - let request = server.received_requests().await.unwrap().remove(0); - assert_eq!(header(&request, "authorization"), "Bearer test-key"); - assert_eq!(header(&request, "x-custom"), "value"); - assert_eq!(header(&request, "x-request-header"), "request-value"); - } - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should throw when no request_id returned - #[tokio::test] - async fn do_start_fails_without_request_id() { - let server = server_with(json!({}), json_response(done_status())).await; - let err = model_with(&server, V1, &[]) - .do_start(&base()) - .await - .unwrap_err(); - match err { - AiMuxError::Other(message) => { - assert_eq!(message, "No request_id returned from xAI API.") - } - other => panic!("unexpected error: {other:?}"), - } - } - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should return warnings for unsupported features - #[tokio::test] - async fn warns_for_fps_and_seed() { - let mut o = base(); - o.fps = Some(30); - o.seed = Some(42); - let s = start(V1, o).await; - assert!(warned(&s.start.warnings, "fps")); - assert!(warned(&s.start.warnings, "seed")); - } - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should warn when explicit R2V has no references at all - #[tokio::test] - async fn warns_when_reference_to_video_has_no_references() { - let s = start(V15, xai(base(), json!({"mode": "reference-to-video"}))).await; - assert!(s.body.get("reference_images").is_none()); - assert!( - details(&s.start.warnings, "referenceImages")[0] - .contains("without reference images") - ); - } - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should allow duration and aspectRatio with reference images - #[tokio::test] - async fn allows_duration_and_aspect_ratio_with_reference_images() { - let mut o = xai( - base(), - json!({"referenceImageUrls": ["https://example.com/ref.jpg"]}), - ); - o.duration = Some(8); - o.aspect_ratio = ar("16:9"); - let s = start(V1, o).await; - assert_eq!(s.body["duration"], 8); - assert_eq!(s.body["aspect_ratio"], "16:9"); - assert!(!warned(&s.start.warnings, "duration")); - assert!(!warned(&s.start.warnings, "aspectRatio")); - } - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should omit reference_audios for an empty referenceVoiceIds array - #[tokio::test] - async fn omits_reference_audios_for_empty_voice_ids() { - let o = xai( - base(), - json!({ - "mode": "reference-to-video", - "referenceImageUrls": ["https://example.com/ref1.jpg"], - "referenceVoiceIds": [] - }), - ); - let s = start(V15, o).await; - assert!(s.body.get("reference_audios").is_none()); - assert!(!warned(&s.start.warnings, "referenceVoiceIds")); - } - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStart › should warn and omit referenceVoiceIds outside reference-to-video - #[tokio::test] - async fn warns_and_omits_voice_ids_outside_reference_to_video() { - let s = start(V15, xai(base(), json!({"referenceVoiceIds": ["eve"]}))).await; - assert!(s.body.get("reference_audios").is_none()); - assert!(warned(&s.start.warnings, "referenceVoiceIds")); - } - } - - mod do_status { - use super::*; - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStatus › should return pending when status is pending - #[tokio::test] - async fn do_status_returns_pending() { - for template in [ - json_response(json!({"status": "pending"})), - // xAI answers 202 while running, sometimes with an empty body. - ResponseTemplate::new(202), - ResponseTemplate::new(202) - .set_body_json(json!({"status": "pending", "progress": 40})), - ] { - let (result, _) = poll_once(template, "req-123").await; - assert!( - matches!(result, Ok(VideoOperationStatus::Pending)), - "{result:?}" - ); - } - } - - /// upstream: @ai-sdk/xai src/xai-video-model.test.ts › XaiVideoModel › doStatus › should pass headers to request - #[tokio::test] - async fn do_status_passes_headers() { - let server = server_with(json!({}), json_response(done_status())).await; - let model = model_with(&server, V1, &[("X-Custom", "value")]); - let mut o = base(); - o.headers = Some(HashMap::from([( - "X-Request-Header".to_string(), - "request-value".to_string(), - )])); - model - .do_status(&json!({"requestId": "req-123"}), &o) - .await - .unwrap(); - let request = server.received_requests().await.unwrap().remove(0); - assert_eq!(header(&request, "authorization"), "Bearer test-key"); - assert_eq!(header(&request, "x-custom"), "value"); - assert_eq!(header(&request, "x-request-header"), "request-value"); - } - } -} - -/// Checks the provider id, model id and maxVideosPerCall; specificationVersion is not checked. partial: should expose correct provider and model information; full port in src/xai/tests/ -#[tokio::test] -async fn constructor_exposes_provider_and_model_information() { - let server = MockServer::start().await; - let model = model_with(&server, V1, &[]); - assert_eq!(model.provider(), "xai.video"); - assert_eq!(model.model_id(), V1); - assert_eq!(model.max_videos_per_call(), Some(1)); -} - -/// Checks the requestId operation and the response model id, timestamp and headers. partial: should return operation with requestId; should include response metadata; full port in src/xai/tests/ -#[tokio::test] -async fn do_start_returns_request_id_and_response_metadata() { - let s = start(V1, base()).await; - assert_eq!(s.start.operation, json!({"requestId": "req-123"})); - let response = &s.start.response; - assert_eq!(response.model_id.as_deref(), Some(V1)); - assert!(response.timestamp.is_some()); - assert!(response.headers.is_some()); -} - -/// Checks the POST to /videos/generations with the minimal body and no warnings. partial: should pass correct request body; should return empty warnings for supported features; full port in src/xai/tests/ -#[tokio::test] -async fn do_start_posts_minimal_generation_body() { - let s = start(V1, base()).await; - assert_eq!(s.method, "POST"); - assert_eq!(s.path, "/videos/generations"); - assert_eq!(s.body, json!({"model": V1, "prompt": PROMPT})); - assert!(s.start.warnings.is_empty()); -} - -/// Checks the POST path and video field for legacy, edit and extend modes. partial: should use edits endpoint for video editing (legacy videoUrl without mode); should use edits endpoint for video editing with explicit mode; should use extensions endpoint for extend-video mode; full port in src/xai/tests/ -#[tokio::test] -async fn selects_endpoint_by_mode() { - let cases = [ - (json!({"videoUrl": SRC}), "/videos/edits"), - ( - json!({"mode": "edit-video", "videoUrl": SRC}), - "/videos/edits", - ), - ( - json!({"mode": "extend-video", "videoUrl": SRC}), - "/videos/extensions", - ), - ]; - for (options, path) in cases { - let s = start(V1, xai(base(), options)).await; - assert_eq!(s.method, "POST"); - assert_eq!(s.path, path); - assert_eq!(s.body["video"], json!({"url": SRC})); - } -} - -/// Checks the n warning for n=3 and its absence for n=1. partial: should warn when n > 1; should not warn when n is 1; full port in src/xai/tests/ -#[tokio::test] -async fn warns_only_when_n_exceeds_one() { - for (n, expected) in [(3, true), (1, false)] { - let mut o = base(); - o.n = n; - let s = start(V1, o).await; - assert_eq!(warned(&s.start.warnings, "n"), expected, "n = {n}"); - } -} - -/// Checks that duration and aspect_ratio are sent in the body. partial: should send duration in request body; should send aspect_ratio in request body; full port in src/xai/tests/ -#[tokio::test] -async fn sends_duration_and_aspect_ratio() { - let mut o = base(); - o.duration = Some(10); - o.aspect_ratio = ar("9:16"); - let s = start(V1, o).await; - assert_eq!(s.body["duration"], 10); - assert_eq!(s.body["aspect_ratio"], "9:16"); -} - -/// Checks that SDK sizes map to 720p, 480p and 1080p. partial: should map SDK resolution 1280x720 to 720p; should map SDK resolution 854x480 to 480p; should map SDK resolution 640x480 to 480p; should map SDK resolution 1920x1080 to 1080p; full port in src/xai/tests/ -#[tokio::test] -async fn maps_sdk_resolution_to_xai_resolution() { - let cases = [ - (V1, "1280x720", "720p"), - (V1, "854x480", "480p"), - (V1, "640x480", "480p"), - (V15, "1920x1080", "1080p"), - ]; - for (model, resolution, expected) in cases { - let mut o = base(); - o.resolution = size(resolution); - let s = start(model, o).await; - assert_eq!(s.body["resolution"], expected, "{resolution}"); - } -} - -/// Checks that a provider resolution is sent and wins over the SDK size. partial: should pass through provider option resolution 1080p; should prefer provider option resolution over SDK resolution; full port in src/xai/tests/ -#[tokio::test] -async fn provider_option_resolution_wins() { - let s = start(V15, xai(base(), json!({"resolution": "1080p"}))).await; - assert_eq!(s.body["resolution"], "1080p"); - - let mut o = xai(base(), json!({"resolution": "480p"})); - o.resolution = size("1280x720"); - let s = start(V1, o).await; - assert_eq!(s.body["resolution"], "480p"); -} - -/// Checks the 1080p warning on grok-imagine-video (SDK and provider sources) and none on 1.5. partial: should warn when SDK resolution 1920x1080 is used with grok-imagine-video; should warn when provider resolution 1080p is used with grok-imagine-video; should not warn about 1080p with grok-imagine-video-1.5; full port in src/xai/tests/ -#[tokio::test] -async fn warns_about_1080p_only_on_grok_imagine_video() { - let mut sdk = base(); - sdk.resolution = size("1920x1080"); - let provider = xai(base(), json!({"resolution": "1080p"})); - for options in [sdk, provider] { - let s = start(V1, options).await; - assert_eq!(s.body["resolution"], "1080p"); - assert!( - details(&s.start.warnings, "resolution")[0].contains("does not support 1080p"), - "{:?}", - s.start.warnings - ); - } - let mut o = base(); - o.resolution = size("1920x1080"); - assert!(start(V15, o).await.start.warnings.is_empty()); -} - -/// Checks that unrecognized resolutions are left out of the body and warned about. partial: should warn for unrecognized resolution format; should warn and omit body resolution for completely unknown format; full port in src/xai/tests/ -#[tokio::test] -async fn warns_and_omits_unrecognized_resolution() { - for resolution in ["2560x1440", "3840x2160"] { - let mut o = base(); - o.resolution = size(resolution); - let s = start(V1, o).await; - assert!(s.body.get("resolution").is_none()); - assert!(warned(&s.start.warnings, "resolution"), "{resolution}"); - } -} - -/// Checks that URL, binary and base64 image inputs are sent as image.url. partial: should send image object from URL-based image input; should send image object with data URI from file data; should send image object with data URI from base64 string; full port in src/xai/tests/ -#[tokio::test] -async fn sends_image_object_from_url_and_file_inputs() { - let cases = [ - ( - url_file("https://example.com/image.png", None), - "https://example.com/image.png", - ), - ( - VideoFile::File { - media_type: "image/png".into(), - data: VideoFileData::Binary(vec![137, 80, 78, 71]), - }, - "data:image/png;base64,iVBORw==", - ), - ( - VideoFile::File { - media_type: "image/jpeg".into(), - data: VideoFileData::Base64("aGVsbG8=".into()), - }, - "data:image/jpeg;base64,aGVsbG8=", - ), - ]; - for (image, expected) in cases { - let mut o = base(); - o.image = Some(image); - let s = start(V1, o).await; - assert_eq!(s.body["image"], json!({"url": expected})); - } -} - -/// Checks which options warn in default, edit and extend modes. partial: should warn about duration in edit mode; should warn about aspectRatio in edit mode; should warn about resolution in edit mode; should not warn about duration outside edit mode; should not warn about aspectRatio outside edit mode; should not warn about resolution outside edit mode; should warn about aspectRatio in extension mode; should warn about resolution in extension mode; should not warn about duration in extension mode; full port in src/xai/tests/ -#[tokio::test] -async fn warns_about_unsupported_options_per_mode() { - type Setter = fn(&mut VideoCallOptions); - let setters: [(&str, Setter); 3] = [ - ("duration", |o| o.duration = Some(10)), - ("aspectRatio", |o| o.aspect_ratio = ar("16:9")), - ("resolution", |o| o.resolution = size("1280x720")), - ]; - for mode in [None, Some("edit-video"), Some("extend-video")] { - for (feature, set) in setters { - let mut o = base(); - set(&mut o); - if let Some(mode) = mode { - o = xai(o, json!({"mode": mode, "videoUrl": SRC})); - } - let expected = matches!( - (mode, feature), - (Some("edit-video"), _) | (Some("extend-video"), "aspectRatio" | "resolution") - ); - let s = start(V1, o).await; - assert_eq!( - warned(&s.start.warnings, feature), - expected, - "{mode:?} {feature}" - ); - } - } -} - -/// Checks that unsupported fields are left out of edit and extend bodies. partial: should omit duration, aspect_ratio, and resolution from body in edit mode; should allow duration in extension mode; should omit aspect_ratio and resolution from body in extension mode; should warn about provider-level resolution in extension mode; full port in src/xai/tests/ -#[tokio::test] -async fn omits_unsupported_fields_from_edit_and_extension_bodies() { - let mut edit = xai(base(), json!({"mode": "edit-video", "videoUrl": SRC})); - edit.duration = Some(10); - edit.aspect_ratio = ar("16:9"); - edit.resolution = size("1280x720"); - let s = start(V1, edit).await; - for key in ["duration", "aspect_ratio", "resolution"] { - assert!(s.body.get(key).is_none(), "edit body has {key}"); - } - - let mut extend = xai(base(), json!({"mode": "extend-video", "videoUrl": SRC})); - extend.duration = Some(6); - extend.aspect_ratio = ar("16:9"); - extend.resolution = size("1280x720"); - let s = start(V1, extend).await; - assert_eq!(s.body["duration"], 6); - assert!(s.body.get("aspect_ratio").is_none()); - assert!(s.body.get("resolution").is_none()); - - let provider_resolution = xai( - base(), - json!({"mode": "extend-video", "videoUrl": SRC, "resolution": "720p"}), - ); - let s = start(V1, provider_resolution).await; - assert!(s.body.get("resolution").is_none()); - assert!(warned(&s.start.warnings, "resolution")); -} - -/// Checks reference_images for explicit and implicit reference-to-video. partial: should send reference_images array to /videos/generations with explicit mode; should fallback to reference-to-video mode when referenceImageUrls is set without mode; full port in src/xai/tests/ -#[tokio::test] -async fn sends_reference_images_for_reference_to_video() { - let urls = [ - "https://example.com/ref1.jpg", - "https://example.com/ref2.jpg", - ]; - let explicit = xai( - base(), - json!({"mode": "reference-to-video", "referenceImageUrls": urls}), - ); - let implicit = xai(base(), json!({"referenceImageUrls": urls})); - for options in [explicit, implicit] { - let s = start(V1, options).await; - assert_eq!(s.method, "POST"); - assert_eq!(s.path, "/videos/generations"); - assert_eq!( - s.body["reference_images"], - json!([{"url": urls[0]}, {"url": urls[1]}]) - ); - } -} - -/// Checks that 1080p is downgraded to 720p in reference-to-video with a warning. partial: should downgrade R2V 1080p to 720p with a warning; should downgrade R2V 1920x1080 from the SDK resolution to 720p; full port in src/xai/tests/ -#[tokio::test] -async fn downgrades_reference_to_video_1080p() { - let refs = json!(["https://example.com/ref1.jpg"]); - let provider = xai( - base(), - json!({"mode": "reference-to-video", "referenceImageUrls": refs, "resolution": "1080p"}), - ); - let mut sdk = xai( - base(), - json!({"mode": "reference-to-video", "referenceImageUrls": refs}), - ); - sdk.resolution = size("1920x1080"); - for options in [provider, sdk] { - let s = start(V15, options).await; - assert_eq!(s.body["resolution"], "720p"); - assert!(warned(&s.start.warnings, "resolution")); - } -} - -/// Checks that inputReferences split into reference_images and reference_audios by media type. partial: should separate image and audio inputReferences; should support audio-only reference-to-video; should not send an empty reference_images array for video-only inputReferences; should combine a pinned first frame with an audio reference; should support explicit audio-only R2V without reference_images; full port in src/xai/tests/ -#[tokio::test] -async fn splits_input_references_by_media_type() { - let image = url_file("https://example.com/ref1.jpg", None); - let audio = url_file("https://example.com/voice.mp3", Some("audio/mpeg")); - let video = url_file("https://example.com/clip.mp4", Some("video/mp4")); - let start_image = url_file("https://example.com/start.jpg", Some("image/jpeg")); - let voice = json!([{"url": "https://example.com/voice.mp3"}]); - - let mut o = base(); - o.input_references = Some(vec![image, audio.clone()]); - let s = start(V15, o).await; - assert_eq!( - s.body["reference_images"], - json!([{"url": "https://example.com/ref1.jpg"}]) - ); - assert_eq!(s.body["reference_audios"], voice); - assert!(s.start.warnings.is_empty()); - - let mut o = base(); - o.input_references = Some(vec![audio.clone()]); - let s = start(V15, o).await; - assert_eq!(s.body["reference_audios"], voice); - assert!(s.body.get("reference_images").is_none()); - assert!(s.start.warnings.is_empty()); - - let mut o = base(); - o.input_references = Some(vec![video]); - let s = start(V15, o).await; - assert!(s.body.get("reference_images").is_none()); - - let mut o = base(); - o.image = Some(start_image); - o.input_references = Some(vec![audio.clone()]); - let s = start(V15, o).await; - assert_eq!( - s.body["image"], - json!({"url": "https://example.com/start.jpg"}) - ); - assert_eq!(s.body["reference_audios"], voice); - assert!(s.body.get("reference_images").is_none()); - assert!(s.start.warnings.is_empty()); - - let mut o = xai(base(), json!({"mode": "reference-to-video"})); - o.input_references = Some(vec![audio]); - let s = start(V15, o).await; - assert_eq!(s.body["reference_audios"], voice); - assert!(s.body.get("reference_images").is_none()); - assert!(s.start.warnings.is_empty()); -} - -/// Checks that referenceVoiceIds become reference_audios in order and are never sent raw. partial: should send reference_audios for preset reference voices; should send up to 3 preset reference voices in order; should never forward referenceVoiceIds as a raw body key; full port in src/xai/tests/ -#[tokio::test] -async fn sends_reference_voice_ids_as_reference_audios() { - let cases = [ - (json!(["eve"]), json!([{"voice_id": "eve"}])), - ( - json!(["eve", "leo", "rex"]), - json!([{"voice_id": "eve"}, {"voice_id": "leo"}, {"voice_id": "rex"}]), - ), - ]; - for (ids, expected) in cases { - let o = xai( - base(), - json!({ - "mode": "reference-to-video", - "referenceImageUrls": ["https://example.com/ref1.jpg"], - "referenceVoiceIds": ids - }), - ); - let s = start(V15, o).await; - assert_eq!(s.path, "/videos/generations"); - assert_eq!( - s.body["reference_images"], - json!([{"url": "https://example.com/ref1.jpg"}]) - ); - assert_eq!(s.body["reference_audios"], expected); - assert!(s.body.get("referenceVoiceIds").is_none()); - } -} - -/// Checks that more than 3 voices, or an empty voice id, gives InvalidArgument. partial: should reject more than 3 preset reference voices; should reject empty-string preset reference voice ids; full port in src/xai/tests/ -#[tokio::test] -async fn rejects_invalid_reference_voice_ids() { - for ids in [json!(["ara", "eve", "leo", "rex"]), json!([""])] { - let o = xai( - base(), - json!({ - "mode": "reference-to-video", - "referenceImageUrls": ["https://example.com/ref1.jpg"], - "referenceVoiceIds": ids - }), - ); - match try_start(V15, &o).await { - Err(AiMuxError::InvalidArgument(_)) => {} - Err(other) => panic!("unexpected error: {other:?}"), - Ok(_) => panic!("expected InvalidArgument for {ids}"), - } - } -} - -/// Checks that the requestId is encoded as one URL path segment, dot segments included. partial: should encode the request ID as a single URL path segment; should preserve the $requestId request ID as a URL path segment; full port in src/xai/tests/ -#[tokio::test] -async fn do_status_encodes_request_id_as_one_path_segment() { - let cases = [ - ("abc/../../internal", "/videos/abc%2F..%2F..%2Finternal"), - (".", "/videos/%252E"), - ("..", "/videos/%252E%252E"), - ]; - for (request_id, expected) in cases { - let (_, requests) = poll_once(json_response(done_status()), request_id).await; - assert_eq!(requests[0].method.as_str(), "GET"); - assert_eq!(requests[0].url.path(), expected); - } -} - -/// Checks the completed video, its xai metadata and the response. partial: should return completed with video data when done; should include response metadata; full port in src/xai/tests/ -#[tokio::test] -async fn do_status_returns_completed_video_and_metadata() { - let (result, _) = poll_once(json_response(done_status()), "req-123").await; - let done = completed(result); - assert_eq!(done.videos.len(), 1); - match &done.videos[0] { - VideoData::Url { url, media_type } => { - assert_eq!(url, VIDEO_URL); - assert_eq!(media_type, "video/mp4"); - } - other => panic!("expected url, got {other:?}"), - } - let meta = xai_meta(&done); - assert_eq!(meta.len(), 4); - assert_eq!(meta["requestId"], "req-123"); - assert_eq!(meta["videoUrl"], VIDEO_URL); - assert_eq!(meta["duration"].as_f64(), Some(5.0)); - assert_eq!(meta["progress"].as_f64(), Some(100.0)); - assert_eq!(done.response.model_id.as_deref(), Some(V1)); - assert!(done.response.timestamp.is_some()); - assert!(done.response.headers.is_some()); -} - -/// Checks that expired, failed, missing-URL and blocked statuses give non-retryable errors. partial: should return error status on expired; should return error status on failed; should report an error status when video URL missing on done; should report an error status when respect_moderation is false; full port in src/xai/tests/ -#[tokio::test] -async fn do_status_reports_terminal_failures_as_non_retryable_errors() { - let cases = [ - ( - json!({"status": "expired", "model": V1}), - "Video generation request expired.", - ), - ( - json!({"status": "failed", "model": V1, "progress": 0, - "error": {"message": "Content policy violation"}}), - "Video generation failed: Content policy violation", - ), - ( - json!({"status": "done", "video": null, "model": V1}), - "Video generation completed but no video URL was returned.", - ), - ( - json!({"status": "done", "video": {"url": "", "respect_moderation": false}, "model": V1}), - "Video generation was blocked due to a content policy violation.", - ), - ]; - for (body, message) in cases { - let (result, _) = poll_once(json_response(body), "req-123").await; - assert_terminal(result, message); - } -} - -/// Checks that costInUsdTicks is added and that duration and progress are omitted when absent. partial: should include costInUsdTicks when returned in usage; should omit duration from metadata when absent in response; should omit progress from metadata when absent in response; full port in src/xai/tests/ -#[tokio::test] -async fn do_status_metadata_includes_optional_fields_only_when_present() { - let mut with_usage = done_status(); - with_usage["usage"] = json!({"cost_in_usd_ticks": 4_000_000_000_u64}); - let mut no_duration = done_status(); - no_duration["video"] - .as_object_mut() - .unwrap() - .remove("duration"); - let mut no_progress = done_status(); - no_progress.as_object_mut().unwrap().remove("progress"); - - let meta = xai_meta(&completed( - poll_once(json_response(with_usage), "req-123").await.0, - )); - assert_eq!(meta.len(), 5); - assert_eq!(meta["costInUsdTicks"].as_f64(), Some(4_000_000_000.0)); - let meta = xai_meta(&completed( - poll_once(json_response(no_duration), "req-123").await.0, - )); - assert!(!meta.contains_key("duration")); - let meta = xai_meta(&completed( - poll_once(json_response(no_progress), "req-123").await.0, - )); - assert!(!meta.contains_key("progress")); -} From dd0d5df975b535c566ab1ff99323fe1a3b28ada0 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 09:12:23 +0000 Subject: [PATCH 2/2] docs(changelog): list the xAI module paths the move changes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RFC-0046 §4.5 item 2: the vendor move PR lists the old public sub-module paths and the replacement entry under Unreleased > Breaking. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0116HPMoqaLfFr1TyqDq1hrW --- CHANGELOG.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3912ed030..361bd71ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1135,6 +1135,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 -- --check` and `python3 scripts/check_dto_vectors.py ` guard the files and the codecs. +- `aimux_providers::xai` follows the upstream `@ai-sdk/xai` file layout + (RFC-0046). The sub-modules `xai::convert`, `xai::responses::convert` and + `xai::responses::types` no longer exist; their public items are re-exported + instead. `xai::convert::supports_reasoning_effort` is now + `xai::supports_reasoning_effort`. + `xai::convert::resolve_provider_reference`, the items of + `xai::responses::convert` (`build_responses_request_body`, + `convert_to_xai_responses_input`, `prepare_responses_tools`, + `convert_xai_responses_usage`, `map_xai_responses_finish_reason`, ...) and + the items of `xai::responses::types` (`XaiResponsesResponse`, + `XaiResponsesUsage`, ...) are now at `xai::responses`. Exports at the crate + root and the `xai` module root are unchanged. + ### Added - Vertex AI MaaS and xAI sub-providers, as `@ai-sdk/google-vertex` 5.0.98's