From d0120544e71bed38a2dc2f64a7501b7489757e83 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 11:29:12 +0800 Subject: [PATCH 01/16] refactor(openai-compatible)!: configure compatible endpoints with data, not closures The compatible chat config holds data where the AI SDK passes functions: - `ConvertUsage` is an enum (OpenAICompatible / Alibaba / MoonshotAI); the Alibaba and Moonshot AI rows select it through `profile.convert_usage` in the registry instead of a name match in preset.rs. - `supported_urls` is a `SupportedUrls` value. - `ChatDialect` is folded into `ChatSettings`; the error structure, which only ever had the default value, is no longer a field. - `MetadataExtractor` / `StreamMetadataExtractor` are removed: nothing implemented them. - Provider headers are a `ProviderHeaders` value resolved per request, and a preset's lazy base URL is `BaseUrl::Preset { descriptor, params }`. Behaviour is unchanged: provider-utils and providers tests pass as before. Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-providers/src/openai_compatible/chat.rs | 53 +---- .../src/openai_compatible/config.rs | 183 ++++++++++-------- .../src/openai_compatible/convert.rs | 22 +-- aimux-providers/src/openai_compatible/mod.rs | 109 +++-------- aimux-providers/src/preset.rs | 73 +++---- aimux-providers/src/provider_registry.json | 8 +- aimux-providers/src/shared/mod.rs | 97 ++++++---- 7 files changed, 243 insertions(+), 302 deletions(-) diff --git a/aimux-providers/src/openai_compatible/chat.rs b/aimux-providers/src/openai_compatible/chat.rs index 7b94fb9b3..860e75ef5 100644 --- a/aimux-providers/src/openai_compatible/chat.rs +++ b/aimux-providers/src/openai_compatible/chat.rs @@ -2,7 +2,7 @@ //! //! The Rust form of `OpenAICompatibleChatLanguageModel`. Request construction //! lives in [`super::convert`]; vendor differences arrive through the -//! [`ChatDialect`](super::config::ChatDialect) of the model's config. +//! [`ChatSettings`](super::config::ChatSettings) of the model's config. use aimux_core::tool::RawToolCall; use std::collections::{HashMap, HashSet}; @@ -58,9 +58,7 @@ impl OpenAICompatibleChatModel { stream, &ChatBodySpec { provider_options_name: self.config.provider_options_name(), - include_usage: self.config.chat.include_usage, - supports_structured_outputs: self.config.chat.supports_structured_outputs, - dialect: &self.config.chat.dialect, + chat: &self.config.chat, }, ) } @@ -189,15 +187,7 @@ impl LanguageModel for OpenAICompatibleChatModel { } fn supported_urls(&self) -> SupportedUrls { - self.config - .chat - .dialect - .supported_urls - .as_ref() - .map_or_else( - || self.config.chat.supported_urls.clone(), - |supported_urls| supported_urls(), - ) + self.config.chat.supported_urls.clone() } async fn do_generate(&self, options: &CallOptions) -> Result { @@ -287,21 +277,10 @@ impl LanguageModel for OpenAICompatibleChatModel { raw: None, }); - let dialect = &self.config.chat.dialect; - let usage = dialect.convert_usage.as_ref().map_or_else( - || usage_from_raw(raw.get("usage")), - |convert| convert(raw.get("usage")), - ); + let usage = self.config.chat.convert_usage.convert(raw.get("usage")); let mut metadata = HashMap::new(); metadata.insert(metadata_key.clone(), Map::new()); - if let Some(extractor) = &dialect.metadata_extractor - && let Some(extra) = extractor.extract_metadata(&raw) - { - for (key, value) in extra { - metadata.insert(key, value); - } - } prediction_tokens( data.usage.as_ref(), @@ -353,11 +332,8 @@ impl LanguageModel for OpenAICompatibleChatModel { let response_headers = resp.response_headers; let mut sse_stream = resp.value; - let dialect = self.config.chat.dialect.clone(); - let mut metadata_extractor = dialect - .metadata_extractor - .as_ref() - .map(|extractor| extractor.create_stream_extractor()); + let convert_usage = self.config.chat.convert_usage; + let stream_usage_key = self.config.chat.stream_usage_key.clone(); let emit_raw_chunks = options.include_raw_chunks == Some(true); let stream_error_url = endpoint; let stream_error_body = body.clone(); @@ -395,9 +371,6 @@ impl LanguageModel for OpenAICompatibleChatModel { if emit_raw_chunks { yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); } - if let Some(extractor) = metadata_extractor.as_mut() { - extractor.process_chunk(&parsed); - } if let Some(error) = parsed.get("error") { yield Ok(StreamPart::Error { error: stream_error( @@ -412,7 +385,7 @@ impl LanguageModel for OpenAICompatibleChatModel { continue; } - let chunk_usage_raw: Option = match &dialect.stream_usage_key { + let chunk_usage_raw: Option = match &stream_usage_key { Some(key) => parsed.get(key).and_then(|v| v.get("usage")).cloned(), None => parsed.get("usage").cloned(), } @@ -623,13 +596,6 @@ impl LanguageModel for OpenAICompatibleChatModel { let mut metadata = HashMap::new(); metadata.insert(metadata_key.clone(), Map::new()); - if let Some(extractor) = &metadata_extractor - && let Some(extra) = extractor.build_metadata() - { - for (key, value) in extra { - metadata.insert(key, value); - } - } prediction_tokens(final_usage_parsed.as_ref(), metadata.entry(metadata_key).or_default()); @@ -643,10 +609,7 @@ impl LanguageModel for OpenAICompatibleChatModel { unified: FinishReasonUnified::Error, raw: None, }), - usage: dialect.convert_usage.as_ref().map_or_else( - || usage_from_raw(final_usage_raw.as_ref()), - |convert| convert(final_usage_raw.as_ref()), - ), + usage: convert_usage.convert(final_usage_raw.as_ref()), provider_metadata: Some(metadata), }); }; diff --git a/aimux-providers/src/openai_compatible/config.rs b/aimux-providers/src/openai_compatible/config.rs index 8a367057c..4284465b5 100644 --- a/aimux-providers/src/openai_compatible/config.rs +++ b/aimux-providers/src/openai_compatible/config.rs @@ -3,112 +3,107 @@ //! Like the native OpenAI config this is the Rust shape of the object //! `createOpenAICompatible` hands each model class (`{ provider, url, headers, //! fetch, ... }`). It has no getters for the credential or the settings that -//! produced it. Compatible endpoint behavior is data (flags and a usage-key -//! name); the shared model never compares a provider name. +//! produced it. Where the AI SDK passes functions (`url`, `headers`, +//! `convertUsage`, `supportedUrls`, `errorStructure`), this config holds data; +//! the shared model never compares a provider name. use std::collections::HashMap; use std::sync::Arc; -use aimux_core::types::{ProviderMetadata, Usage}; +use aimux_core::types::Usage; use serde_json::Value; use aimux_core::AiMuxError; use aimux_core::language_model::SupportedUrls; use aimux_provider_utils::{ - ExchangeContext, FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, ProviderErrorParts, - combine_headers, normalize_headers, + ExchangeContext, FetchFunction, HeaderMapOpt, HttpRequest, ProviderErrorParts, combine_headers, + normalize_headers, }; -pub use crate::shared::TransformRequestBody; +use crate::preset::PresetDescriptor; +use crate::shared::ProviderHeaders; -/// Converts the raw optional token usage of a chat response. -pub type ConvertUsage = Arc) -> Usage + Send + Sync>; +pub use crate::shared::TransformRequestBody; /// Where the base URL comes from. /// /// A plain provider has a [`Fixed`](Self::Fixed) URL. A preset whose URL is /// read from the environment or expanded from template parameters is -/// [`Lazy`](Self::Lazy): it is evaluated per request, so creating the provider -/// (and the preset default instances) never reads the environment and never -/// fails. -#[derive(Clone)] +/// [`Preset`](Self::Preset): it is evaluated per request, so creating the +/// provider (and the preset default instances) never reads the environment +/// and never fails. +#[derive(Clone, Debug)] pub(crate) enum BaseUrl { Fixed(String), - Lazy(Arc Result + Send + Sync>), + Preset { + descriptor: &'static PresetDescriptor, + /// The explicit template parameters of the settings. + params: Arc>, + }, } impl BaseUrl { pub(crate) fn resolve(&self) -> Result { match self { Self::Fixed(url) => Ok(url.clone()), - Self::Lazy(produce) => produce(), + Self::Preset { descriptor, params } => { + crate::preset::resolve_base_url(descriptor, params) + } } } } -/// Maps a provider error payload to the message and code of the API error. -pub(crate) type ErrorStructure = Arc ProviderErrorParts + Send + Sync>; - -/// Captures vendor-specific metadata from responses (the AI SDK's -/// `MetadataExtractor`). The returned map is merged into `provider_metadata` -/// next to the namespace entry. -pub trait MetadataExtractor: Send + Sync { - /// Metadata of a non-streaming response body. - fn extract_metadata(&self, parsed_body: &Value) -> Option; - - /// A fresh extractor for one streaming response. - fn create_stream_extractor(&self) -> Box; -} - -/// The per-stream half of [`MetadataExtractor`]. -pub trait StreamMetadataExtractor: Send { - /// Called with every parsed chunk (including the one carrying usage). - fn process_chunk(&mut self, parsed_chunk: &Value); - - /// Metadata merged into the finish part's `provider_metadata`. - fn build_metadata(&self) -> Option; -} - -/// Everything that distinguishes one compatible vendor's chat endpoint from -/// the baseline, other than the public provider settings. -#[derive(Clone)] -pub(crate) struct ChatDialect { - /// Send `top_k`. The AI SDK baseline never does (it warns and drops it). - pub supports_top_k: bool, - /// Send `tools` / `tool_choice`. - pub supports_tools: bool, - /// Send `response_format`. - pub supports_response_format: bool, - /// The only `max-tokens` key the vendor accepts: `"max_tokens"` or - /// `"max_completion_tokens"`. Preset capability data (RFC-0036 section 5). - pub max_tokens_key: Option<&'static str>, - /// Streaming usage rides in `chunk[key].usage` instead of `chunk.usage`. - pub stream_usage_key: Option, - pub metadata_extractor: Option>, - pub supported_urls: Option SupportedUrls + Send + Sync>>, - pub convert_usage: Option, - pub error_structure: ErrorStructure, +/// How the raw `usage` object of a chat response becomes core [`Usage`] +/// (the AI SDK's `convertUsage`). +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ConvertUsage { + /// The generic OpenAI-shaped usage (`convertOpenAICompatibleChatUsage`). + #[default] + OpenAICompatible, + /// Alibaba: cache writes from `prompt_tokens_details.cache_creation_input_tokens` + /// (or `cache_write_tokens`), taken out of the uncached input. + Alibaba, + /// Moonshot AI: cache reads from the top-level `cached_tokens`. + MoonshotAI, } -impl ChatDialect { - /// The AI SDK baseline. - pub(crate) fn baseline() -> Self { - Self { - supports_top_k: false, - supports_tools: true, - supports_response_format: true, - max_tokens_key: None, - stream_usage_key: None, - metadata_extractor: None, - supported_urls: None, - convert_usage: None, - error_structure: Arc::new(default_error_structure), +impl ConvertUsage { + pub(crate) fn convert(self, raw: Option<&Value>) -> Usage { + let mut usage = super::chat::usage_from_raw(raw); + let Some(raw) = raw.filter(|raw| !raw.is_null()) else { + return usage; + }; + let tokens = |value: &Value| value.as_u64().and_then(|value| u32::try_from(value).ok()); + match self { + Self::OpenAICompatible => {} + Self::Alibaba => { + let details = &raw["prompt_tokens_details"]; + let cache_write = tokens(&details["cache_creation_input_tokens"]) + .or_else(|| tokens(&details["cache_write_tokens"])) + .unwrap_or(0); + usage.input_tokens.cache_write = Some(cache_write); + usage.input_tokens.no_cache = usage + .input_tokens + .no_cache + .map(|tokens| tokens.saturating_sub(cache_write)); + } + Self::MoonshotAI => { + if let Some(cached) = tokens(&raw["cached_tokens"]) { + usage.input_tokens.cache_read = Some(cached); + usage.input_tokens.no_cache = usage + .input_tokens + .total + .map(|tokens| tokens.saturating_sub(cached)); + } + } } + usage } } /// `{ error: { message, type?, param?, code? } }`, the structure nearly every -/// compatible server answers errors with. +/// compatible server answers errors with (the AI SDK's default +/// `errorStructure`). pub(crate) fn default_error_structure(data: &Value) -> ProviderErrorParts { let error = data.get("error").unwrap_or(data); ProviderErrorParts { @@ -127,16 +122,43 @@ pub(crate) fn default_error_structure(data: &Value) -> ProviderErrorParts { } } -/// Returns the URLs supported by chat models. -pub type SupportedUrlsFn = Arc SupportedUrls + Send + Sync>; - -/// Chat-only settings of a provider. -#[derive(Clone)] +/// The chat settings of a compatible endpoint: the AI SDK's +/// `includeUsage` / `supportsStructuredOutputs` / `supportedUrls` / +/// `convertUsage`, plus the preset capability data of RFC-0036 section 5. +#[derive(Clone, Debug)] pub(crate) struct ChatSettings { pub include_usage: bool, pub supports_structured_outputs: bool, pub supported_urls: SupportedUrls, - pub dialect: Arc, + pub convert_usage: ConvertUsage, + /// Send `top_k`. The AI SDK baseline never does (it warns and drops it). + pub supports_top_k: bool, + /// Send `tools` / `tool_choice`. + pub supports_tools: bool, + /// Send `response_format`. + pub supports_response_format: bool, + /// The only `max-tokens` key the vendor accepts: `"max_tokens"` or + /// `"max_completion_tokens"`. + pub max_tokens_key: Option<&'static str>, + /// Streaming usage rides in `chunk[key].usage` instead of `chunk.usage`. + pub stream_usage_key: Option, +} + +impl Default for ChatSettings { + /// The AI SDK baseline. + fn default() -> Self { + Self { + include_usage: false, + supports_structured_outputs: false, + supported_urls: SupportedUrls::default(), + convert_usage: ConvertUsage::default(), + supports_top_k: false, + supports_tools: true, + supports_response_format: true, + max_tokens_key: None, + stream_usage_key: None, + } + } } /// What a model needs to talk to the API. @@ -147,8 +169,8 @@ pub(crate) struct CompatModelConfig { pub base_url: BaseUrl, /// Appended to every request URL, in this order. pub query_params: Option>>, - /// Provider headers (credential, user headers), resolved on every request. - pub headers: HeadersFn, + /// Provider headers (credential, user headers), produced on every request. + pub headers: ProviderHeaders, /// Transport; `None` uses the process default, resolved per request. pub fetch: Option, pub transform_request_body: Option, @@ -228,11 +250,10 @@ impl CompatModelConfig { } } - /// The failed-response handler of this provider's error structure. + /// The failed-response handler of the compatible error structure. pub(crate) fn failed_response_handler( &self, ) -> aimux_provider_utils::ResponseHandler { - let structure = self.chat.dialect.error_structure.clone(); - aimux_provider_utils::create_json_error_response_handler(move |data| structure(data)) + aimux_provider_utils::create_json_error_response_handler(default_error_structure) } } diff --git a/aimux-providers/src/openai_compatible/convert.rs b/aimux-providers/src/openai_compatible/convert.rs index af80c6a36..a200f13ff 100644 --- a/aimux-providers/src/openai_compatible/convert.rs +++ b/aimux-providers/src/openai_compatible/convert.rs @@ -5,7 +5,7 @@ //! `user`, sampling settings, `response_format`, `stop`, `seed`, the //! pass-through fields of the provider's own options namespace, then //! `reasoning_effort`, `verbosity`, `messages`, `tools`, `tool_choice`). -//! Compatible endpoint capabilities arrive as [`ChatDialect`] data. +//! Compatible endpoint capabilities arrive as [`ChatSettings`] data. use aimux_core::shared::SharedProviderOptions; use base64::Engine; @@ -22,7 +22,7 @@ use aimux_core::tool::Tool; use aimux_core::types::{FinishReason, FinishReasonUnified, Warning}; use aimux_provider_utils::{get_top_level_media_type, resolve_full_media_type}; -use super::config::ChatDialect; +use super::config::ChatSettings; /// Option keys of the generic compatible schema /// (`openaiCompatibleLanguageModelChatOptions`): consumed by the model, never @@ -113,9 +113,7 @@ pub(crate) fn to_camel_case(name: &str) -> String { pub(crate) struct ChatBodySpec<'a> { /// First segment of the provider string: the providerOptions namespace. pub provider_options_name: &'a str, - pub include_usage: bool, - pub supports_structured_outputs: bool, - pub dialect: &'a ChatDialect, + pub chat: &'a ChatSettings, } /// A request body and what building it produced. @@ -190,7 +188,7 @@ pub(crate) fn build_request_body( ) -> Result { let name = spec.provider_options_name; let camel = to_camel_case(name); - let dialect = spec.dialect; + let chat = spec.chat; let mut warnings = Vec::new(); // Resolved provider options: the generic namespace, then the provider's @@ -238,7 +236,7 @@ pub(crate) fn build_request_body( } if let Some(max_tokens) = options.max_output_tokens { body.insert( - dialect.max_tokens_key.unwrap_or("max_tokens").into(), + chat.max_tokens_key.unwrap_or("max_tokens").into(), json!(max_tokens), ); } @@ -253,7 +251,7 @@ pub(crate) fn build_request_body( } } if let Some(top_k) = options.top_k { - if dialect.supports_top_k { + if chat.supports_top_k { body.insert("top_k".into(), json!(top_k)); } else { warnings.push(Warning::Unsupported { @@ -268,8 +266,8 @@ pub(crate) fn build_request_body( schema, name: format_name, description, - }) if dialect.supports_response_format => match schema { - Some(schema) if spec.supports_structured_outputs => { + }) if chat.supports_response_format => match schema { + Some(schema) if chat.supports_structured_outputs => { let mut json_schema = Map::new(); json_schema.insert("schema".into(), schema.clone()); json_schema.insert("strict".into(), json!(strict_json_schema)); @@ -341,7 +339,7 @@ pub(crate) fn build_request_body( )?; body.insert("messages".into(), Value::Array(messages)); - if dialect.supports_tools { + if chat.supports_tools { let prepared = prepare_function_tools(options); warnings.extend(prepared.warnings); if let Some(tools) = prepared.tools { @@ -363,7 +361,7 @@ pub(crate) fn build_request_body( if stream { body.insert("stream".into(), json!(true)); - if spec.include_usage { + if chat.include_usage { body.insert("stream_options".into(), json!({ "include_usage": true })); } } diff --git a/aimux-providers/src/openai_compatible/mod.rs b/aimux-providers/src/openai_compatible/mod.rs index 54129ca9d..e4b81a79a 100644 --- a/aimux-providers/src/openai_compatible/mod.rs +++ b/aimux-providers/src/openai_compatible/mod.rs @@ -35,9 +35,7 @@ mod types; pub use chat::OpenAICompatibleChatModel; pub use completion::OpenAICompatibleCompletionModel; -pub use config::{ - ConvertUsage, MetadataExtractor, StreamMetadataExtractor, SupportedUrlsFn, TransformRequestBody, -}; +pub use config::{ConvertUsage, TransformRequestBody}; pub use embedding::OpenAICompatibleEmbeddingModel; pub use image::OpenAICompatibleImageModel; @@ -52,12 +50,10 @@ use aimux_core::image_model::ImageModel; use aimux_core::language_model::{LanguageModel, SupportedUrls}; use aimux_core::model_catalogue::RuntimeModel; use aimux_core::provider::{Provider, ProviderDiscovery}; -use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, combine_headers, validate_base_url, -}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, provider_headers}; -use config::{BaseUrl, ChatDialect, ChatSettings, CompatModelConfig}; +use crate::shared::{AuthScheme, Credential, ProviderHeaders}; +use config::{BaseUrl, ChatSettings, CompatModelConfig}; /// Settings of [`create_openai_compatible`] (the AI SDK's /// `OpenAICompatibleProviderSettings`). @@ -88,11 +84,9 @@ pub struct OpenAICompatibleProviderSettings { /// Rewrites each chat JSON request body once, after it is serialized and /// before it is sent. pub transform_request_body: Option, - /// Extracts metadata from chat responses and streaming chunks. - pub metadata_extractor: Option>, - /// URL patterns supported by chat models, evaluated when requested. - pub supported_urls: Option, - /// Custom chat token accounting. + /// URL patterns supported by chat models. + pub supported_urls: Option, + /// How chat token usage is read; the generic OpenAI shape by default. pub convert_usage: Option, } @@ -118,9 +112,8 @@ impl std::fmt::Debug for OpenAICompatibleProviderSettings { "transform_request_body", &self.transform_request_body.is_some(), ) - .field("metadata_extractor", &self.metadata_extractor.is_some()) .field("supported_urls", &self.supported_urls.is_some()) - .field("convert_usage", &self.convert_usage.is_some()) + .field("convert_usage", &self.convert_usage) .finish() } } @@ -138,53 +131,32 @@ pub fn create_openai_compatible( let base_url = validate_base_url(&settings.base_url)?; let credential = settings .api_key - .clone() .filter(|key| !key.is_empty()) .map_or(Credential::None, |key| { Credential::Explicit(Resolvable::Value(key)) }); - let user_headers = settings.headers.clone(); - let mut provider = OpenAICompatibleProvider::assemble(Assembly { + let defaults = ChatSettings::default(); + OpenAICompatibleProvider::assemble(Assembly { name: settings.name, base_url: BaseUrl::Fixed(base_url), - credential: credential.clone(), + credential, fixed_headers: Vec::new(), headers: settings.headers, query_params: settings.query_params, fetch: settings.fetch, transform_request_body: settings.transform_request_body, - profile: ChatProfile { + chat: ChatSettings { include_usage: settings.include_usage.unwrap_or(false), supports_structured_outputs: settings.supports_structured_outputs.unwrap_or(false), - dialect: ChatDialect::baseline(), + supported_urls: settings.supported_urls.unwrap_or_default(), + convert_usage: settings.convert_usage.unwrap_or_default(), + ..defaults }, - })?; - let dialect = Arc::make_mut(&mut provider.chat.dialect); - dialect.metadata_extractor = settings.metadata_extractor; - dialect.supported_urls = settings.supported_urls; - dialect.convert_usage = settings.convert_usage; - let headers = compatible_headers(credential, user_headers); - provider.headers = aimux_provider_utils::headers::with_user_agent_suffix_fn( - headers, - "openai-compatible", - "3.0.59", - ); - Ok(provider) -} - -/// The chat-endpoint behavior of a compatible vendor: the capability flags the -/// public settings expose, and the [`ChatDialect`] they cannot. A vendor -/// package or preset supplies its own; the public factory uses the AI SDK -/// baseline. -#[derive(Clone)] -pub(crate) struct ChatProfile { - pub include_usage: bool, - pub supports_structured_outputs: bool, - pub dialect: ChatDialect, + }) } /// Everything a compatible provider is made of. The public factory fills it -/// from the settings; vendor packages and presets fill it themselves. +/// from the settings; presets fill it from their registry row. pub(crate) struct Assembly { pub name: String, pub base_url: BaseUrl, @@ -195,7 +167,7 @@ pub(crate) struct Assembly { pub query_params: Option>, pub fetch: Option, pub transform_request_body: Option, - pub profile: ChatProfile, + pub chat: ChatSettings, } /// An OpenAI-compatible provider (the AI SDK's `OpenAICompatibleProvider`). @@ -204,58 +176,33 @@ pub struct OpenAICompatibleProvider { name: String, base_url: BaseUrl, query_params: Option>>, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, transform_request_body: Option, chat: ChatSettings, } -fn compatible_headers(credential: Credential, user: Option) -> HeadersFn { - Resolvable::from_async_fn(move || { - let credential = credential.clone(); - let user = user.clone(); - async move { - let mut headers = HeaderMapOpt::new(); - if let Some(key) = credential.secret().await?.filter(|key| !key.is_empty()) { - headers.insert("Authorization".into(), Some(format!("Bearer {key}"))); - } - Ok(match user { - Some(user) => combine_headers(&[&headers, &user]), - None => headers, - }) - } - }) -} - impl OpenAICompatibleProvider { pub(crate) fn assemble(assembly: Assembly) -> Result { - let name = assembly.name; let query_params = assembly.query_params.map(|params| { let mut pairs: Vec<(String, String)> = params.into_iter().collect(); pairs.sort(); Arc::new(pairs) }); Ok(Self { - name, + name: assembly.name, base_url: assembly.base_url, query_params, - headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( - provider_headers( - assembly.credential, - assembly.fixed_headers, - assembly.headers, - ), - "openai-compatible", - "3.0.59", - ), + headers: ProviderHeaders { + credential: assembly.credential, + scheme: AuthScheme::Bearer, + fixed: assembly.fixed_headers, + user: assembly.headers, + user_agent: Some(("openai-compatible", "3.0.59")), + }, fetch: assembly.fetch, transform_request_body: assembly.transform_request_body, - chat: ChatSettings { - include_usage: assembly.profile.include_usage, - supports_structured_outputs: assembly.profile.supports_structured_outputs, - supported_urls: SupportedUrls::default(), - dialect: Arc::new(assembly.profile.dialect), - }, + chat: assembly.chat, }) } diff --git a/aimux-providers/src/preset.rs b/aimux-providers/src/preset.rs index 17b13894c..0b9ad1152 100644 --- a/aimux-providers/src/preset.rs +++ b/aimux-providers/src/preset.rs @@ -33,9 +33,9 @@ use serde::Deserialize; use aimux_core::error::AiMuxError; use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::openai_compatible::config::{BaseUrl, ChatDialect}; +use crate::openai_compatible::config::{BaseUrl, ChatSettings}; use crate::openai_compatible::{ - Assembly, ChatProfile, OpenAICompatibleProvider, TransformRequestBody, + Assembly, ConvertUsage, OpenAICompatibleProvider, TransformRequestBody, }; use crate::shared::Credential; @@ -85,6 +85,8 @@ pub struct PresetDescriptor { /// The only max-token key the vendor accepts (`"max_tokens"` or /// `"max_completion_tokens"`); `None` sends `max_tokens`. pub max_tokens_key: Option<&'static str>, + /// How chat usage is read (`profile.convert_usage`). + pub convert_usage: ConvertUsage, pub params: &'static [ParamSpec], } @@ -232,7 +234,9 @@ fn load_registry() -> Vec { "invalid base_url_env name", ); check( - row.profile.keys().all(|key| *key == "max_tokens_key"), + row.profile + .keys() + .all(|key| matches!(*key, "max_tokens_key" | "convert_usage")), "unknown profile key", ); let max_tokens_key = row.profile.get("max_tokens_key").copied(); @@ -240,6 +244,15 @@ fn load_registry() -> Vec { max_tokens_key.is_none_or(|key| matches!(key, "max_tokens" | "max_completion_tokens")), "invalid max_tokens_key", ); + let convert_usage = match row.profile.get("convert_usage").copied() { + None => ConvertUsage::OpenAICompatible, + Some("alibaba") => ConvertUsage::Alibaba, + Some("moonshotai") => ConvertUsage::MoonshotAI, + Some(value) => panic!( + "registry row '{}': unknown convert_usage {value:?}", + row.name + ), + }; let mut params = Vec::with_capacity(row.params.len()); for param in row.params { check(identifier.is_match(param.name), "bad param name"); @@ -287,6 +300,7 @@ fn load_registry() -> Vec { auth, base_url_env: row.base_url_env, max_tokens_key, + convert_usage, params: params.leak(), })), }); @@ -295,7 +309,7 @@ fn load_registry() -> Vec { } /// Create the provider of the registry row `name`: the row's base URL, key -/// variable and dialect, with `settings` overriding individual fields. The +/// variable and chat settings, with `settings` overriding individual fields. The /// result is an ordinary [`OpenAICompatibleProvider`]; nothing is read from /// the environment here. /// @@ -360,10 +374,10 @@ fn assemble( None if descriptor.params.is_empty() && descriptor.base_url_env.is_none() => { BaseUrl::Fixed(validate_base_url(descriptor.base_url)?) } - None => { - let explicit = settings.params; - BaseUrl::Lazy(Arc::new(move || resolve_base_url(descriptor, &explicit))) - } + None => BaseUrl::Preset { + descriptor, + params: Arc::new(settings.params), + }, }; let credential = match (descriptor.auth, settings.api_key) { @@ -375,42 +389,13 @@ fn assemble( (AuthMode::None, None) => Credential::None, }; - let mut dialect = ChatDialect::baseline(); - dialect.supports_top_k = true; - dialect.max_tokens_key = descriptor.max_tokens_key; - if matches!(descriptor.name, "alibaba" | "moonshotai") { - dialect.convert_usage = Some(Arc::new(move |raw| { - let mut usage = crate::openai_compatible::chat::usage_from_raw(raw); - let Some(raw) = raw.filter(|raw| !raw.is_null()) else { - return usage; - }; - let tokens = |value: &serde_json::Value| { - value.as_u64().and_then(|value| u32::try_from(value).ok()) - }; - if descriptor.name == "alibaba" { - let details = &raw["prompt_tokens_details"]; - let cache_write = tokens(&details["cache_creation_input_tokens"]) - .or_else(|| tokens(&details["cache_write_tokens"])) - .unwrap_or(0); - usage.input_tokens.cache_write = Some(cache_write); - usage.input_tokens.no_cache = usage - .input_tokens - .no_cache - .map(|tokens| tokens.saturating_sub(cache_write)); - } else if let Some(cached) = tokens(&raw["cached_tokens"]) { - usage.input_tokens.cache_read = Some(cached); - usage.input_tokens.no_cache = usage - .input_tokens - .total - .map(|tokens| tokens.saturating_sub(cached)); - } - usage - })); - } - let profile = ChatProfile { + let chat = ChatSettings { include_usage: true, supports_structured_outputs: true, - dialect, + supports_top_k: true, + max_tokens_key: descriptor.max_tokens_key, + convert_usage: descriptor.convert_usage, + ..ChatSettings::default() }; OpenAICompatibleProvider::assemble(Assembly { @@ -422,7 +407,7 @@ fn assemble( query_params: None, fetch: settings.fetch, transform_request_body: settings.transform_request_body, - profile, + chat, }) } @@ -473,7 +458,7 @@ fn env_value(names: &[&str]) -> Option { /// The base URL a request goes to: the descriptor's base-URL variable, else /// the template expanded from `explicit`, the environment and the defaults. -fn resolve_base_url( +pub(crate) fn resolve_base_url( descriptor: &PresetDescriptor, explicit: &HashMap, ) -> Result { diff --git a/aimux-providers/src/provider_registry.json b/aimux-providers/src/provider_registry.json index 8fa5f295e..c4b0c54c2 100644 --- a/aimux-providers/src/provider_registry.json +++ b/aimux-providers/src/provider_registry.json @@ -95,7 +95,9 @@ "display": "Alibaba Cloud (DashScope)", "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1", "env_var": "ALIBABA_API_KEY", - "profile": {} + "profile": { + "convert_usage": "alibaba" + } }, { "name": "alibaba_coding_plan", @@ -1157,7 +1159,9 @@ "display": "Moonshot AI", "base_url": "https://api.moonshot.cn/v1", "env_var": "MOONSHOT_API_KEY", - "profile": {} + "profile": { + "convert_usage": "moonshotai" + } }, { "name": "moonshotai_cn", diff --git a/aimux-providers/src/shared/mod.rs b/aimux-providers/src/shared/mod.rs index 6ca274530..232105cbc 100644 --- a/aimux-providers/src/shared/mod.rs +++ b/aimux-providers/src/shared/mod.rs @@ -19,7 +19,10 @@ use std::sync::Arc; use serde_json::Value; use aimux_core::AiMuxError; -use aimux_provider_utils::{HeaderMapOpt, HeadersFn, Resolvable, combine_headers, load_api_key}; +use aimux_provider_utils::{ + HeaderMapOpt, HeadersFn, Resolvable, combine_headers, load_api_key, normalize_headers, + with_user_agent_suffix, +}; /// A host part (a Vertex location, an Azure resource name) is one DNS label: letters, digits and hyphens, not starting or /// ending with a hyphen (`isValidHostnamePart`). @@ -102,48 +105,68 @@ pub(crate) fn provider_headers( credential_headers(credential, AuthScheme::Bearer, fixed, user) } -/// Provider headers, evaluated on every request: the credential first, then -/// the fixed headers (organization, project, version, ...), then the user's -/// headers, which may override or remove any of them (case-insensitively, -/// `None` removes). +/// The provider-level request headers, as data. They are produced on every +/// request: the credential first, then the fixed headers (organization, +/// project, version, ...), then the user's headers, which may override or +/// remove any of them (case-insensitively, `None` removes), and finally the +/// `ai-sdk-{package}/{version}` user-agent suffix when one is set. +#[derive(Clone, Debug)] +pub(crate) struct ProviderHeaders { + pub credential: Credential, + pub scheme: AuthScheme, + pub fixed: Vec<(String, String)>, + pub user: Option, + /// `(package, version)` of the user-agent suffix. + pub user_agent: Option<(&'static str, &'static str)>, +} + +impl ProviderHeaders { + pub(crate) async fn resolve(&self) -> Result { + let mut layer = HeaderMapOpt::new(); + if let Some(secret) = self.credential.secret().await? { + let (name, value) = match self.scheme { + AuthScheme::Bearer => ("Authorization", format!("Bearer {secret}")), + AuthScheme::Header(name) => (name, secret), + AuthScheme::Scheme(scheme) => ("Authorization", format!("{scheme} {secret}")), + }; + layer.insert(name.to_string(), Some(value)); + } + for (name, value) in &self.fixed { + layer.insert(name.clone(), Some(value.clone())); + } + let headers = match &self.user { + Some(user) => combine_headers(&[&layer, user]), + None => layer, + }; + let Some((package, version)) = self.user_agent else { + return Ok(headers); + }; + let mut headers = normalize_headers(headers).into_iter().collect(); + with_user_agent_suffix(&mut headers, &format!("ai-sdk-{package}/{version}")); + Ok(headers + .into_iter() + .map(|(name, value)| (name, Some(value))) + .collect()) + } +} + +/// [`ProviderHeaders`] without a user-agent suffix, as a [`HeadersFn`] for +/// the packages that still take one. pub(crate) fn credential_headers( credential: Credential, scheme: AuthScheme, fixed: Vec<(String, String)>, user: Option, ) -> HeadersFn { + let headers = ProviderHeaders { + credential, + scheme, + fixed, + user, + user_agent: None, + }; Resolvable::from_async_fn(move || { - let credential = credential.clone(); - let fixed = fixed.clone(); - let user = user.clone(); - async move { - let mut layer = HeaderMapOpt::new(); - if let Some(secret) = credential.secret().await? { - match scheme { - AuthScheme::Bearer => { - layer.insert( - "Authorization".to_string(), - Some(format!("Bearer {secret}")), - ); - } - AuthScheme::Header(name) => { - layer.insert(name.to_string(), Some(secret)); - } - AuthScheme::Scheme(scheme) => { - layer.insert( - "Authorization".to_string(), - Some(format!("{scheme} {secret}")), - ); - } - } - } - for (name, value) in fixed { - layer.insert(name, Some(value)); - } - Ok(match &user { - Some(user) => combine_headers(&[&layer, user]), - None => layer, - }) - } + let headers = headers.clone(); + async move { headers.resolve().await } }) } From 12068f878adb379a150f58cc9ca249f8575bf5be Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 11:37:02 +0800 Subject: [PATCH 02/16] refactor(providers)!: id generators are plain function pointers `generate_id` settings (anthropic, google, cohere, mistral) and the tool-call tracker take `IdGenerator = fn() -> String` (the AI SDK's `IdGenerator`) instead of boxed closures: no captured state, usable in const configs. Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-provider-utils/src/generate_id.rs | 4 +++ aimux-provider-utils/src/lib.rs | 2 +- .../src/streaming_tool_call_tracker.rs | 9 ++--- .../tests/streaming_tool_call_tracker_test.rs | 34 +++++++++---------- aimux-providers/src/anthropic/mod.rs | 6 ++-- aimux-providers/src/anthropic/model.rs | 12 +++---- aimux-providers/src/anthropic/stream.rs | 16 ++++----- aimux-providers/src/cohere/mod.rs | 6 ++-- aimux-providers/src/cohere/model.rs | 7 ++-- aimux-providers/src/google/mod.rs | 6 ++-- aimux-providers/src/google/model.rs | 23 ++++++------- aimux-providers/src/mistral/mod.rs | 6 ++-- aimux-providers/src/mistral/model.rs | 13 +++---- .../tests/google_upstream_e2e_test.rs | 2 +- 14 files changed, 71 insertions(+), 75 deletions(-) diff --git a/aimux-provider-utils/src/generate_id.rs b/aimux-provider-utils/src/generate_id.rs index 70c0211d8..c6629697e 100644 --- a/aimux-provider-utils/src/generate_id.rs +++ b/aimux-provider-utils/src/generate_id.rs @@ -1,6 +1,10 @@ use std::hash::{BuildHasher, RandomState}; use std::sync::atomic::{AtomicU64, Ordering}; +/// A function that returns a fresh id (the AI SDK's `IdGenerator`). A plain +/// function pointer: it captures no state and can sit in a `const` config. +pub type IdGenerator = fn() -> String; + /// Generate a non-cryptographic, 16-character alphanumeric ID, like `generateId`. #[must_use] pub fn generate_id() -> String { diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs index dfe7b8fc6..f5b003eae 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -43,7 +43,7 @@ pub use fetch::{ Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, PinnedFetch, ReqwestFetch, default_fetch, }; -pub use generate_id::generate_id; +pub use generate_id::{IdGenerator, generate_id}; pub use get_from_api::get_from_api; pub use headers::{ HeaderMapOpt, HeadersFn, combine_headers, normalize_headers, with_user_agent_suffix, diff --git a/aimux-provider-utils/src/streaming_tool_call_tracker.rs b/aimux-provider-utils/src/streaming_tool_call_tracker.rs index fb862a71a..c4603239e 100644 --- a/aimux-provider-utils/src/streaming_tool_call_tracker.rs +++ b/aimux-provider-utils/src/streaming_tool_call_tracker.rs @@ -34,6 +34,7 @@ use aimux_core::stream_part::StreamPart; use aimux_core::tool::RawToolCall; use aimux_core::types::ProviderMetadata; +use crate::generate_id::IdGenerator; use crate::streaming_tool_call_argument_state::{ StreamingToolCallArgumentState, starts_with_structured_value, }; @@ -110,7 +111,7 @@ enum Resolution { /// Tracks the tool calls of one streamed response. See the module docs. pub struct StreamingToolCallTracker { calls: Vec, - generate_id: Box String + Send + Sync>, + generate_id: IdGenerator, type_validation: TypeValidation, } @@ -127,7 +128,7 @@ impl StreamingToolCallTracker { pub fn new() -> Self { Self { calls: Vec::new(), - generate_id: Box::new(|| FALLBACK_TOOL_CALL_ID.to_string()), + generate_id: || FALLBACK_TOOL_CALL_ID.to_string(), type_validation: TypeValidation::None, } } @@ -135,8 +136,8 @@ impl StreamingToolCallTracker { /// Set the generator for calls without a usable wire id. Blank or /// repeated outputs are turned into unique ids. #[must_use] - pub fn with_generate_id(mut self, f: impl Fn() -> String + Send + Sync + 'static) -> Self { - self.generate_id = Box::new(f); + pub fn with_generate_id(mut self, generate_id: IdGenerator) -> Self { + self.generate_id = generate_id; self } diff --git a/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs index 1687dd785..28c6d070e 100644 --- a/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs +++ b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs @@ -4,7 +4,6 @@ //! Tracker parts are projected to a local `Part` enum so the assertions stay //! byte-for-byte comparable with the upstream expectations. -use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use aimux_core::shared::provider_namespace; @@ -12,7 +11,7 @@ use aimux_core::stream_part::StreamPart; use aimux_core::tool::RawToolCall; use aimux_core::types::ProviderMetadata; use aimux_provider_utils::{ - StreamingToolCallDelta, StreamingToolCallTracker, TrackerError, TypeValidation, + IdGenerator, StreamingToolCallDelta, StreamingToolCallTracker, TrackerError, TypeValidation, }; use serde_json::json; @@ -207,17 +206,18 @@ fn tool_call(id: &str, name: &str, input: &str) -> Part { } } -/// Deterministic id generator: `prefix-1`, `prefix-2`, … and a call counter. -fn counting_generator(ids: &'static [&'static str]) -> (impl Fn() -> String, Arc) { - let calls = Arc::new(AtomicUsize::new(0)); - let counter = Arc::clone(&calls); - ( - move || { - let n = counter.fetch_add(1, Ordering::SeqCst); - ids[n.min(ids.len() - 1)].to_string() - }, - calls, - ) +/// A deterministic id generator returning the given ids in order (the last +/// one repeats) and its call counter. Each use gets its own counter. +macro_rules! counting_generator { + ($($id:literal),+ $(,)?) => {{ + static CALLS: AtomicUsize = AtomicUsize::new(0); + fn generate() -> String { + const IDS: &[&str] = &[$($id),+]; + let n = CALLS.fetch_add(1, Ordering::SeqCst); + IDS[n.min(IDS.len() - 1)].to_string() + } + (generate as IdGenerator, &CALLS) + }}; } mod process_delta { @@ -521,7 +521,7 @@ mod process_delta { #[test] fn keeps_id_less_calls_distinct_when_an_index_is_reused_and_type_is_omitted() { - let (generate, _) = counting_generator(&["generated-1", "generated-2", "generated-3"]); + let (generate, _) = counting_generator!("generated-1", "generated-2", "generated-3"); let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); h.delta(delta( @@ -672,7 +672,7 @@ mod process_delta { #[test] fn generates_unique_ids_for_blank_and_repeated_ids() { - let (generate, _) = counting_generator(&["generated-1", "generated-2"]); + let (generate, _) = counting_generator!("generated-1", "generated-2"); let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); h.delta(start(0, "", "read_file", "{}")).unwrap(); @@ -734,7 +734,7 @@ mod process_delta { #[test] fn creates_bounded_unique_ids_when_generate_id_returns_duplicates() { - let (generate, calls) = counting_generator(&["generated-id"]); + let (generate, calls) = counting_generator!("generated-id"); let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); for (index, name) in [(0, "first"), (1, "second"), (2, "third")] { @@ -759,7 +759,7 @@ mod process_delta { #[test] fn creates_usable_ids_when_generate_id_returns_blank_values() { - let (generate, calls) = counting_generator(&[" "]); + let (generate, calls) = counting_generator!(" "); let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); for (index, name) in [(0, "first"), (1, "second")] { diff --git a/aimux-providers/src/anthropic/mod.rs b/aimux-providers/src/anthropic/mod.rs index 8351f7629..511665a0c 100644 --- a/aimux-providers/src/anthropic/mod.rs +++ b/aimux-providers/src/anthropic/mod.rs @@ -132,7 +132,7 @@ pub struct AnthropicProviderSettings { /// `None` uses the process default, resolved per request. pub fetch: Option, /// Generates identifiers for returned sources. - pub generate_id: Option String + Send + Sync>>, + pub generate_id: Option, } impl std::fmt::Debug for AnthropicProviderSettings { @@ -230,7 +230,7 @@ pub struct AnthropicProvider { headers: HeadersFn, fetch: Option, supported_urls: SupportedUrls, - generate_id: Option String + Send + Sync>>, + generate_id: Option, } impl AnthropicProvider { @@ -265,7 +265,7 @@ impl AnthropicProvider { model_id.to_string(), self.model_config(self.name.clone()), ) - .with_generate_id(self.generate_id.clone()) + .with_generate_id(self.generate_id) } /// An alias for the Messages model. diff --git a/aimux-providers/src/anthropic/model.rs b/aimux-providers/src/anthropic/model.rs index bc157bbae..e9f0f5fba 100644 --- a/aimux-providers/src/anthropic/model.rs +++ b/aimux-providers/src/anthropic/model.rs @@ -6,8 +6,6 @@ //! and hands the result to the shared [`super::stream`] core, which sends it //! and parses the response or the SSE stream. -use std::sync::Arc; - use async_trait::async_trait; use serde_json::Value; @@ -38,7 +36,7 @@ struct PreparedCall { pub struct AnthropicMessagesModel { model_id: String, config: AnthropicModelConfig, - generate_id: Arc String + Send + Sync>, + generate_id: aimux_provider_utils::IdGenerator, } impl AnthropicMessagesModel { @@ -47,14 +45,14 @@ impl AnthropicMessagesModel { Self { model_id, config, - generate_id: Arc::new(aimux_provider_utils::generate_id), + generate_id: aimux_provider_utils::generate_id, } } #[must_use] pub(crate) fn with_generate_id( mut self, - generate_id: Option String + Send + Sync>>, + generate_id: Option, ) -> Self { if let Some(generate_id) = generate_id { self.generate_id = generate_id; @@ -160,7 +158,7 @@ impl LanguageModel for AnthropicMessagesModel { self.uses_custom_options(options), call.uses_json_response_tool, Self::citation_documents(options), - self.generate_id.clone(), + self.generate_id, ) .await } @@ -176,7 +174,7 @@ impl LanguageModel for AnthropicMessagesModel { self.uses_custom_options(options), call.uses_json_response_tool, Self::citation_documents(options), - self.generate_id.clone(), + self.generate_id, ) .await } diff --git a/aimux-providers/src/anthropic/stream.rs b/aimux-providers/src/anthropic/stream.rs index 1449bbd70..1a5b0e063 100644 --- a/aimux-providers/src/anthropic/stream.rs +++ b/aimux-providers/src/anthropic/stream.rs @@ -364,7 +364,7 @@ pub(crate) fn stream_parts_for_result_block( names: &ToolNameMapping, mcp_tool_calls: &HashMap, server_tool_calls: &HashMap, - generate_id: &(dyn Fn() -> String + Send + Sync), + generate_id: aimux_provider_utils::IdGenerator, ) -> Vec { let caller = match block { ContentBlock::WebSearchToolResult { caller, .. } @@ -534,7 +534,7 @@ pub(crate) struct CitationDocument { fn citation_source( citation: &Value, documents: &[CitationDocument], - generate_id: &(dyn Fn() -> String + Send + Sync), + generate_id: aimux_provider_utils::IdGenerator, ) -> Option { match citation["type"].as_str()? { "web_search_result_location" => Some(Source::Url { @@ -617,7 +617,7 @@ pub(crate) fn parse_anthropic_content( names: &ToolNameMapping, uses_json_response_tool: bool, mut citation_documents: Vec, - generate_id: &(dyn Fn() -> String + Send + Sync), + generate_id: aimux_provider_utils::IdGenerator, ) -> Vec { let options_name = CANONICAL; let mut content = Vec::new(); @@ -976,7 +976,7 @@ pub(crate) async fn anthropic_generate_core( used_custom_options_key: bool, uses_json_response_tool: bool, citation_documents: Vec, - generate_id: std::sync::Arc String + Send + Sync>, + generate_id: aimux_provider_utils::IdGenerator, ) -> Result { let resp = aimux_provider_utils::post_json_to_api( request, @@ -1000,7 +1000,7 @@ pub(crate) async fn anthropic_generate_core( tool_names, uses_json_response_tool, citation_documents, - generate_id.as_ref(), + generate_id, ); let mut finish_reason = data @@ -1093,7 +1093,7 @@ pub(crate) async fn anthropic_stream_core( used_custom_options_key: bool, uses_json_response_tool: bool, mut citation_documents: Vec, - generate_id: std::sync::Arc String + Send + Sync>, + generate_id: aimux_provider_utils::IdGenerator, ) -> Result { let endpoint = request.url.clone(); let options_name = config.provider_options_name.clone(); @@ -1342,7 +1342,7 @@ pub(crate) async fn anthropic_stream_core( &tool_names, &mcp_tool_calls, &server_tool_calls, - generate_id.as_ref(), + generate_id, ) { yield Ok(part); } @@ -1364,7 +1364,7 @@ pub(crate) async fn anthropic_stream_core( if let Some(BlockState::Text { citations }) = blocks.get_mut(&index) { citations.push(citation.clone()); } - if let Some(source) = citation_source(&citation, &citation_documents, generate_id.as_ref()) { + if let Some(source) = citation_source(&citation, &citation_documents, generate_id) { yield Ok(StreamPart::Source(source)); } } diff --git a/aimux-providers/src/cohere/mod.rs b/aimux-providers/src/cohere/mod.rs index eadfd8302..7c8af5869 100644 --- a/aimux-providers/src/cohere/mod.rs +++ b/aimux-providers/src/cohere/mod.rs @@ -76,7 +76,7 @@ pub struct CohereProviderSettings { /// including `Authorization`. Per-call headers win over these. pub headers: Option, /// Generates a unique identifier for each citation. - pub generate_id: Option String + Send + Sync>>, + pub generate_id: Option, /// The transport: a mock, a signing decorator, a proxy-aware client. /// `None` uses the process default, resolved per request. pub fetch: Option, @@ -147,7 +147,7 @@ pub fn cohere() -> &'static CohereProvider { /// A Cohere provider (the AI SDK's `CohereProvider`). Cheap to clone the /// models out of; it holds no HTTP client. pub struct CohereProvider { - generate_id: Option String + Send + Sync>>, + generate_id: Option, base_url: String, headers: HeadersFn, fetch: Option, @@ -168,7 +168,7 @@ impl CohereProvider { #[must_use] pub fn chat(&self, model_id: &str) -> CohereModel { CohereModel::from_config(model_id.to_string(), self.model_config("chat")) - .with_generate_id(self.generate_id.clone()) + .with_generate_id(self.generate_id) } /// An embedding model (e.g. `"embed-english-v3.0"`); `provider()` is diff --git a/aimux-providers/src/cohere/model.rs b/aimux-providers/src/cohere/model.rs index 9c3a1e825..48d41fa71 100644 --- a/aimux-providers/src/cohere/model.rs +++ b/aimux-providers/src/cohere/model.rs @@ -20,13 +20,12 @@ use crate::shared::EndpointConfig; use super::convert::{build_request_body, parse_finish_reason}; use super::types::{ChatResponse, StreamEvent, UsageResponse}; -use std::sync::Arc; /// A Cohere language model. pub struct CohereModel { model_id: String, config: EndpointConfig, - generate_id: Arc String + Send + Sync>, + generate_id: aimux_provider_utils::IdGenerator, } impl CohereModel { @@ -34,12 +33,12 @@ impl CohereModel { Self { model_id, config, - generate_id: Arc::new(aimux_provider_utils::generate_id), + generate_id: aimux_provider_utils::generate_id, } } pub(crate) fn with_generate_id( mut self, - generate_id: Option String + Send + Sync>>, + generate_id: Option, ) -> Self { if let Some(generate_id) = generate_id { self.generate_id = generate_id; diff --git a/aimux-providers/src/google/mod.rs b/aimux-providers/src/google/mod.rs index 4368d814b..1f0977495 100644 --- a/aimux-providers/src/google/mod.rs +++ b/aimux-providers/src/google/mod.rs @@ -201,7 +201,7 @@ pub struct GoogleProviderSettings { /// `None` uses the process default, resolved per request. pub fetch: Option, /// Generates IDs for tool calls and sources that have no provider ID. - pub generate_id: Option String + Send + Sync>>, + pub generate_id: Option, /// Opens the WebSocket of live transcription (the AI SDK's `webSocket`). /// `None` uses the built-in client. #[cfg(feature = "realtime")] @@ -282,7 +282,7 @@ pub struct GoogleProvider { base_url: String, headers: HeadersFn, fetch: Option, - generate_id: Option String + Send + Sync>>, + generate_id: Option, #[cfg(feature = "realtime")] web_socket: Option>, } @@ -316,7 +316,7 @@ impl GoogleProvider { #[must_use] pub fn chat(&self, model_id: &str) -> GoogleModel { GoogleModel::from_config(model_id.to_string(), self.model_config(self.name.clone())) - .with_generate_id(self.generate_id.clone()) + .with_generate_id(self.generate_id) } /// Alias of [`chat`](Self::chat), matching `generativeAI`. diff --git a/aimux-providers/src/google/model.rs b/aimux-providers/src/google/model.rs index 9ef7e7d09..2a264b3ec 100644 --- a/aimux-providers/src/google/model.rs +++ b/aimux-providers/src/google/model.rs @@ -35,7 +35,7 @@ use crate::shared::EndpointConfig; pub struct GoogleModel { model_id: String, config: EndpointConfig, - generate_id: std::sync::Arc String + Send + Sync>, + generate_id: aimux_provider_utils::IdGenerator, } impl GoogleModel { @@ -43,13 +43,13 @@ impl GoogleModel { Self { model_id, config, - generate_id: std::sync::Arc::new(aimux_provider_utils::generate_id), + generate_id: aimux_provider_utils::generate_id, } } pub(crate) fn with_generate_id( mut self, - generate_id: Option String + Send + Sync>>, + generate_id: Option, ) -> Self { if let Some(generate_id) = generate_id { self.generate_id = generate_id; @@ -110,11 +110,8 @@ impl LanguageModel for GoogleModel { let candidate = data.candidates.into_iter().next().unwrap_or_default(); let block_reason = confirmed_prompt_block_reason(data.prompt_feedback.as_ref()); - let (content, has_tool_calls) = extract_content_from_candidate( - &candidate, - &code_execution_tool_name, - &*self.generate_id, - ); + let (content, has_tool_calls) = + extract_content_from_candidate(&candidate, &code_execution_tool_name, self.generate_id); let finish_reason = if candidate.finish_reason.is_none() && block_reason.is_some() { FinishReason { @@ -201,7 +198,7 @@ impl LanguageModel for GoogleModel { let stream_request_body = body.clone(); let stream_error_headers = response_headers.clone(); - let generate_id = self.generate_id.clone(); + let generate_id = self.generate_id; let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings: tool_warnings }); @@ -454,7 +451,7 @@ impl LanguageModel for GoogleModel { .and_then(|v| v.as_str()) .filter(|id| !id.is_empty()) .map(std::string::ToString::to_string) - .unwrap_or_else(|| generate_id()); + .unwrap_or_else(generate_id); last_server_tool_call_id = Some(id.clone()); let args = tc.get("args").cloned().unwrap_or(json!({})); let server_meta = server_tool_metadata( @@ -482,7 +479,7 @@ impl LanguageModel for GoogleModel { .map(std::string::ToString::to_string) }) .filter(|id| !id.is_empty()) - .unwrap_or_else(|| generate_id()); + .unwrap_or_else(generate_id); let response = tr.get("response").cloned().unwrap_or(json!({})); let server_meta = server_tool_metadata( @@ -539,7 +536,7 @@ impl LanguageModel for GoogleModel { .filter(|id| !id.is_empty()) .filter(|id| !id.is_empty()) .map(std::string::ToString::to_string) - .unwrap_or_else(|| generate_id()); + .unwrap_or_else(generate_id); let args = fc.get("args").filter(|value| !value.is_null()).cloned().unwrap_or(json!({})); yield Ok(StreamPart::ToolInputStart { @@ -674,7 +671,7 @@ fn server_tool_metadata( fn extract_content_from_candidate( candidate: &Candidate, code_execution_tool_name: &str, - generate_id: &dyn Fn() -> String, + generate_id: aimux_provider_utils::IdGenerator, ) -> (Vec, bool) { let mut content = Vec::new(); let mut has_tool_calls = false; diff --git a/aimux-providers/src/mistral/mod.rs b/aimux-providers/src/mistral/mod.rs index 10ad2e89c..57283ff0f 100644 --- a/aimux-providers/src/mistral/mod.rs +++ b/aimux-providers/src/mistral/mod.rs @@ -130,7 +130,7 @@ pub struct MistralProviderSettings { /// `None` uses the process default, resolved per request. pub fetch: Option, /// Generates identifiers for reasoning and streamed tool calls. - pub generate_id: Option String + Send + Sync>>, + pub generate_id: Option, } impl std::fmt::Debug for MistralProviderSettings { @@ -199,7 +199,7 @@ pub struct MistralProvider { base_url: String, headers: HeadersFn, fetch: Option, - generate_id: Option String + Send + Sync>>, + generate_id: Option, } impl MistralProvider { @@ -220,7 +220,7 @@ impl MistralProvider { model_id.to_string(), self.model_config("chat") .with_supported_urls(Arc::new(|_| chat_supported_urls())), - self.generate_id.clone(), + self.generate_id, ) } diff --git a/aimux-providers/src/mistral/model.rs b/aimux-providers/src/mistral/model.rs index 703ce3c64..6d3283573 100644 --- a/aimux-providers/src/mistral/model.rs +++ b/aimux-providers/src/mistral/model.rs @@ -13,7 +13,6 @@ use aimux_core::tool::RawToolCall; use async_trait::async_trait; use futures::StreamExt; use serde_json::Value; -use std::sync::Arc; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; @@ -33,20 +32,19 @@ use super::types::{ChatCompletionResponse, StreamChunk, UsageResponse}; pub struct MistralModel { model_id: String, config: EndpointConfig, - generate_id: Arc String + Send + Sync>, + generate_id: aimux_provider_utils::IdGenerator, } impl MistralModel { pub(crate) fn from_config( model_id: String, config: EndpointConfig, - generate_id_fn: Option String + Send + Sync>>, + generate_id_fn: Option, ) -> Self { Self { model_id, config, - generate_id: generate_id_fn - .unwrap_or_else(|| Arc::new(aimux_provider_utils::generate_id)), + generate_id: generate_id_fn.unwrap_or(aimux_provider_utils::generate_id), } } } @@ -296,15 +294,14 @@ impl LanguageModel for MistralModel { let stream_error_body = body.clone(); let stream_response_headers = response_headers.clone(); - let generate_id_fn = self.generate_id.clone(); + let generate_id_fn = self.generate_id; let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings }); let text_id = 0usize; let mut text_started = false; let mut reasoning_started = false; - let tool_generate_id = generate_id_fn.clone(); - let mut tool_calls = StreamingToolCallTracker::new().with_generate_id(move || tool_generate_id()); + let mut tool_calls = StreamingToolCallTracker::new().with_generate_id(generate_id_fn); let mut tool_parts = Vec::new(); let mut reasoning_id: Option = None; let mut final_usage = Usage::default(); diff --git a/aimux-providers/tests/google_upstream_e2e_test.rs b/aimux-providers/tests/google_upstream_e2e_test.rs index 96e84e729..e3b0c996b 100644 --- a/aimux-providers/tests/google_upstream_e2e_test.rs +++ b/aimux-providers/tests/google_upstream_e2e_test.rs @@ -26,7 +26,7 @@ fn settings(mock: &Arc) -> GoogleProviderSettings { base_url: Some(BASE.into()), api_key: Some("test-api-key".into()), fetch: Some(mock.transport()), - generate_id: Some(Arc::new(|| "test-id".into())), + generate_id: Some(|| "test-id".into()), ..Default::default() } } From ce88c051521288d70006154f1122bd94b2f9bfc7 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 11:50:00 +0800 Subject: [PATCH 03/16] refactor(providers)!: remove the transform_request_body closure setting Provider configuration is data; the request-body rewrite closure is gone from every provider's settings, PresetSettings and the internal configs. The OpenAI-compatible provider still passes unknown providerOptions fields of its namespace through to the body. Callers that did not set a transform see identical requests. Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-providers/src/anthropic/config.rs | 17 +-- aimux-providers/src/anthropic/mod.rs | 1 - aimux-providers/src/anthropic/model.rs | 7 +- aimux-providers/src/anthropic_aws/mod.rs | 12 -- aimux-providers/src/assemblyai.rs | 1 - aimux-providers/src/azure/mod.rs | 1 - aimux-providers/src/bedrock/embedding.rs | 1 - aimux-providers/src/bedrock/image.rs | 2 +- aimux-providers/src/bedrock/mod.rs | 12 +- aimux-providers/src/bedrock/model.rs | 2 - aimux-providers/src/bedrock/reranking.rs | 1 - aimux-providers/src/black_forest_labs.rs | 1 - aimux-providers/src/cartesia.rs | 1 - aimux-providers/src/codex.rs | 14 +-- aimux-providers/src/cohere/embedding.rs | 2 +- aimux-providers/src/cohere/mod.rs | 1 - aimux-providers/src/cohere/model.rs | 4 +- aimux-providers/src/cohere/reranking.rs | 1 - aimux-providers/src/dataforseo.rs | 1 - aimux-providers/src/deepgram.rs | 1 - aimux-providers/src/deepseek/mod.rs | 1 - aimux-providers/src/deepseek/model.rs | 11 +- aimux-providers/src/default_providers.rs | 106 ++++++++---------- aimux-providers/src/elevenlabs.rs | 18 +-- aimux-providers/src/exa_ai.rs | 1 - aimux-providers/src/fal.rs | 1 - aimux-providers/src/firecrawl.rs | 1 - aimux-providers/src/gladia.rs | 1 - aimux-providers/src/google/embedding.rs | 2 +- aimux-providers/src/google/files.rs | 2 +- aimux-providers/src/google/image.rs | 2 +- aimux-providers/src/google/mod.rs | 1 - aimux-providers/src/google/model.rs | 2 - aimux-providers/src/google/speech.rs | 2 +- aimux-providers/src/google/transcription.rs | 2 +- aimux-providers/src/google/video.rs | 2 +- aimux-providers/src/google_pse.rs | 1 - aimux-providers/src/groq/mod.rs | 1 - aimux-providers/src/groq/model.rs | 4 +- aimux-providers/src/huggingface.rs | 12 -- aimux-providers/src/huggingface/responses.rs | 4 +- aimux-providers/src/hume.rs | 1 - aimux-providers/src/jina_ai.rs | 1 - aimux-providers/src/klingai.rs | 1 - aimux-providers/src/linkup.rs | 1 - aimux-providers/src/lmnt.rs | 1 - aimux-providers/src/luma.rs | 1 - aimux-providers/src/mistral/embedding.rs | 2 +- aimux-providers/src/mistral/mod.rs | 1 - aimux-providers/src/mistral/model.rs | 4 +- aimux-providers/src/open_responses.rs | 16 +-- aimux-providers/src/openai/config.rs | 16 --- aimux-providers/src/openai/embedding.rs | 2 +- aimux-providers/src/openai/image.rs | 12 +- aimux-providers/src/openai/mod.rs | 1 - aimux-providers/src/openai/model.rs | 36 +----- aimux-providers/src/openai/responses/mod.rs | 1 - aimux-providers/src/openai/speech.rs | 2 +- aimux-providers/src/openai_compatible/chat.rs | 11 +- .../src/openai_compatible/completion.rs | 12 -- .../src/openai_compatible/config.rs | 11 -- .../src/openai_compatible/convert.rs | 2 +- aimux-providers/src/openai_compatible/mod.rs | 21 +--- aimux-providers/src/parallel_ai.rs | 1 - aimux-providers/src/preset.rs | 11 +- aimux-providers/src/prodia.rs | 1 - aimux-providers/src/recraft.rs | 1 - aimux-providers/src/replicate.rs | 1 - aimux-providers/src/revai.rs | 1 - aimux-providers/src/runwayml.rs | 1 - aimux-providers/src/serper.rs | 1 - aimux-providers/src/shared/exchange.rs | 18 --- aimux-providers/src/shared/mod.rs | 8 -- aimux-providers/src/stability.rs | 1 - aimux-providers/src/tavily.rs | 1 - aimux-providers/src/tinyfish.rs | 1 - aimux-providers/src/vertex/anthropic_model.rs | 15 +-- .../src/vertex/anthropic_provider.rs | 1 - aimux-providers/src/vertex/embedding.rs | 4 +- .../src/vertex/gemini_transcription.rs | 2 +- aimux-providers/src/vertex/image.rs | 2 +- aimux-providers/src/vertex/mod.rs | 12 +- aimux-providers/src/vertex/model.rs | 2 - aimux-providers/src/vertex/transcription.rs | 2 +- aimux-providers/src/vertex/video.rs | 4 +- aimux-providers/src/voyage/embedding.rs | 2 +- aimux-providers/src/voyage/mod.rs | 11 -- aimux-providers/src/voyage/reranking.rs | 1 - aimux-providers/src/xai/mod.rs | 1 - aimux-providers/src/xai/responses/mod.rs | 4 +- aimux-providers/src/you_com.rs | 1 - 91 files changed, 119 insertions(+), 408 deletions(-) diff --git a/aimux-providers/src/anthropic/config.rs b/aimux-providers/src/anthropic/config.rs index 2e72a921b..142fa5757 100644 --- a/aimux-providers/src/anthropic/config.rs +++ b/aimux-providers/src/anthropic/config.rs @@ -23,7 +23,6 @@ use aimux_provider_utils::{ }; use super::convert::RequestProfile; -pub use crate::shared::TransformRequestBody; /// Maps an endpoint path (`"/messages"`) to the full request URL. pub(crate) type UrlFn = Arc String + Send + Sync>; @@ -52,8 +51,7 @@ pub(crate) struct AnthropicModelHooks { /// The messages URL when it depends on the model and the mode. `None`: /// `url("/messages")`. pub(crate) request_url: Option, - /// Applied to the finished request body before the provider's own - /// `transform_request_body`. + /// Applied to the finished request body before it is sent. pub(crate) prepare_body: Option, /// Parses a failed response. pub(crate) failed_response_handler: ErrorHandlerFn, @@ -92,8 +90,6 @@ pub(crate) struct AnthropicModelConfig { pub(crate) fetch: Option, /// URLs the model fetches itself. Empty: the caller downloads them. pub(crate) supported_urls: SupportedUrls, - /// Provider-level request-body rewrite. - pub(crate) transform_request_body: Option, /// The origin credentialed headers may be sent to. Never read for URLs. pub(crate) base_url: String, /// The providerOptions key read in addition to `anthropic` and written @@ -216,13 +212,10 @@ impl AnthropicModelConfig { request } - /// Run the host's body preparation, then the provider-level rewrite. - pub(crate) fn transform_body(&self, mut body: Value) -> Value { - if let Some(prepare) = &self.hooks.prepare_body { - body = prepare(body); - } - match &self.transform_request_body { - Some(transform) => transform(body), + /// Run the host's body preparation, when there is one. + pub(crate) fn prepare_body(&self, body: Value) -> Value { + match &self.hooks.prepare_body { + Some(prepare) => prepare(body), None => body, } } diff --git a/aimux-providers/src/anthropic/mod.rs b/aimux-providers/src/anthropic/mod.rs index 511665a0c..4fe28093c 100644 --- a/aimux-providers/src/anthropic/mod.rs +++ b/aimux-providers/src/anthropic/mod.rs @@ -243,7 +243,6 @@ impl AnthropicProvider { headers: self.headers.clone(), fetch: self.fetch.clone(), supported_urls: self.supported_urls.clone(), - transform_request_body: None, base_url: self.base_url.clone(), provider_options_name: options::options_name_of(&self.name), hooks: AnthropicModelHooks::default(), diff --git a/aimux-providers/src/anthropic/model.rs b/aimux-providers/src/anthropic/model.rs index e9f0f5fba..a082d3558 100644 --- a/aimux-providers/src/anthropic/model.rs +++ b/aimux-providers/src/anthropic/model.rs @@ -103,9 +103,8 @@ impl AnthropicMessagesModel { .is_some_and(|value| value.contains_key(&self.config.provider_options_name)) } - /// Build the request for one call: the body (host preparation and the - /// provider's `transform_request_body` applied), its warnings, and the - /// URL/headers/transport the exchange goes through. + /// Build the request for one call: the body (host preparation applied), + /// its warnings, and the URL/headers/transport the exchange goes through. async fn prepare( &self, options: &CallOptions, @@ -114,7 +113,7 @@ impl AnthropicMessagesModel { let config = self.config.resolved().await?; let built = build_request_body_for(&self.model_id, options, stream, &config.request_profile())?; - let body = config.transform_body(built.body); + let body = config.prepare_body(built.body); let headers = config .request_headers(options.headers.as_ref(), &built.betas) .await?; diff --git a/aimux-providers/src/anthropic_aws/mod.rs b/aimux-providers/src/anthropic_aws/mod.rs index 12bca8694..1bb90aab9 100644 --- a/aimux-providers/src/anthropic_aws/mod.rs +++ b/aimux-providers/src/anthropic_aws/mod.rs @@ -32,8 +32,6 @@ use crate::anthropic::AnthropicMessagesModel; use crate::anthropic::config::{AnthropicModelConfig, AnthropicModelHooks}; use crate::anthropic::options::{CANONICAL, options_name_of}; -pub use crate::shared::TransformRequestBody; - /// The region used when the settings name none and SigV4 credentials do not /// carry one. const DEFAULT_REGION: &str = "us-east-1"; @@ -86,9 +84,6 @@ pub struct AnthropicAwsProviderSettings { /// The transport SigV4 signing (when used) wraps. `None` uses the process /// default, resolved per request. pub fetch: Option, - /// Rewrites every JSON request body once, after it is serialized and - /// before it is sent (and before it is signed). - pub transform_request_body: Option, } impl std::fmt::Debug for AnthropicAwsProviderSettings { @@ -105,10 +100,6 @@ impl std::fmt::Debug for AnthropicAwsProviderSettings { ) .field("name", &self.name) .field("fetch", &self.fetch.is_some()) - .field( - "transform_request_body", - &self.transform_request_body.is_some(), - ) .finish() } } @@ -176,7 +167,6 @@ pub fn create_anthropic_aws( base_url, headers, fetch, - transform_request_body: settings.transform_request_body, }) } @@ -199,7 +189,6 @@ pub struct AnthropicAwsProvider { base_url: String, headers: HeadersFn, fetch: Option, - transform_request_body: Option, } impl AnthropicAwsProvider { @@ -211,7 +200,6 @@ impl AnthropicAwsProvider { headers: self.headers.clone(), fetch: self.fetch.clone(), supported_urls: SupportedUrls::default(), - transform_request_body: self.transform_request_body.clone(), base_url: self.base_url.clone(), provider_options_name: self.provider_options_name.clone(), hooks: AnthropicModelHooks::default(), diff --git a/aimux-providers/src/assemblyai.rs b/aimux-providers/src/assemblyai.rs index 5b825cb59..143c91925 100644 --- a/aimux-providers/src/assemblyai.rs +++ b/aimux-providers/src/assemblyai.rs @@ -179,7 +179,6 @@ impl AssemblyAIProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/azure/mod.rs b/aimux-providers/src/azure/mod.rs index 04fbe7eb7..cc66a90a8 100644 --- a/aimux-providers/src/azure/mod.rs +++ b/aimux-providers/src/azure/mod.rs @@ -414,7 +414,6 @@ impl AzureOpenAIProvider { token_provider: None, fetch: self.fetch.clone(), supported_urls: crate::openai::config::supported_urls(provider), - transform_request_body: None, responses: ResponsesProfile::default(), chat_options: options::CHAT_OPTIONS, } diff --git a/aimux-providers/src/bedrock/embedding.rs b/aimux-providers/src/bedrock/embedding.rs index 1f64f6bec..bcf688094 100644 --- a/aimux-providers/src/bedrock/embedding.rs +++ b/aimux-providers/src/bedrock/embedding.rs @@ -181,7 +181,6 @@ impl EmbeddingModel for BedrockEmbeddingModel { }; let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body); let body_str = serde_json::to_string(&body).unwrap_or_default(); let url = exchange.url(&format!( "/model/{}/invoke", diff --git a/aimux-providers/src/bedrock/image.rs b/aimux-providers/src/bedrock/image.rs index 04f0cbda0..1023bde1c 100644 --- a/aimux-providers/src/bedrock/image.rs +++ b/aimux-providers/src/bedrock/image.rs @@ -326,7 +326,7 @@ impl ImageModel for BedrockImageModel { let current_date = chrono::Utc::now().to_rfc3339(); let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(Value::Object(args)); + let body = Value::Object(args); let body_str = serde_json::to_string(&body).map_err(|e| AiMuxError::JsonParse(e.to_string()))?; let url = exchange.url(&format!( diff --git a/aimux-providers/src/bedrock/mod.rs b/aimux-providers/src/bedrock/mod.rs index 3e838a57d..713055e6c 100644 --- a/aimux-providers/src/bedrock/mod.rs +++ b/aimux-providers/src/bedrock/mod.rs @@ -56,7 +56,7 @@ use aimux_provider_utils::{ load_setting, validate_base_url, without_trailing_slash, }; -use crate::shared::{Endpoint, EndpointConfig, TransformRequestBody, is_valid_hostname_part}; +use crate::shared::{Endpoint, EndpointConfig, is_valid_hostname_part}; pub use embedding::BedrockEmbeddingModel; pub use image::BedrockImageModel; @@ -199,9 +199,6 @@ pub struct AmazonBedrockProviderSettings { /// credentials. Its `region` is used only when neither `region` nor /// `AWS_REGION` names one. pub credential_provider: Option>, - /// Rewrites every JSON request body once, after it is serialized and - /// before it is sent (and before it is signed). - pub transform_request_body: Option, } impl std::fmt::Debug for AmazonBedrockProviderSettings { @@ -220,10 +217,6 @@ impl std::fmt::Debug for AmazonBedrockProviderSettings { ) .field("fetch", &self.fetch.is_some()) .field("credential_provider", &self.credential_provider.is_some()) - .field( - "transform_request_body", - &self.transform_request_body.is_some(), - ) .finish() } } @@ -453,7 +446,6 @@ pub fn create_amazon_bedrock( "5.0.100", ), fetch, - transform_request_body: settings.transform_request_body, }) } @@ -476,7 +468,6 @@ pub struct AmazonBedrockProvider { base_url: Option, headers: HeadersFn, fetch: FetchFunction, - transform_request_body: Option, } impl AmazonBedrockProvider { @@ -512,7 +503,6 @@ impl AmazonBedrockProvider { }), fetch: Some(self.fetch.clone()), supported_urls: Arc::new(|_| SupportedUrls::default()), - transform_request_body: self.transform_request_body.clone(), } } diff --git a/aimux-providers/src/bedrock/model.rs b/aimux-providers/src/bedrock/model.rs index 3e40090b5..ea9534adf 100644 --- a/aimux-providers/src/bedrock/model.rs +++ b/aimux-providers/src/bedrock/model.rs @@ -148,7 +148,6 @@ impl LanguageModel for BedrockModel { let (body, warnings, uses_json_instruction, uses_json_tool) = build_request_body_checked(&self.model_id, options, self.model_family.as_deref())?; let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body); let body_str = serde_json::to_string(&body).unwrap_or_default(); let url = exchange.url(&self.path(false)); let resp = aimux_provider_utils::post_to_api( @@ -253,7 +252,6 @@ impl LanguageModel for BedrockModel { let (body, warnings, uses_json_instruction, uses_json_tool) = build_request_body_checked(&self.model_id, options, self.model_family.as_deref())?; let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body); let body_str = serde_json::to_string(&body).unwrap_or_default(); let url = exchange.url(&self.path(true)); let error_url = url.clone(); diff --git a/aimux-providers/src/bedrock/reranking.rs b/aimux-providers/src/bedrock/reranking.rs index 6bc9d2b95..49ad3f157 100644 --- a/aimux-providers/src/bedrock/reranking.rs +++ b/aimux-providers/src/bedrock/reranking.rs @@ -171,7 +171,6 @@ impl RerankingModel for BedrockRerankingModel { } let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body); let body_str = serde_json::to_string(&body).unwrap_or_default(); let url = exchange.url("/rerank"); diff --git a/aimux-providers/src/black_forest_labs.rs b/aimux-providers/src/black_forest_labs.rs index be3b50f6e..e397277f1 100644 --- a/aimux-providers/src/black_forest_labs.rs +++ b/aimux-providers/src/black_forest_labs.rs @@ -192,7 +192,6 @@ impl BlackForestLabsProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/cartesia.rs b/aimux-providers/src/cartesia.rs index b760f4eaf..57ae2d463 100644 --- a/aimux-providers/src/cartesia.rs +++ b/aimux-providers/src/cartesia.rs @@ -172,7 +172,6 @@ impl CartesiaProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/codex.rs b/aimux-providers/src/codex.rs index 31b08c679..02a6b89ac 100644 --- a/aimux-providers/src/codex.rs +++ b/aimux-providers/src/codex.rs @@ -59,7 +59,6 @@ use crate::openai::responses::responses_convert::{ use crate::openai::responses::{ OpenAIResponsesModel, ResponsesNamespace, build_responses_request_body, }; -pub use crate::shared::TransformRequestBody; use crate::shared::{AuthScheme, Credential, credential_headers}; /// Default API-key base URL (official OpenAI Responses endpoint). @@ -134,10 +133,6 @@ pub struct CodexProviderSettings { /// The transport: a mock, a signing decorator, a proxy-aware client. /// `None` uses the process default, resolved per request. pub fetch: Option, - /// Rewrites every JSON request body once, after it is serialized (and - /// after the `store: false` rule of the subscription channel) and before it - /// is sent. - pub transform_request_body: Option, } impl std::fmt::Debug for CodexProviderSettings { @@ -153,10 +148,6 @@ impl std::fmt::Debug for CodexProviderSettings { .field("name", &self.name) .field("originator", &self.originator) .field("fetch", &self.fetch.is_some()) - .field( - "transform_request_body", - &self.transform_request_body.is_some(), - ) .finish() } } @@ -211,7 +202,6 @@ pub fn create_codex(settings: CodexProviderSettings) -> Result, - transform_request_body: Option, } impl CodexProvider { @@ -245,7 +234,6 @@ impl CodexProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - self.transform_request_body.clone(), ) } @@ -318,7 +306,7 @@ impl CodexModel { ) -> Result<(Value, Vec), AiMuxError> { let mut result = build_responses_request_body(&self.model_id, options, true)?; result.body["store"] = Value::Bool(false); - Ok((self.config.transform_body(result.body), result.warnings)) + Ok((result.body, result.warnings)) } /// Map a subscription-channel authentication failure to diff --git a/aimux-providers/src/cohere/embedding.rs b/aimux-providers/src/cohere/embedding.rs index 2e2eef584..8e91dba35 100644 --- a/aimux-providers/src/cohere/embedding.rs +++ b/aimux-providers/src/cohere/embedding.rs @@ -84,7 +84,7 @@ impl EmbeddingModel for CohereEmbeddingModel { let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/embed"), options), - exchange.transform_body(Value::Object(body)), + Value::Object(body), aimux_provider_utils::create_json_response_handler::(), super::cohere_failed_response_handler(), ) diff --git a/aimux-providers/src/cohere/mod.rs b/aimux-providers/src/cohere/mod.rs index 7c8af5869..5fe4c451a 100644 --- a/aimux-providers/src/cohere/mod.rs +++ b/aimux-providers/src/cohere/mod.rs @@ -160,7 +160,6 @@ impl CohereProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/cohere/model.rs b/aimux-providers/src/cohere/model.rs index 48d41fa71..a129e4e59 100644 --- a/aimux-providers/src/cohere/model.rs +++ b/aimux-providers/src/cohere/model.rs @@ -96,7 +96,7 @@ impl LanguageModel for CohereModel { async fn do_generate(&self, options: &CallOptions) -> Result { let body_result = build_request_body(&self.model_id, options, false)?; let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body_result.body.clone()); + let body = body_result.body.clone(); let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/chat"), options), body.clone(), @@ -232,7 +232,7 @@ impl LanguageModel for CohereModel { async fn do_stream(&self, options: &CallOptions) -> Result { let body_result = build_request_body(&self.model_id, options, true)?; let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body_result.body.clone()); + let body = body_result.body.clone(); let endpoint = exchange.url("/chat"); let resp = aimux_provider_utils::post_json_to_api( exchange.request(endpoint.clone(), options), diff --git a/aimux-providers/src/cohere/reranking.rs b/aimux-providers/src/cohere/reranking.rs index ac25fe153..de9cedec1 100644 --- a/aimux-providers/src/cohere/reranking.rs +++ b/aimux-providers/src/cohere/reranking.rs @@ -119,7 +119,6 @@ impl RerankingModel for CohereRerankingModel { } let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body); let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/rerank"), options), diff --git a/aimux-providers/src/dataforseo.rs b/aimux-providers/src/dataforseo.rs index 01672c1d8..eb84838bf 100644 --- a/aimux-providers/src/dataforseo.rs +++ b/aimux-providers/src/dataforseo.rs @@ -191,7 +191,6 @@ impl DataforseoProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/deepgram.rs b/aimux-providers/src/deepgram.rs index 29edd2701..9eef9ec96 100644 --- a/aimux-providers/src/deepgram.rs +++ b/aimux-providers/src/deepgram.rs @@ -168,7 +168,6 @@ impl DeepgramProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/deepseek/mod.rs b/aimux-providers/src/deepseek/mod.rs index 2597c283e..5706df4f6 100644 --- a/aimux-providers/src/deepseek/mod.rs +++ b/aimux-providers/src/deepseek/mod.rs @@ -134,7 +134,6 @@ impl DeepSeekProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/deepseek/model.rs b/aimux-providers/src/deepseek/model.rs index 29e9fbdbb..0ffbf96aa 100644 --- a/aimux-providers/src/deepseek/model.rs +++ b/aimux-providers/src/deepseek/model.rs @@ -46,7 +46,7 @@ pub struct DeepSeekChatLanguageModel { /// A request body and the warnings raised while building it. struct RequestBodyResult { - /// The JSON body, before the provider's `transform_request_body`. + /// The JSON body. body: Value, /// Warnings raised while building it. warnings: Vec, @@ -149,9 +149,8 @@ impl DeepSeekChatLanguageModel { .trim() } - /// The JSON body a call would send, before the provider's - /// `transform_request_body`, with the warnings raised while building it - /// (`getArgs`, plus the stream fields). + /// The JSON body a call would send, with the warnings raised while + /// building it (`getArgs`, plus the stream fields). fn request_body( &self, options: &CallOptions, @@ -369,7 +368,7 @@ impl LanguageModel for DeepSeekChatLanguageModel { .exchange(options.headers.as_ref()) .await?; let built = self.request_body(options, false)?; - let body = exchange.transform_body(built.body); + let body = built.body; let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/chat/completions"), options), body.clone(), @@ -487,7 +486,7 @@ impl LanguageModel for DeepSeekChatLanguageModel { .await?; let built = self.request_body(options, true)?; let warnings = built.warnings; - let body = exchange.transform_body(built.body); + let body = built.body; let endpoint = exchange.url("/chat/completions"); let resp = aimux_provider_utils::post_json_to_api( exchange.request(endpoint.clone(), options), diff --git a/aimux-providers/src/default_providers.rs b/aimux-providers/src/default_providers.rs index 64fd45eeb..3e54db8a9 100644 --- a/aimux-providers/src/default_providers.rs +++ b/aimux-providers/src/default_providers.rs @@ -110,7 +110,6 @@ pub fn create_provider( headers, fetch, params, - transform_request_body, } = settings; let unsupported = |field: &str| { Err(AiMuxError::InvalidArgument(format!( @@ -120,28 +119,10 @@ pub fn create_provider( if !params.is_empty() { return unsupported("params"); } - // The packages whose settings also support a body transform. macro_rules! package { - ($module:ident, $create:ident, $settings:ident) => { - Ok(Arc::new(crate::$module::$create( - #[allow(clippy::needless_update)] - crate::$module::$settings { - api_key, - base_url, - headers, - fetch, - transform_request_body, - ..Default::default() - }, - )?)) - }; - } - macro_rules! package_without_transform { ($module:ident, $create:ident, $settings:ident $(, $api_key:ident)?) => {{ - if transform_request_body.is_some() { - return unsupported("transform_request_body"); - } Ok(Arc::new(crate::$module::$create( + #[allow(clippy::needless_update)] crate::$module::$settings { $($api_key,)? base_url, @@ -154,9 +135,6 @@ pub fn create_provider( } macro_rules! package_with_string_key { ($module:ident, $create:ident, $settings:ident $(, $defaults:expr)?) => {{ - if transform_request_body.is_some() { - return unsupported("transform_request_body"); - } let api_key = match api_key { Some(aimux_provider_utils::Resolvable::Value(key)) => Some(key), None => None, @@ -178,7 +156,8 @@ pub fn create_provider( "amazon_bedrock" => package!( bedrock, create_amazon_bedrock, - AmazonBedrockProviderSettings + AmazonBedrockProviderSettings, + api_key ), "anthropic" => { package_with_string_key!( @@ -189,7 +168,7 @@ pub fn create_provider( ) } "azure" => { - package_without_transform!(azure, create_azure, AzureOpenAIProviderSettings, api_key) + package!(azure, create_azure, AzureOpenAIProviderSettings, api_key) } "cohere" => package_with_string_key!( cohere, @@ -198,7 +177,12 @@ pub fn create_provider( Default::default() ), "deepseek" => package_with_string_key!(deepseek, create_deepseek, DeepSeekProviderSettings), - "elevenlabs" => package!(elevenlabs, create_elevenlabs, ElevenLabsProviderSettings), + "elevenlabs" => package!( + elevenlabs, + create_elevenlabs, + ElevenLabsProviderSettings, + api_key + ), "google" => package_with_string_key!( google, create_google, @@ -206,7 +190,12 @@ pub fn create_provider( Default::default() ), "groq" => package_with_string_key!(groq, create_groq, GroqProviderSettings), - "huggingface" => package!(huggingface, create_huggingface, HuggingFaceProviderSettings), + "huggingface" => package!( + huggingface, + create_huggingface, + HuggingFaceProviderSettings, + api_key + ), "mistral" => { package_with_string_key!( mistral, @@ -223,113 +212,113 @@ pub fn create_provider( Default::default() ) } - "voyage" => package!(voyage, create_voyage, VoyageProviderSettings), + "voyage" => package!(voyage, create_voyage, VoyageProviderSettings, api_key), "xai" => { package_with_string_key!(xai, create_xai, XAIProviderSettings, Default::default()) } - "assemblyai" => package_without_transform!( + "assemblyai" => package!( assemblyai, create_assemblyai, AssemblyAIProviderSettings, api_key ), - "black_forest_labs" => package_without_transform!( + "black_forest_labs" => package!( black_forest_labs, create_black_forest_labs, BlackForestLabsProviderSettings, api_key ), "cartesia" => { - package_without_transform!(cartesia, create_cartesia, CartesiaProviderSettings, api_key) + package!(cartesia, create_cartesia, CartesiaProviderSettings, api_key) } "deepgram" => { - package_without_transform!(deepgram, create_deepgram, DeepgramProviderSettings, api_key) + package!(deepgram, create_deepgram, DeepgramProviderSettings, api_key) } "exa_ai" => { - package_without_transform!(exa_ai, create_exa_ai, ExaAiProviderSettings, api_key) + package!(exa_ai, create_exa_ai, ExaAiProviderSettings, api_key) } - "fal" => package_without_transform!(fal, create_fal, FalProviderSettings, api_key), - "firecrawl" => package_without_transform!( + "fal" => package!(fal, create_fal, FalProviderSettings, api_key), + "firecrawl" => package!( firecrawl, create_firecrawl, FirecrawlProviderSettings, api_key ), "gladia" => { - package_without_transform!(gladia, create_gladia, GladiaProviderSettings, api_key) + package!(gladia, create_gladia, GladiaProviderSettings, api_key) } - "google_pse" => package_without_transform!( + "google_pse" => package!( google_pse, create_google_pse, GooglePseProviderSettings, api_key ), - "hume" => package_without_transform!(hume, create_hume, HumeProviderSettings, api_key), + "hume" => package!(hume, create_hume, HumeProviderSettings, api_key), "jina_ai" => { - package_without_transform!(jina_ai, create_jina_ai, JinaAiProviderSettings, api_key) + package!(jina_ai, create_jina_ai, JinaAiProviderSettings, api_key) } "klingai" => { - package_without_transform!(klingai, create_klingai, KlingAIProviderSettings, api_key) + package!(klingai, create_klingai, KlingAIProviderSettings, api_key) } "linkup" => { - package_without_transform!(linkup, create_linkup, LinkupProviderSettings, api_key) + package!(linkup, create_linkup, LinkupProviderSettings, api_key) } - "lmnt" => package_without_transform!(lmnt, create_lmnt, LMNTProviderSettings, api_key), - "luma" => package_without_transform!(luma, create_luma, LumaProviderSettings, api_key), - "parallel_ai" => package_without_transform!( + "lmnt" => package!(lmnt, create_lmnt, LMNTProviderSettings, api_key), + "luma" => package!(luma, create_luma, LumaProviderSettings, api_key), + "parallel_ai" => package!( parallel_ai, create_parallel_ai, ParallelAiProviderSettings, api_key ), "prodia" => { - package_without_transform!(prodia, create_prodia, ProdiaProviderSettings, api_key) + package!(prodia, create_prodia, ProdiaProviderSettings, api_key) } "recraft" => { - package_without_transform!(recraft, create_recraft, RecraftProviderSettings, api_key) + package!(recraft, create_recraft, RecraftProviderSettings, api_key) } - "replicate" => package_without_transform!( + "replicate" => package!( replicate, create_replicate, ReplicateProviderSettings, api_key ), - "revai" => package_without_transform!(revai, create_revai, RevaiProviderSettings, api_key), + "revai" => package!(revai, create_revai, RevaiProviderSettings, api_key), "runwayml" => { - package_without_transform!(runwayml, create_runwayml, RunwaymlProviderSettings, api_key) + package!(runwayml, create_runwayml, RunwaymlProviderSettings, api_key) } "searxng" => { - package_without_transform!(searxng, create_searxng, SearxngProviderSettings, api_key) + package!(searxng, create_searxng, SearxngProviderSettings, api_key) } "serper" => { - package_without_transform!(serper, create_serper, SerperProviderSettings, api_key) + package!(serper, create_serper, SerperProviderSettings, api_key) } - "stability" => package_without_transform!( + "stability" => package!( stability, create_stability, StabilityProviderSettings, api_key ), "tavily" => { - package_without_transform!(tavily, create_tavily, TavilyProviderSettings, api_key) + package!(tavily, create_tavily, TavilyProviderSettings, api_key) } "tinyfish" => { - package_without_transform!(tinyfish, create_tinyfish, TinyfishProviderSettings, api_key) + package!(tinyfish, create_tinyfish, TinyfishProviderSettings, api_key) } "you_com" => { - package_without_transform!(you_com, create_you_com, YouComProviderSettings, api_key) + package!(you_com, create_you_com, YouComProviderSettings, api_key) } "aws_polly" => { if api_key.is_some() { return unsupported("api_key"); } - package_without_transform!(aws_polly, create_aws_polly, AwsPollyProviderSettings) + package!(aws_polly, create_aws_polly, AwsPollyProviderSettings) } "dataforseo" => { if api_key.is_some() { return unsupported("api_key"); } - package_without_transform!(dataforseo, create_dataforseo, DataforseoProviderSettings) + package!(dataforseo, create_dataforseo, DataforseoProviderSettings) } "google_vertex" => Ok(Arc::new(crate::vertex::create_google_vertex( crate::vertex::VertexProviderSettings { @@ -337,7 +326,6 @@ pub fn create_provider( base_url, headers: headers.map(Into::into), fetch, - transform_request_body, ..Default::default() }, )?)), @@ -351,7 +339,6 @@ pub fn create_provider( base_url, headers, fetch, - transform_request_body, ..Default::default() }, )?)) @@ -361,7 +348,6 @@ pub fn create_provider( base_url, headers, fetch, - transform_request_body, ..Default::default() }, )?)) diff --git a/aimux-providers/src/elevenlabs.rs b/aimux-providers/src/elevenlabs.rs index 399695fc1..6688ba1af 100644 --- a/aimux-providers/src/elevenlabs.rs +++ b/aimux-providers/src/elevenlabs.rs @@ -32,9 +32,7 @@ use aimux_provider_utils::{ FetchFunction, HeaderMapOpt, HeadersFn, HttpBody, Resolvable, validate_base_url, }; -use crate::shared::{ - AuthScheme, Credential, EndpointConfig, TransformRequestBody, credential_headers, -}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; pub(crate) mod options; @@ -110,9 +108,6 @@ pub struct ElevenLabsProviderSettings { /// The transport: a mock, a signing decorator, a proxy-aware client. /// `None` uses the process default, resolved per request. pub fetch: Option, - /// Rewrites every JSON request body once, after it is serialized and - /// before it is sent. - pub transform_request_body: Option, /// Opens the WebSocket of realtime transcription (the AI SDK's /// `webSocket`). `None` uses the built-in tungstenite connector. #[cfg(feature = "realtime")] @@ -131,11 +126,7 @@ impl std::fmt::Debug for ElevenLabsProviderSettings { &self.headers.as_ref().map(std::collections::HashMap::len), ) .field("name", &self.name) - .field("fetch", &self.fetch.is_some()) - .field( - "transform_request_body", - &self.transform_request_body.is_some(), - ); + .field("fetch", &self.fetch.is_some()); #[cfg(feature = "realtime")] debug.field("web_socket", &self.web_socket.is_some()); debug.finish() @@ -166,7 +157,6 @@ pub fn create_elevenlabs( settings.headers, ), fetch: settings.fetch, - transform_request_body: settings.transform_request_body, #[cfg(feature = "realtime")] web_socket: settings.web_socket, }) @@ -193,7 +183,6 @@ pub struct ElevenLabsProvider { base_url: String, headers: HeadersFn, fetch: Option, - transform_request_body: Option, #[cfg(feature = "realtime")] web_socket: Option>, } @@ -205,7 +194,6 @@ impl ElevenLabsProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - self.transform_request_body.clone(), ) } @@ -302,7 +290,7 @@ impl SpeechModel for ElevenLabsSpeechModel { let resp = aimux_provider_utils::post_json_to_api( exchange.request(url, options), - exchange.transform_body(Value::Object(body.clone())), + Value::Object(body.clone()), aimux_provider_utils::create_binary_response_handler(), elevenlabs_failed_response_handler(), ) diff --git a/aimux-providers/src/exa_ai.rs b/aimux-providers/src/exa_ai.rs index 9f5472049..c1cd99cbe 100644 --- a/aimux-providers/src/exa_ai.rs +++ b/aimux-providers/src/exa_ai.rs @@ -141,7 +141,6 @@ impl ExaAiProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/fal.rs b/aimux-providers/src/fal.rs index d18a45069..c38820160 100644 --- a/aimux-providers/src/fal.rs +++ b/aimux-providers/src/fal.rs @@ -182,7 +182,6 @@ impl FalProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/firecrawl.rs b/aimux-providers/src/firecrawl.rs index 04ed8affa..408cd831d 100644 --- a/aimux-providers/src/firecrawl.rs +++ b/aimux-providers/src/firecrawl.rs @@ -139,7 +139,6 @@ impl FirecrawlProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/gladia.rs b/aimux-providers/src/gladia.rs index d3fd38c16..410a28ea3 100644 --- a/aimux-providers/src/gladia.rs +++ b/aimux-providers/src/gladia.rs @@ -163,7 +163,6 @@ impl GladiaProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/google/embedding.rs b/aimux-providers/src/google/embedding.rs index 9384c14ec..2d7b09582 100644 --- a/aimux-providers/src/google/embedding.rs +++ b/aimux-providers/src/google/embedding.rs @@ -180,7 +180,7 @@ impl EmbeddingModel for GoogleEmbeddingModel { exchange.url(&format!("/models/{}:{endpoint}", self.model_id)), options, ), - exchange.transform_body(body), + body, aimux_provider_utils::create_json_response_handler::(), super::google_failed_response_handler(), ) diff --git a/aimux-providers/src/google/files.rs b/aimux-providers/src/google/files.rs index d30a17004..96008377e 100644 --- a/aimux-providers/src/google/files.rs +++ b/aimux-providers/src/google/files.rs @@ -210,7 +210,7 @@ impl Files for GoogleFiles { abort_signal: options.abort_signal.clone(), ..Default::default() }), - exchange.transform_body(init_body_value), + init_body_value, aimux_provider_utils::ResponseHandler::new(|input| async move { let headers = aimux_provider_utils::extract_response_headers::extract_response_headers( diff --git a/aimux-providers/src/google/image.rs b/aimux-providers/src/google/image.rs index 4ae1a757e..df936950b 100644 --- a/aimux-providers/src/google/image.rs +++ b/aimux-providers/src/google/image.rs @@ -164,7 +164,7 @@ impl ImageModel for GoogleImageModel { exchange.url(&format!("/models/{}:generateContent", self.model_id)), options, ), - exchange.transform_body(body), + body, aimux_provider_utils::create_json_response_handler::(), super::google_failed_response_handler(), ) diff --git a/aimux-providers/src/google/mod.rs b/aimux-providers/src/google/mod.rs index 1f0977495..489755717 100644 --- a/aimux-providers/src/google/mod.rs +++ b/aimux-providers/src/google/mod.rs @@ -307,7 +307,6 @@ impl GoogleProvider { }), fetch: self.fetch.clone(), supported_urls: Arc::new(move |model_id| supported_urls(&urls_base, Some(model_id))), - transform_request_body: None, } } diff --git a/aimux-providers/src/google/model.rs b/aimux-providers/src/google/model.rs index 2a264b3ec..f5cb9947d 100644 --- a/aimux-providers/src/google/model.rs +++ b/aimux-providers/src/google/model.rs @@ -92,7 +92,6 @@ impl LanguageModel for GoogleModel { } } let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body); let url = exchange.url(&format!("/{}:generateContent", self.model_path())); let resp = aimux_provider_utils::post_json_to_api( exchange.request(url, options), @@ -174,7 +173,6 @@ impl LanguageModel for GoogleModel { } } let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body); let endpoint = exchange.url(&format!( "/{}:streamGenerateContent?alt=sse", self.model_path() diff --git a/aimux-providers/src/google/speech.rs b/aimux-providers/src/google/speech.rs index e6ad09bae..6bbcecc1d 100644 --- a/aimux-providers/src/google/speech.rs +++ b/aimux-providers/src/google/speech.rs @@ -458,7 +458,7 @@ impl SpeechModel for GoogleSpeechModel { exchange.url(&format!("/models/{}:generateContent", self.model_id)), options, ), - exchange.transform_body(body.clone()), + body.clone(), aimux_provider_utils::create_json_response_handler::(), super::google_failed_response_handler(), ) diff --git a/aimux-providers/src/google/transcription.rs b/aimux-providers/src/google/transcription.rs index 3c3700701..5a32254cb 100644 --- a/aimux-providers/src/google/transcription.rs +++ b/aimux-providers/src/google/transcription.rs @@ -238,7 +238,7 @@ impl TranscriptionModel for GoogleTranscriptionModel { let exchange = self.config.exchange(options.headers.as_ref()).await?; let response = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/interactions"), options), - exchange.transform_body(body), + body, aimux_provider_utils::create_json_response_handler::(), super::google_failed_response_handler(), ) diff --git a/aimux-providers/src/google/video.rs b/aimux-providers/src/google/video.rs index c0598ba70..c4438c7d5 100644 --- a/aimux-providers/src/google/video.rs +++ b/aimux-providers/src/google/video.rs @@ -173,7 +173,7 @@ impl VideoModel for GoogleVideoModel { let resp = aimux_provider_utils::post_json_to_api( exchange.request(url, options), - exchange.transform_body(body), + body, aimux_provider_utils::create_json_response_handler(), super::google_failed_response_handler(), ) diff --git a/aimux-providers/src/google_pse.rs b/aimux-providers/src/google_pse.rs index d61f04a52..60d82ece5 100644 --- a/aimux-providers/src/google_pse.rs +++ b/aimux-providers/src/google_pse.rs @@ -170,7 +170,6 @@ impl GooglePseProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ), self.credential.clone(), self.cx.clone(), diff --git a/aimux-providers/src/groq/mod.rs b/aimux-providers/src/groq/mod.rs index 7e5067191..207101a10 100644 --- a/aimux-providers/src/groq/mod.rs +++ b/aimux-providers/src/groq/mod.rs @@ -138,7 +138,6 @@ impl GroqProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/groq/model.rs b/aimux-providers/src/groq/model.rs index 04fdd0714..4271f34c3 100644 --- a/aimux-providers/src/groq/model.rs +++ b/aimux-providers/src/groq/model.rs @@ -280,7 +280,7 @@ impl LanguageModel for GroqChatLanguageModel { async fn do_generate(&self, options: &CallOptions) -> Result { let (args, warnings) = self.get_args(options)?; let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(Value::Object(args)); + let body = Value::Object(args); let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/chat/completions"), options), @@ -352,7 +352,7 @@ impl LanguageModel for GroqChatLanguageModel { let (mut args, warnings) = self.get_args(options)?; args.insert("stream".into(), json!(true)); let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(Value::Object(args)); + let body = Value::Object(args); let endpoint = exchange.url("/chat/completions"); let resp = aimux_provider_utils::post_json_to_api( diff --git a/aimux-providers/src/huggingface.rs b/aimux-providers/src/huggingface.rs index 72d49a97f..8343e4ccc 100644 --- a/aimux-providers/src/huggingface.rs +++ b/aimux-providers/src/huggingface.rs @@ -17,7 +17,6 @@ pub(crate) mod options; pub mod responses; -pub use crate::shared::TransformRequestBody; pub use responses::HuggingFaceResponsesModel; use std::sync::{Arc, OnceLock}; @@ -66,9 +65,6 @@ pub struct HuggingFaceProviderSettings { /// The transport: a mock, a signing decorator, a proxy-aware client. /// `None` uses the process default, resolved per request. pub fetch: Option, - /// Rewrites every JSON request body once, after it is serialized and - /// before it is sent. - pub transform_request_body: Option, } impl std::fmt::Debug for HuggingFaceProviderSettings { @@ -83,10 +79,6 @@ impl std::fmt::Debug for HuggingFaceProviderSettings { ) .field("name", &self.name) .field("fetch", &self.fetch.is_some()) - .field( - "transform_request_body", - &self.transform_request_body.is_some(), - ) .finish() } } @@ -114,7 +106,6 @@ pub fn create_huggingface( settings.headers, ), fetch: settings.fetch, - transform_request_body: settings.transform_request_body, }) } @@ -136,7 +127,6 @@ pub struct HuggingFaceProvider { base_url: String, headers: HeadersFn, fetch: Option, - transform_request_body: Option, } impl HuggingFaceProvider { @@ -146,7 +136,6 @@ impl HuggingFaceProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - self.transform_request_body.clone(), ) } @@ -156,7 +145,6 @@ impl HuggingFaceProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - self.transform_request_body.clone(), ) } diff --git a/aimux-providers/src/huggingface/responses.rs b/aimux-providers/src/huggingface/responses.rs index e4b4614d5..3b9792e04 100644 --- a/aimux-providers/src/huggingface/responses.rs +++ b/aimux-providers/src/huggingface/responses.rs @@ -152,7 +152,7 @@ impl LanguageModel for HuggingFaceResponsesModel { async fn do_generate(&self, options: &CallOptions) -> Result { let exchange = self.config.exchange(options.headers.as_ref()).await?; let request = build_request_body_with_warnings(&self.model_id, options, false)?; - let body = exchange.transform_body(request.body); + let body = request.body; let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/responses"), options), @@ -205,7 +205,7 @@ impl LanguageModel for HuggingFaceResponsesModel { async fn do_stream(&self, options: &CallOptions) -> Result { let exchange = self.config.exchange(options.headers.as_ref()).await?; let request = build_request_body_with_warnings(&self.model_id, options, true)?; - let body = exchange.transform_body(request.body); + let body = request.body; let endpoint = exchange.url("/responses"); let resp = aimux_provider_utils::post_json_to_api( diff --git a/aimux-providers/src/hume.rs b/aimux-providers/src/hume.rs index 9300b9b9a..4292a5f20 100644 --- a/aimux-providers/src/hume.rs +++ b/aimux-providers/src/hume.rs @@ -170,7 +170,6 @@ impl HumeProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/jina_ai.rs b/aimux-providers/src/jina_ai.rs index c042176b6..784f75fa8 100644 --- a/aimux-providers/src/jina_ai.rs +++ b/aimux-providers/src/jina_ai.rs @@ -149,7 +149,6 @@ impl JinaAiProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/klingai.rs b/aimux-providers/src/klingai.rs index 71aa4aa3f..624ac12c3 100644 --- a/aimux-providers/src/klingai.rs +++ b/aimux-providers/src/klingai.rs @@ -148,7 +148,6 @@ impl KlingAIProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/linkup.rs b/aimux-providers/src/linkup.rs index 0b2bc93a4..102c95b13 100644 --- a/aimux-providers/src/linkup.rs +++ b/aimux-providers/src/linkup.rs @@ -154,7 +154,6 @@ impl LinkupProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/lmnt.rs b/aimux-providers/src/lmnt.rs index 2937cfb91..aa514c22b 100644 --- a/aimux-providers/src/lmnt.rs +++ b/aimux-providers/src/lmnt.rs @@ -155,7 +155,6 @@ impl LMNTProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/luma.rs b/aimux-providers/src/luma.rs index d98875565..063a6822f 100644 --- a/aimux-providers/src/luma.rs +++ b/aimux-providers/src/luma.rs @@ -154,7 +154,6 @@ impl LumaProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/mistral/embedding.rs b/aimux-providers/src/mistral/embedding.rs index 2243690c7..56c242c65 100644 --- a/aimux-providers/src/mistral/embedding.rs +++ b/aimux-providers/src/mistral/embedding.rs @@ -78,7 +78,7 @@ impl EmbeddingModel for MistralEmbeddingModel { let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/embeddings"), options), - exchange.transform_body(Value::Object(body)), + Value::Object(body), aimux_provider_utils::create_json_response_handler(), super::mistral_failed_response_handler(), ) diff --git a/aimux-providers/src/mistral/mod.rs b/aimux-providers/src/mistral/mod.rs index 57283ff0f..2e1cefb27 100644 --- a/aimux-providers/src/mistral/mod.rs +++ b/aimux-providers/src/mistral/mod.rs @@ -209,7 +209,6 @@ impl MistralProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/mistral/model.rs b/aimux-providers/src/mistral/model.rs index 6d3283573..c2a70dba3 100644 --- a/aimux-providers/src/mistral/model.rs +++ b/aimux-providers/src/mistral/model.rs @@ -172,7 +172,7 @@ impl LanguageModel for MistralModel { async fn do_generate(&self, options: &CallOptions) -> Result { let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(build_request_body(&self.model_id, options, false)?); + let body = build_request_body(&self.model_id, options, false)?; let endpoint = exchange.url("/chat/completions"); let resp = aimux_provider_utils::post_json_to_api( exchange.request(endpoint.clone(), options), @@ -274,7 +274,7 @@ impl LanguageModel for MistralModel { async fn do_stream(&self, options: &CallOptions) -> Result { let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(build_request_body(&self.model_id, options, true)?); + let body = build_request_body(&self.model_id, options, true)?; let endpoint = exchange.url("/chat/completions"); let resp = aimux_provider_utils::post_json_to_api( exchange.request(endpoint.clone(), options), diff --git a/aimux-providers/src/open_responses.rs b/aimux-providers/src/open_responses.rs index 20e5c5ec4..f75a5fd4b 100644 --- a/aimux-providers/src/open_responses.rs +++ b/aimux-providers/src/open_responses.rs @@ -37,7 +37,7 @@ use aimux_provider_utils::{ FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, combine_headers, validate_base_url, }; -use crate::shared::{Credential, EndpointConfig, TransformRequestBody, provider_headers}; +use crate::shared::{Credential, EndpointConfig, provider_headers}; fn open_responses_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -114,9 +114,6 @@ pub struct OpenResponsesProviderSettings { /// The transport: a mock, a signing decorator, a proxy-aware client. /// `None` uses the process default, resolved per request. pub fetch: Option, - /// Rewrites every JSON request body once, after it is serialized and - /// before it is sent. - pub transform_request_body: Option, } impl OpenResponsesProviderSettings { @@ -129,7 +126,6 @@ impl OpenResponsesProviderSettings { api_key: None, headers: None, fetch: None, - transform_request_body: None, } } } @@ -143,10 +139,6 @@ impl std::fmt::Debug for OpenResponsesProviderSettings { .field("api_key", &self.api_key) .field("headers", &self.headers.is_some()) .field("fetch", &self.fetch.is_some()) - .field( - "transform_request_body", - &self.transform_request_body.is_some(), - ) .finish() } } @@ -199,7 +191,6 @@ pub fn create_open_responses( base_url, headers, fetch: settings.fetch, - transform_request_body: settings.transform_request_body, }) } @@ -211,7 +202,6 @@ pub struct OpenResponsesProvider { base_url: String, headers: HeadersFn, fetch: Option, - transform_request_body: Option, } impl OpenResponsesProvider { @@ -226,7 +216,6 @@ impl OpenResponsesProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - self.transform_request_body.clone(), ), } } @@ -279,7 +268,6 @@ impl LanguageModel for OpenResponsesModel { validate_tool_output_media(&options.prompt)?; let (body, warnings) = build_request_body(&self.model_id, options, &self.provider_options_name); - let body = exchange.transform_body(body); let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/responses"), options), @@ -425,7 +413,7 @@ impl LanguageModel for OpenResponsesModel { if let Some(obj) = b.as_object_mut() { obj.insert("stream".to_string(), json!(true)); } - exchange.transform_body(b) + b }; let endpoint = exchange.url("/responses"); diff --git a/aimux-providers/src/openai/config.rs b/aimux-providers/src/openai/config.rs index f39d20f0f..c66ffea55 100644 --- a/aimux-providers/src/openai/config.rs +++ b/aimux-providers/src/openai/config.rs @@ -11,8 +11,6 @@ use std::collections::HashMap; use std::sync::Arc; -use serde_json::Value; - use aimux_core::AiMuxError; use aimux_core::language_model::SupportedUrls; use aimux_provider_utils::{ @@ -22,8 +20,6 @@ use aimux_provider_utils::{ use super::responses::ResponsesProfile; -pub use crate::shared::TransformRequestBody; - /// Maps an endpoint path (`"/chat/completions"`) to the full request URL. It /// can fail because a host may only know the address when a call is made (an /// Azure resource name read from the environment). @@ -45,8 +41,6 @@ pub(crate) struct OpenAIModelConfig { pub(crate) fetch: Option, /// URLs the model fetches itself. Empty: the caller downloads them. pub(crate) supported_urls: SupportedUrls, - /// Provider-level request-body rewrite. - pub(crate) transform_request_body: Option, /// How the Responses API model reads and writes providerOptions and file /// ids for this host. pub(crate) responses: ResponsesProfile, @@ -60,7 +54,6 @@ impl OpenAIModelConfig { base_url: String, headers: HeadersFn, fetch: Option, - transform_request_body: Option, ) -> Self { Self { provider, @@ -69,7 +62,6 @@ impl OpenAIModelConfig { token_provider: None, fetch, supported_urls: SupportedUrls::default(), - transform_request_body, responses: ResponsesProfile::default(), chat_options: super::options::parse_chat_options, } @@ -157,14 +149,6 @@ impl OpenAIModelConfig { request.credentialed_origin = Some(request.url.clone()); request } - - /// Run the provider-level body rewrite, when there is one. - pub(crate) fn transform_body(&self, body: Value) -> Value { - match &self.transform_request_body { - Some(transform) => transform(body), - None => body, - } - } } /// Native language model URL patterns from the upstream model classes. diff --git a/aimux-providers/src/openai/embedding.rs b/aimux-providers/src/openai/embedding.rs index 8b4bef080..10a1e3ab4 100644 --- a/aimux-providers/src/openai/embedding.rs +++ b/aimux-providers/src/openai/embedding.rs @@ -76,7 +76,7 @@ impl EmbeddingModel for OpenAIEmbeddingModel { .config .request_headers(options.headers.as_ref()) .await?; - let body = self.config.transform_body(Value::Object(body)); + let body = Value::Object(body); let resp = aimux_provider_utils::post_json_to_api( self.config diff --git a/aimux-providers/src/openai/image.rs b/aimux-providers/src/openai/image.rs index 7319299b8..1ed060fd4 100644 --- a/aimux-providers/src/openai/image.rs +++ b/aimux-providers/src/openai/image.rs @@ -135,13 +135,11 @@ impl ImageModel for OpenAIImageModel { // ── Generation path: JSON body ── validate_image_options(&options.provider_options, false)?; let openai_options = parse_generation_provider_options(&options.provider_options); - let body = self - .config - .transform_body(Value::Object(build_generation_body( - &self.model_id, - options, - &openai_options, - ))); + let body = Value::Object(build_generation_body( + &self.model_id, + options, + &openai_options, + )); let resp = aimux_provider_utils::post_json_to_api( self.config diff --git a/aimux-providers/src/openai/mod.rs b/aimux-providers/src/openai/mod.rs index 1cb44f17b..8e8aebd21 100644 --- a/aimux-providers/src/openai/mod.rs +++ b/aimux-providers/src/openai/mod.rs @@ -296,7 +296,6 @@ impl OpenAIProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ); config.supported_urls = config::supported_urls(method); if method == "responses" { diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs index 064200ba3..8e41b8c27 100644 --- a/aimux-providers/src/openai/model.rs +++ b/aimux-providers/src/openai/model.rs @@ -24,7 +24,7 @@ use aimux_provider_utils::{ HttpRequest, StreamingToolCallDelta, StreamingToolCallTracker, TypeValidation, generate_id, }; -use super::config::{OpenAIModelConfig, TransformRequestBody}; +use super::config::OpenAIModelConfig; use super::convert::{ RequestBodyResult, build_request_body_with_chat_options, parse_finish_reason, }; @@ -137,14 +137,7 @@ impl LanguageModel for OpenAIModel { let http = self.config .http_request(self.config.url("/chat/completions")?, headers, options); - execute_generate( - http, - &self.model_id, - options, - self.config.transform_request_body.as_ref(), - self.config.chat_options, - ) - .await + execute_generate(http, &self.model_id, options, self.config.chat_options).await } async fn do_stream(&self, options: &CallOptions) -> Result { @@ -155,14 +148,7 @@ impl LanguageModel for OpenAIModel { let http = self.config .http_request(self.config.url("/chat/completions")?, headers, options); - execute_stream( - http, - &self.model_id, - options, - self.config.transform_request_body.as_ref(), - self.config.chat_options, - ) - .await + execute_stream(http, &self.model_id, options, self.config.chat_options).await } } @@ -176,8 +162,7 @@ impl LanguageModel for OpenAIModel { /// /// `http` carries the full chat-completions URL, the auth and request headers /// and the transport; `model_id` is placed in the request body's `model` -/// field. `transform_request_body`, when set, rewrites the finished body once -/// before it is sent (and the rewritten body is what the result reports). +/// field. /// /// # Errors /// @@ -188,15 +173,11 @@ pub(crate) async fn execute_generate( http: HttpRequest, model_id: &str, options: &CallOptions, - transform_request_body: Option<&TransformRequestBody>, parse_chat_options: super::options::ChatOptionsParser, ) -> Result { let request_result = build_request_body_with_chat_options(model_id, options, false, parse_chat_options)?; - let body = match transform_request_body { - Some(transform) => transform(request_result.body), - None => request_result.body, - }; + let body = request_result.body; let resp = aimux_provider_utils::post_json_to_api( http, @@ -339,7 +320,7 @@ pub(crate) async fn execute_generate( /// /// `http` carries the full chat-completions URL, the auth and request headers /// and the transport; `model_id` is placed in the request body's `model` -/// field. `transform_request_body` is applied as in [`execute_generate`]. +/// field. /// /// # Errors /// @@ -349,7 +330,6 @@ pub(crate) async fn execute_stream( http: HttpRequest, model_id: &str, options: &CallOptions, - transform_request_body: Option<&TransformRequestBody>, parse_chat_options: super::options::ChatOptionsParser, ) -> Result { let request_result = @@ -357,10 +337,6 @@ pub(crate) async fn execute_stream( // M9 (RFC-0016): keep the warnings computed while building the body — // they are emitted in `StreamStart` below instead of being dropped. let RequestBodyResult { body, warnings } = request_result; - let body = match transform_request_body { - Some(transform) => transform(body), - None => body, - }; let endpoint = http.url.clone(); let resp = aimux_provider_utils::post_json_to_api( diff --git a/aimux-providers/src/openai/responses/mod.rs b/aimux-providers/src/openai/responses/mod.rs index 1bd543a18..c53a28c44 100644 --- a/aimux-providers/src/openai/responses/mod.rs +++ b/aimux-providers/src/openai/responses/mod.rs @@ -234,7 +234,6 @@ impl OpenAIResponsesModel { "Adjacent reasoning effort configuration updates".into(), )); } - result.body = self.config.transform_body(result.body); Ok(result) } diff --git a/aimux-providers/src/openai/speech.rs b/aimux-providers/src/openai/speech.rs index 8f8ffc26c..1993e28f3 100644 --- a/aimux-providers/src/openai/speech.rs +++ b/aimux-providers/src/openai/speech.rs @@ -55,7 +55,7 @@ impl SpeechModel for OpenAISpeechModel { .config .request_headers(options.headers.as_ref()) .await?; - let body = self.config.transform_body(Value::Object(body)); + let body = Value::Object(body); let resp = aimux_provider_utils::post_json_to_api( self.config diff --git a/aimux-providers/src/openai_compatible/chat.rs b/aimux-providers/src/openai_compatible/chat.rs index 860e75ef5..08991b4a8 100644 --- a/aimux-providers/src/openai_compatible/chat.rs +++ b/aimux-providers/src/openai_compatible/chat.rs @@ -38,10 +38,9 @@ impl OpenAICompatibleChatModel { Self { model_id, config } } - /// The JSON body a call would send, before the provider's - /// `transform_request_body`, with the warnings raised while building it. - /// Nothing is read from the environment and nothing is sent: a key is not - /// resolved here. + /// The JSON body a call would send, with the warnings raised while + /// building it. Nothing is read from the environment and nothing is sent: + /// a key is not resolved here. /// /// # Errors /// @@ -193,7 +192,7 @@ impl LanguageModel for OpenAICompatibleChatModel { async fn do_generate(&self, options: &CallOptions) -> Result { let built = self.request_body(options, false)?; let metadata_key = built.metadata_key; - let body = self.config.transform_body(built.body); + let body = built.body; let headers = self .config .request_headers(options.headers.as_ref()) @@ -312,7 +311,7 @@ impl LanguageModel for OpenAICompatibleChatModel { let built = self.request_body(options, true)?; let metadata_key = built.metadata_key; let warnings = built.warnings; - let body = self.config.transform_body(built.body); + let body = built.body; let headers = self .config .request_headers(options.headers.as_ref()) diff --git a/aimux-providers/src/openai_compatible/completion.rs b/aimux-providers/src/openai_compatible/completion.rs index 10bc9c161..ea5508854 100644 --- a/aimux-providers/src/openai_compatible/completion.rs +++ b/aimux-providers/src/openai_compatible/completion.rs @@ -184,16 +184,6 @@ impl OpenAICompatibleCompletionModel { } } - fn transform_body(&self, body: Value) -> Value { - match &self.config { - CompletionConfig::Compatible(_) => body, - CompletionConfig::Native(config) => match &config.transform_request_body { - Some(transform) => transform(body), - None => body, - }, - } - } - fn failed_response_handler(&self) -> aimux_provider_utils::ResponseHandler { match &self.config { CompletionConfig::Compatible(config) => config.failed_response_handler(), @@ -394,7 +384,6 @@ impl LanguageModel for OpenAICompatibleCompletionModel { async fn do_generate(&self, options: &CallOptions) -> Result { let (body, warnings) = self.arguments(options)?; - let body = self.transform_body(body); let response = aimux_provider_utils::post_json_to_api( self.request(options).await?, body.clone(), @@ -461,7 +450,6 @@ impl LanguageModel for OpenAICompatibleCompletionModel { if include_usage { body["stream_options"] = json!({ "include_usage": true }); } - let body = self.transform_body(body); let request = self.request(options).await?; let error_url = request.url.clone(); let response = aimux_provider_utils::post_json_to_api( diff --git a/aimux-providers/src/openai_compatible/config.rs b/aimux-providers/src/openai_compatible/config.rs index 4284465b5..c3d6fb1c7 100644 --- a/aimux-providers/src/openai_compatible/config.rs +++ b/aimux-providers/src/openai_compatible/config.rs @@ -23,8 +23,6 @@ use aimux_provider_utils::{ use crate::preset::PresetDescriptor; use crate::shared::ProviderHeaders; -pub use crate::shared::TransformRequestBody; - /// Where the base URL comes from. /// /// A plain provider has a [`Fixed`](Self::Fixed) URL. A preset whose URL is @@ -173,7 +171,6 @@ pub(crate) struct CompatModelConfig { pub headers: ProviderHeaders, /// Transport; `None` uses the process default, resolved per request. pub fetch: Option, - pub transform_request_body: Option, pub chat: ChatSettings, } @@ -242,14 +239,6 @@ impl CompatModelConfig { request } - /// Run the provider-level body rewrite, when there is one. - pub(crate) fn transform_body(&self, body: Value) -> Value { - match &self.transform_request_body { - Some(transform) => transform(body), - None => body, - } - } - /// The failed-response handler of the compatible error structure. pub(crate) fn failed_response_handler( &self, diff --git a/aimux-providers/src/openai_compatible/convert.rs b/aimux-providers/src/openai_compatible/convert.rs index a200f13ff..3562d615a 100644 --- a/aimux-providers/src/openai_compatible/convert.rs +++ b/aimux-providers/src/openai_compatible/convert.rs @@ -119,7 +119,7 @@ pub(crate) struct ChatBodySpec<'a> { /// A request body and what building it produced. #[derive(Debug, Clone)] pub(crate) struct RequestBodyResult { - /// The JSON body, before the provider's `transform_request_body`. + /// The JSON body. pub body: Value, /// Warnings raised while building it. pub warnings: Vec, diff --git a/aimux-providers/src/openai_compatible/mod.rs b/aimux-providers/src/openai_compatible/mod.rs index e4b81a79a..65992d16b 100644 --- a/aimux-providers/src/openai_compatible/mod.rs +++ b/aimux-providers/src/openai_compatible/mod.rs @@ -9,8 +9,7 @@ //! What is fixed when the provider is created, and what is not: //! //! - The explicit settings (`name`, `base_url`, `headers`, `query_params`, -//! `fetch`, the capability flags, `transform_request_body`) are fixed in the -//! factory. +//! `fetch`, the capability flags) are fixed in the factory. //! - A non-empty `api_key` produces a bearer header. `None` or `""` sends //! no `Authorization` header, which is what a local server wants. //! @@ -35,7 +34,7 @@ mod types; pub use chat::OpenAICompatibleChatModel; pub use completion::OpenAICompatibleCompletionModel; -pub use config::{ConvertUsage, TransformRequestBody}; +pub use config::ConvertUsage; pub use embedding::OpenAICompatibleEmbeddingModel; pub use image::OpenAICompatibleImageModel; @@ -81,9 +80,6 @@ pub struct OpenAICompatibleProviderSettings { /// Whether the chat endpoint accepts `json_schema` response formats; when /// not, a schema degrades to `json_object` with a warning. pub supports_structured_outputs: Option, - /// Rewrites each chat JSON request body once, after it is serialized and - /// before it is sent. - pub transform_request_body: Option, /// URL patterns supported by chat models. pub supported_urls: Option, /// How chat token usage is read; the generic OpenAI shape by default. @@ -108,10 +104,6 @@ impl std::fmt::Debug for OpenAICompatibleProviderSettings { "supports_structured_outputs", &self.supports_structured_outputs, ) - .field( - "transform_request_body", - &self.transform_request_body.is_some(), - ) .field("supported_urls", &self.supported_urls.is_some()) .field("convert_usage", &self.convert_usage) .finish() @@ -144,7 +136,6 @@ pub fn create_openai_compatible( headers: settings.headers, query_params: settings.query_params, fetch: settings.fetch, - transform_request_body: settings.transform_request_body, chat: ChatSettings { include_usage: settings.include_usage.unwrap_or(false), supports_structured_outputs: settings.supports_structured_outputs.unwrap_or(false), @@ -166,7 +157,6 @@ pub(crate) struct Assembly { pub headers: Option, pub query_params: Option>, pub fetch: Option, - pub transform_request_body: Option, pub chat: ChatSettings, } @@ -178,7 +168,6 @@ pub struct OpenAICompatibleProvider { query_params: Option>>, headers: ProviderHeaders, fetch: Option, - transform_request_body: Option, chat: ChatSettings, } @@ -201,7 +190,6 @@ impl OpenAICompatibleProvider { user_agent: Some(("openai-compatible", "3.0.59")), }, fetch: assembly.fetch, - transform_request_body: assembly.transform_request_body, chat: assembly.chat, }) } @@ -213,11 +201,6 @@ impl OpenAICompatibleProvider { query_params: self.query_params.clone(), headers: self.headers.clone(), fetch: self.fetch.clone(), - transform_request_body: if method == "chat" { - self.transform_request_body.clone() - } else { - None - }, chat: self.chat.clone(), } } diff --git a/aimux-providers/src/parallel_ai.rs b/aimux-providers/src/parallel_ai.rs index 12365ba7d..63d4bd10c 100644 --- a/aimux-providers/src/parallel_ai.rs +++ b/aimux-providers/src/parallel_ai.rs @@ -131,7 +131,6 @@ impl ParallelAiProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/preset.rs b/aimux-providers/src/preset.rs index 0b9ad1152..b3e9cb837 100644 --- a/aimux-providers/src/preset.rs +++ b/aimux-providers/src/preset.rs @@ -34,9 +34,7 @@ use aimux_core::error::AiMuxError; use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use crate::openai_compatible::config::{BaseUrl, ChatSettings}; -use crate::openai_compatible::{ - Assembly, ConvertUsage, OpenAICompatibleProvider, TransformRequestBody, -}; +use crate::openai_compatible::{Assembly, ConvertUsage, OpenAICompatibleProvider}; use crate::shared::Credential; /// How a preset authenticates. @@ -109,8 +107,6 @@ pub struct PresetSettings { /// non-derived parameters are accepted; a value must be a plain segment /// (`A-Z a-z 0-9 . _ -`). pub params: HashMap, - /// Rewrites every JSON request body once, before it is sent. - pub transform_request_body: Option, } impl std::fmt::Debug for PresetSettings { @@ -125,10 +121,6 @@ impl std::fmt::Debug for PresetSettings { ) .field("fetch", &self.fetch.is_some()) .field("params", &self.params) - .field( - "transform_request_body", - &self.transform_request_body.is_some(), - ) .finish() } } @@ -406,7 +398,6 @@ fn assemble( headers: settings.headers, query_params: None, fetch: settings.fetch, - transform_request_body: settings.transform_request_body, chat, }) } diff --git a/aimux-providers/src/prodia.rs b/aimux-providers/src/prodia.rs index f5a58f184..c69ceb313 100644 --- a/aimux-providers/src/prodia.rs +++ b/aimux-providers/src/prodia.rs @@ -169,7 +169,6 @@ impl ProdiaProvider { ), }, self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/recraft.rs b/aimux-providers/src/recraft.rs index e0513f8ca..ffcfbe14e 100644 --- a/aimux-providers/src/recraft.rs +++ b/aimux-providers/src/recraft.rs @@ -155,7 +155,6 @@ impl RecraftProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/replicate.rs b/aimux-providers/src/replicate.rs index 1d4cc4ee6..0c0a29bdf 100644 --- a/aimux-providers/src/replicate.rs +++ b/aimux-providers/src/replicate.rs @@ -167,7 +167,6 @@ impl ReplicateProvider { ), }, self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/revai.rs b/aimux-providers/src/revai.rs index 5fff1f651..5bcbdaeca 100644 --- a/aimux-providers/src/revai.rs +++ b/aimux-providers/src/revai.rs @@ -164,7 +164,6 @@ impl RevaiProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/runwayml.rs b/aimux-providers/src/runwayml.rs index 9c5a8c19b..452fed5d1 100644 --- a/aimux-providers/src/runwayml.rs +++ b/aimux-providers/src/runwayml.rs @@ -153,7 +153,6 @@ impl RunwaymlProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/serper.rs b/aimux-providers/src/serper.rs index 41b45a1da..18aceae9e 100644 --- a/aimux-providers/src/serper.rs +++ b/aimux-providers/src/serper.rs @@ -124,7 +124,6 @@ impl SerperProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/shared/exchange.rs b/aimux-providers/src/shared/exchange.rs index d6a941890..5dfebcc60 100644 --- a/aimux-providers/src/shared/exchange.rs +++ b/aimux-providers/src/shared/exchange.rs @@ -12,7 +12,6 @@ use std::collections::HashMap; use std::sync::Arc; use futures::future::BoxFuture; -use serde_json::Value; use aimux_core::AiMuxError; use aimux_core::language_model::SupportedUrls; @@ -21,8 +20,6 @@ use aimux_provider_utils::{ normalize_headers, }; -use super::TransformRequestBody; - /// What one request is addressed to: the base URL and the provider headers, /// both resolved for this request. pub(crate) struct Endpoint { @@ -50,8 +47,6 @@ pub(crate) struct EndpointConfig { /// URLs the model fetches itself, by model id. Empty: the caller /// downloads them. pub(crate) supported_urls: SupportedUrlsFn, - /// Provider-level request-body rewrite. - pub(crate) transform_request_body: Option, } impl EndpointConfig { @@ -62,7 +57,6 @@ impl EndpointConfig { base_url: String, headers: HeadersFn, fetch: Option, - transform_request_body: Option, ) -> Self { Self { provider, @@ -78,7 +72,6 @@ impl EndpointConfig { }), fetch, supported_urls: Arc::new(|_| SupportedUrls::default()), - transform_request_body, } } @@ -94,7 +87,6 @@ impl EndpointConfig { endpoint, fetch, supported_urls: Arc::new(|_| SupportedUrls::default()), - transform_request_body: None, } } @@ -130,7 +122,6 @@ impl EndpointConfig { base_url: endpoint.base_url, headers: normalize_headers(combine_headers(&[&endpoint.headers, &call])), fetch: self.fetch.clone(), - transform_request_body: self.transform_request_body.clone(), }) } } @@ -140,7 +131,6 @@ pub(crate) struct Exchange { base_url: String, headers: Vec<(String, String)>, fetch: Option, - transform_request_body: Option, } impl Exchange { @@ -172,12 +162,4 @@ impl Exchange { request.credentialed_origin = Some(self.base_url.clone()); request } - - /// Run the provider-level body rewrite, when there is one. - pub(crate) fn transform_body(&self, body: Value) -> Value { - match &self.transform_request_body { - Some(transform) => transform(body), - None => body, - } - } } diff --git a/aimux-providers/src/shared/mod.rs b/aimux-providers/src/shared/mod.rs index 232105cbc..442f797c2 100644 --- a/aimux-providers/src/shared/mod.rs +++ b/aimux-providers/src/shared/mod.rs @@ -14,10 +14,6 @@ pub(crate) use poll::{ poll_interval_ms, poll_until, retry_download, }; -use std::sync::Arc; - -use serde_json::Value; - use aimux_core::AiMuxError; use aimux_provider_utils::{ HeaderMapOpt, HeadersFn, Resolvable, combine_headers, load_api_key, normalize_headers, @@ -37,10 +33,6 @@ pub(crate) fn is_valid_hostname_part(part: &str) -> bool { && bytes[bytes.len() - 1] != b'-' } -/// A provider-level rewrite of every JSON request body, called once after the -/// body is serialized and before it is sent. -pub type TransformRequestBody = Arc Value + Send + Sync>; - /// Where the credential of a provider comes from. #[derive(Clone, Debug)] pub(crate) enum Credential { diff --git a/aimux-providers/src/stability.rs b/aimux-providers/src/stability.rs index b218cc5ee..9cde9610b 100644 --- a/aimux-providers/src/stability.rs +++ b/aimux-providers/src/stability.rs @@ -216,7 +216,6 @@ impl StabilityProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/tavily.rs b/aimux-providers/src/tavily.rs index 5325da514..ab5af75ec 100644 --- a/aimux-providers/src/tavily.rs +++ b/aimux-providers/src/tavily.rs @@ -137,7 +137,6 @@ impl TavilyProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/tinyfish.rs b/aimux-providers/src/tinyfish.rs index 986967e52..b414b3f53 100644 --- a/aimux-providers/src/tinyfish.rs +++ b/aimux-providers/src/tinyfish.rs @@ -133,7 +133,6 @@ impl TinyfishProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/vertex/anthropic_model.rs b/aimux-providers/src/vertex/anthropic_model.rs index 5f5511e10..d9da000a8 100644 --- a/aimux-providers/src/vertex/anthropic_model.rs +++ b/aimux-providers/src/vertex/anthropic_model.rs @@ -28,7 +28,7 @@ use aimux_core::language_model::SupportedUrls; use aimux_provider_utils::{FetchFunction, Resolvable}; use crate::anthropic::AnthropicMessagesModel; -use crate::anthropic::config::{AnthropicModelConfig, AnthropicModelHooks, TransformRequestBody}; +use crate::anthropic::config::{AnthropicModelConfig, AnthropicModelHooks}; use crate::anthropic::options::options_name_of; use crate::shared::Endpoint; @@ -62,11 +62,7 @@ fn raw_predict_envelope(body: Value) -> Value { } /// The configuration of one request against `endpoint`. -fn request_config( - endpoint: Endpoint, - fetch: Option, - transform_request_body: Option, -) -> AnthropicModelConfig { +fn request_config(endpoint: Endpoint, fetch: Option) -> AnthropicModelConfig { let Endpoint { base_url, headers } = endpoint; let models = base_url.clone(); let url_base = base_url.clone(); @@ -76,7 +72,6 @@ fn request_config( headers: Resolvable::Value(headers), fetch, supported_urls: SupportedUrls::default(), - transform_request_body, base_url, provider_options_name: options_name_of(PROVIDER), hooks: AnthropicModelHooks { @@ -103,25 +98,21 @@ pub(super) fn model( model_id: &str, endpoint: EndpointFn, fetch: Option, - transform_request_body: Option, ) -> VertexAnthropicModel { let resolve_fetch = fetch.clone(); - let resolve_transform = transform_request_body.clone(); let mut config = request_config( Endpoint { base_url: String::new(), headers: Default::default(), }, fetch, - transform_request_body, ); // The model's own fields only describe its identity; every request is // built from the endpoint resolved for it. config.resolve = Some(Arc::new(move || { let endpoint = endpoint.clone(); let fetch = resolve_fetch.clone(); - let transform = resolve_transform.clone(); - Box::pin(async move { Ok(request_config(endpoint().await?, fetch, transform)) }) + Box::pin(async move { Ok(request_config(endpoint().await?, fetch)) }) })); AnthropicMessagesModel::with_config(model_id.to_string(), config) } diff --git a/aimux-providers/src/vertex/anthropic_provider.rs b/aimux-providers/src/vertex/anthropic_provider.rs index 4e731bcd9..ed625bcd0 100644 --- a/aimux-providers/src/vertex/anthropic_provider.rs +++ b/aimux-providers/src/vertex/anthropic_provider.rs @@ -90,7 +90,6 @@ impl VertexAnthropicProvider { Box::pin(async move { resolver.endpoint(super::Publisher::Anthropic).await }) }), self.0.fetch.clone(), - self.0.transform_request_body.clone(), )) } } diff --git a/aimux-providers/src/vertex/embedding.rs b/aimux-providers/src/vertex/embedding.rs index 697d92c81..3ddc706c6 100644 --- a/aimux-providers/src/vertex/embedding.rs +++ b/aimux-providers/src/vertex/embedding.rs @@ -121,7 +121,7 @@ impl EmbeddingModel for VertexEmbeddingModel { let resp = aimux_provider_utils::post_json_to_api( exchange.request(url, options), - exchange.transform_body(Value::Object(body)), + Value::Object(body), aimux_provider_utils::create_json_response_handler(), google_failed_response_handler(), ) @@ -195,7 +195,7 @@ impl EmbeddingModel for VertexEmbeddingModel { let resp = aimux_provider_utils::post_json_to_api( exchange.request(url, options), - exchange.transform_body(Value::Object(body)), + Value::Object(body), aimux_provider_utils::create_json_response_handler(), crate::google::google_failed_response_handler(), ) diff --git a/aimux-providers/src/vertex/gemini_transcription.rs b/aimux-providers/src/vertex/gemini_transcription.rs index 73934964e..0abb78a3d 100644 --- a/aimux-providers/src/vertex/gemini_transcription.rs +++ b/aimux-providers/src/vertex/gemini_transcription.rs @@ -122,7 +122,7 @@ impl TranscriptionModel for VertexGeminiTranscriptionModel { exchange.url(&format!("/models/{}:generateContent", self.model_id)), options, ), - exchange.transform_body(body), + body, aimux_provider_utils::create_json_response_handler::(), crate::google::google_failed_response_handler(), ) diff --git a/aimux-providers/src/vertex/image.rs b/aimux-providers/src/vertex/image.rs index 87b5f3bf3..860072af4 100644 --- a/aimux-providers/src/vertex/image.rs +++ b/aimux-providers/src/vertex/image.rs @@ -146,7 +146,7 @@ impl VertexImageModel { let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url(&self.generate_content_path()), options), - exchange.transform_body(body), + body, aimux_provider_utils::create_json_response_handler(), google_failed_response_handler(), ) diff --git a/aimux-providers/src/vertex/mod.rs b/aimux-providers/src/vertex/mod.rs index 184412bb9..ebaa12cf0 100644 --- a/aimux-providers/src/vertex/mod.rs +++ b/aimux-providers/src/vertex/mod.rs @@ -22,7 +22,7 @@ use aimux_provider_utils::{ use std::sync::OnceLock; use crate::google::GoogleSpeechModel; -use crate::shared::{Endpoint, EndpointConfig, TransformRequestBody}; +use crate::shared::{Endpoint, EndpointConfig}; mod anthropic_model; mod anthropic_provider; @@ -99,9 +99,6 @@ pub struct VertexProviderSettings { /// Custom WebSocket transport for Gemini live transcription. #[cfg(feature = "realtime")] pub web_socket: Option>, - /// Rewrites every JSON request body once, after it is serialized and - /// before it is sent. - pub transform_request_body: Option, } impl std::fmt::Debug for VertexProviderSettings { @@ -116,10 +113,6 @@ impl std::fmt::Debug for VertexProviderSettings { .field("access_token", &self.access_token) .field("google_auth_options", &self.google_auth_options) .field("fetch", &self.fetch.is_some()) - .field( - "transform_request_body", - &self.transform_request_body.is_some(), - ) .finish() } } @@ -303,7 +296,6 @@ pub fn create_google_vertex( fetch: settings.fetch, #[cfg(feature = "realtime")] web_socket: settings.web_socket, - transform_request_body: settings.transform_request_body, }) } @@ -326,7 +318,6 @@ pub struct VertexProvider { fetch: Option, #[cfg(feature = "realtime")] web_socket: Option>, - transform_request_body: Option, } /// The URL patterns the Vertex models fetch themselves @@ -374,7 +365,6 @@ impl VertexProvider { }), fetch: self.fetch.clone(), supported_urls: Arc::new(|_| supported_urls()), - transform_request_body: self.transform_request_body.clone(), } } diff --git a/aimux-providers/src/vertex/model.rs b/aimux-providers/src/vertex/model.rs index 08c7ac2ac..3c3e40315 100644 --- a/aimux-providers/src/vertex/model.rs +++ b/aimux-providers/src/vertex/model.rs @@ -90,7 +90,6 @@ impl LanguageModel for VertexModel { let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); let (body, warnings) = vertex_request_body(&self.model_id, options, false)?; let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body); let resp = aimux_provider_utils::post_json_to_api( exchange.request(self.generate_endpoint(exchange.base_url()), options), body.clone(), @@ -173,7 +172,6 @@ impl LanguageModel for VertexModel { let (body, warnings) = vertex_request_body(&self.model_id, options, true)?; let include_raw_chunks = options.include_raw_chunks == Some(true); let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body); let endpoint = self.stream_endpoint(exchange.base_url()); let resp = aimux_provider_utils::post_json_to_api( exchange.request(endpoint.clone(), options), diff --git a/aimux-providers/src/vertex/transcription.rs b/aimux-providers/src/vertex/transcription.rs index ab3e3d1ef..41f10eae5 100644 --- a/aimux-providers/src/vertex/transcription.rs +++ b/aimux-providers/src/vertex/transcription.rs @@ -220,7 +220,7 @@ impl TranscriptionModel for VertexTranscriptionModel { request.credentialed_origin = Some(url); let resp = aimux_provider_utils::post_json_to_api( request, - exchange.transform_body(request_body.clone()), + request_body.clone(), aimux_provider_utils::create_json_response_handler::(), crate::google::google_failed_response_handler(), ) diff --git a/aimux-providers/src/vertex/video.rs b/aimux-providers/src/vertex/video.rs index fdea68047..93e0c165d 100644 --- a/aimux-providers/src/vertex/video.rs +++ b/aimux-providers/src/vertex/video.rs @@ -184,7 +184,7 @@ impl VideoModel for VertexVideoModel { let resp = aimux_provider_utils::post_json_to_api( exchange.request(url, options), - exchange.transform_body(body), + body, aimux_provider_utils::create_json_response_handler(), crate::google::google_failed_response_handler(), ) @@ -231,7 +231,7 @@ impl VideoModel for VertexVideoModel { let poll_url = exchange.url(&format!("/models/{}:fetchPredictOperation", self.model_id)); let resp = aimux_provider_utils::post_json_to_api( exchange.request(poll_url.clone(), options), - exchange.transform_body(json!({"operationName": operation_name})), + json!({"operationName": operation_name}), aimux_provider_utils::create_json_response_handler::(), crate::google::google_failed_response_handler(), ) diff --git a/aimux-providers/src/voyage/embedding.rs b/aimux-providers/src/voyage/embedding.rs index b52d2aa4f..c58c1949d 100644 --- a/aimux-providers/src/voyage/embedding.rs +++ b/aimux-providers/src/voyage/embedding.rs @@ -81,7 +81,7 @@ impl EmbeddingModel for VoyageEmbeddingModel { // error internally using the shared error structure. let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/embeddings"), options), - exchange.transform_body(Value::Object(body)), + Value::Object(body), aimux_provider_utils::create_json_response_handler(), voyage_failed_response_handler(), ) diff --git a/aimux-providers/src/voyage/mod.rs b/aimux-providers/src/voyage/mod.rs index 0b05fe018..699dc047f 100644 --- a/aimux-providers/src/voyage/mod.rs +++ b/aimux-providers/src/voyage/mod.rs @@ -14,7 +14,6 @@ pub mod embedding; pub(crate) mod options; pub mod reranking; -pub use crate::shared::TransformRequestBody; pub use embedding::VoyageEmbeddingModel; pub use reranking::VoyageRerankingModel; @@ -74,9 +73,6 @@ pub struct VoyageProviderSettings { /// The transport: a mock, a signing decorator, a proxy-aware client. /// `None` uses the process default, resolved per request. pub fetch: Option, - /// Rewrites every JSON request body once, after it is serialized and - /// before it is sent. - pub transform_request_body: Option, } impl std::fmt::Debug for VoyageProviderSettings { @@ -91,10 +87,6 @@ impl std::fmt::Debug for VoyageProviderSettings { ) .field("name", &self.name) .field("fetch", &self.fetch.is_some()) - .field( - "transform_request_body", - &self.transform_request_body.is_some(), - ) .finish() } } @@ -120,7 +112,6 @@ pub fn create_voyage(settings: VoyageProviderSettings) -> Result, - transform_request_body: Option, } impl VoyageProvider { @@ -152,7 +142,6 @@ impl VoyageProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - self.transform_request_body.clone(), ) } diff --git a/aimux-providers/src/voyage/reranking.rs b/aimux-providers/src/voyage/reranking.rs index 834dfa59c..3375c756c 100644 --- a/aimux-providers/src/voyage/reranking.rs +++ b/aimux-providers/src/voyage/reranking.rs @@ -119,7 +119,6 @@ impl RerankingModel for VoyageRerankingModel { } let exchange = self.config.exchange(options.headers.as_ref()).await?; - let body = exchange.transform_body(body); let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/rerank"), options), diff --git a/aimux-providers/src/xai/mod.rs b/aimux-providers/src/xai/mod.rs index 49ab76bbe..49c42d490 100644 --- a/aimux-providers/src/xai/mod.rs +++ b/aimux-providers/src/xai/mod.rs @@ -369,7 +369,6 @@ impl XAIProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } diff --git a/aimux-providers/src/xai/responses/mod.rs b/aimux-providers/src/xai/responses/mod.rs index 0d062001f..84a614b78 100644 --- a/aimux-providers/src/xai/responses/mod.rs +++ b/aimux-providers/src/xai/responses/mod.rs @@ -200,7 +200,7 @@ impl LanguageModel for XaiResponsesModel { 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 body = request_result.body; let provider_tool_names = request_result.provider_tool_names; let resp = aimux_provider_utils::post_json_to_api( @@ -490,7 +490,7 @@ impl LanguageModel for XaiResponsesModel { 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 body = request_result.body; let warnings = request_result.warnings; let provider_tool_names = request_result.provider_tool_names; let endpoint = exchange.url("/responses"); diff --git a/aimux-providers/src/you_com.rs b/aimux-providers/src/you_com.rs index d58b755fc..e69db571d 100644 --- a/aimux-providers/src/you_com.rs +++ b/aimux-providers/src/you_com.rs @@ -134,7 +134,6 @@ impl YouComProvider { self.base_url.clone(), self.headers.clone(), self.fetch.clone(), - None, ) } From b32cf30b1c1c683cd0e1d1f969a4ded916d3013e Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 12:07:20 +0800 Subject: [PATCH 04/16] refactor(providers)!: Anthropic, OpenAI/Azure URL config and SigV4 clock are plain data - anthropic config: drop UrlFn/RequestUrlFn/BodyFn/ResolveFn. url() is base_url + path; MessagesUrl { Messages, RawPredict } selects the messages URL; prepare_body is a fn pointer; AnthropicEndpoint { Fixed, PerRequest(Arc) } replaces the late-binding closure (Vertex implements ResolveEndpoint over its Resolver). - openai config: OpenAIUrl { Base, AzureDeployment, AzureSpeech } replaces UrlFn; azure UrlRules is pub(crate) and gains speech_url(). - SigV4Fetch clock is fn() -> DateTime. Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-provider-utils/src/sigv4_fetch.rs | 13 ++- aimux-providers/src/anthropic/config.rs | 104 +++++++++++------- aimux-providers/src/anthropic/mod.rs | 6 +- aimux-providers/src/anthropic_aws/mod.rs | 6 +- aimux-providers/src/azure/mod.rs | 56 +++++----- aimux-providers/src/openai/config.rs | 47 ++++++-- aimux-providers/src/vertex/anthropic_model.rs | 80 +++++--------- .../src/vertex/anthropic_provider.rs | 6 +- 8 files changed, 177 insertions(+), 141 deletions(-) diff --git a/aimux-provider-utils/src/sigv4_fetch.rs b/aimux-provider-utils/src/sigv4_fetch.rs index 41fa59215..3d019108c 100644 --- a/aimux-provider-utils/src/sigv4_fetch.rs +++ b/aimux-provider-utils/src/sigv4_fetch.rs @@ -13,7 +13,6 @@ //! Reference: use std::fmt; -use std::sync::Arc; use async_trait::async_trait; use chrono::{DateTime, Utc}; @@ -198,7 +197,9 @@ fn hmac_sha256(key: &[u8], data: &[u8]) -> Vec { mac.finalize().into_bytes().to_vec() } -type Clock = Arc DateTime + Send + Sync>; +/// The signing clock: a plain function, so a fixed time in tests is a named +/// function or a non-capturing closure. +type Clock = fn() -> DateTime; /// A [`Fetch`] decorator that signs every request with SigV4 and forwards it /// to the wrapped transport. @@ -229,14 +230,14 @@ impl SigV4Fetch { inner, credentials, service: service.into(), - clock: Arc::new(Utc::now), + clock: Utc::now, } } /// Replace the signing clock (deterministic signatures in tests). #[must_use] - pub fn with_clock(mut self, clock: impl Fn() -> DateTime + Send + Sync + 'static) -> Self { - self.clock = Arc::new(clock); + pub fn with_clock(mut self, clock: fn() -> DateTime) -> Self { + self.clock = clock; self } } @@ -289,7 +290,7 @@ impl Fetch for SigV4Fetch { #[cfg(test)] mod tests { - use std::sync::Mutex; + use std::sync::{Arc, Mutex}; use bytes::Bytes; use chrono::TimeZone; diff --git a/aimux-providers/src/anthropic/config.rs b/aimux-providers/src/anthropic/config.rs index 142fa5757..d42f1305d 100644 --- a/aimux-providers/src/anthropic/config.rs +++ b/aimux-providers/src/anthropic/config.rs @@ -18,27 +18,42 @@ use serde_json::Value; use aimux_core::AiMuxError; use aimux_core::language_model::SupportedUrls; use aimux_provider_utils::{ - ExchangeContext, FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, ResponseHandler, - combine_headers, normalize_headers, + ExchangeContext, FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, Resolvable, + ResponseHandler, combine_headers, normalize_headers, }; use super::convert::RequestProfile; +use crate::shared::Endpoint; + +/// Where the messages call of a model goes. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) enum MessagesUrl { + /// `url("/messages")`: the first-party API and hosts that mirror it. + #[default] + Messages, + /// `{base_url}/{model id}:rawPredict`, or `:streamRawPredict` when + /// streaming: Anthropic on Vertex AI. + RawPredict, +} -/// Maps an endpoint path (`"/messages"`) to the full request URL. -pub(crate) type UrlFn = Arc String + Send + Sync>; - -/// `(model id, streaming)` to the full URL of the messages call, for hosts -/// that put the model and the mode in the path. -pub(crate) type RequestUrlFn = Arc String + Send + Sync>; - -/// Rewrites the request body a host-specific way. -pub(crate) type BodyFn = Arc Value + Send + Sync>; +/// Where a model's endpoint (base URL and provider headers) comes from. +#[derive(Clone, Default)] +pub(crate) enum AnthropicEndpoint { + /// The config's own `base_url` and `headers`. + #[default] + Fixed, + /// Resolved for every request, for hosts whose endpoint or credentials (a + /// Vertex project and location, say) can only be known when a call is + /// made. The config's own `base_url` and `headers` then only describe the + /// model's static identity. + PerRequest(Arc), +} -/// Builds the configuration of one request, for hosts whose endpoint or -/// credentials (a Vertex project and location, say) can only be known when a -/// call is made. -pub(crate) type ResolveFn = - Arc BoxFuture<'static, Result> + Send + Sync>; +/// A host's per-request endpoint lookup (the Vertex resolver, for instance). +pub(crate) trait ResolveEndpoint: Send + Sync { + /// The base URL and provider headers of one request. + fn endpoint(&self) -> BoxFuture<'_, Result>; +} /// Builds the host's error handler for non-2xx responses. pub(crate) type ErrorHandlerFn = fn() -> ResponseHandler; @@ -48,11 +63,10 @@ pub(crate) type ErrorHandlerFn = fn() -> ResponseHandler; /// and `supportsStrictTools` config members, plus its error shape). #[derive(Clone)] pub(crate) struct AnthropicModelHooks { - /// The messages URL when it depends on the model and the mode. `None`: - /// `url("/messages")`. - pub(crate) request_url: Option, + /// Where the messages call goes. + pub(crate) messages_url: MessagesUrl, /// Applied to the finished request body before it is sent. - pub(crate) prepare_body: Option, + pub(crate) prepare_body: Option Value>, /// Parses a failed response. pub(crate) failed_response_handler: ErrorHandlerFn, /// Structured outputs and their beta header. @@ -65,7 +79,7 @@ impl Default for AnthropicModelHooks { /// The first-party API. fn default() -> Self { Self { - request_url: None, + messages_url: MessagesUrl::Messages, prepare_body: None, failed_response_handler: super::anthropic_failed_response_handler, supports_native_structured_output: true, @@ -81,8 +95,6 @@ impl Default for AnthropicModelHooks { pub(crate) struct AnthropicModelConfig { /// The identity the model reports from `provider()`. pub(crate) provider: String, - /// Endpoint path to full URL. - pub(crate) url: UrlFn, /// Provider headers (credential, `anthropic-version`, user headers), /// resolved on every request. pub(crate) headers: HeadersFn, @@ -90,44 +102,58 @@ pub(crate) struct AnthropicModelConfig { pub(crate) fetch: Option, /// URLs the model fetches itself. Empty: the caller downloads them. pub(crate) supported_urls: SupportedUrls, - /// The origin credentialed headers may be sent to. Never read for URLs. + /// The URL every endpoint path is appended to; also the origin + /// credentialed headers may be sent to. pub(crate) base_url: String, /// The providerOptions key read in addition to `anthropic` and written /// on response metadata. pub(crate) provider_options_name: String, /// Host differences. pub(crate) hooks: AnthropicModelHooks, - /// Late binding: when set, every request uses the configuration this - /// returns (its own `resolve` is `None`) instead of the fields above, - /// which then only describe the model's static identity. - pub(crate) resolve: Option, + /// Where `base_url` and `headers` of a request come from. + pub(crate) endpoint: AnthropicEndpoint, } impl AnthropicModelConfig { - /// The configuration to use for one request: itself, or what the host's - /// `resolve` builds now. + /// The configuration to use for one request: itself, or itself with the + /// endpoint the host resolves now. /// /// # Errors /// /// Returns the error of the host's setting or credential loading, for /// instance `AiMuxError::LoadSetting` for an unset Vertex project. pub(crate) async fn resolved(&self) -> Result { - match &self.resolve { - Some(resolve) => resolve().await, - None => Ok(self.clone()), + match &self.endpoint { + AnthropicEndpoint::Fixed => Ok(self.clone()), + AnthropicEndpoint::PerRequest(resolver) => { + let Endpoint { base_url, headers } = resolver.endpoint().await?; + Ok(Self { + base_url, + headers: Resolvable::Value(headers), + endpoint: AnthropicEndpoint::Fixed, + ..self.clone() + }) + } } } - /// The full URL of an endpoint path. + /// The full URL of an endpoint path: `base_url` followed by `path`. pub(crate) fn url(&self, path: &str) -> String { - (self.url)(path) + format!("{}{path}", self.base_url) } /// The URL of the messages call for a model. pub(crate) fn messages_url(&self, model_id: &str, stream: bool) -> String { - match &self.hooks.request_url { - Some(request_url) => request_url(model_id, stream), - None => self.url("/messages"), + match self.hooks.messages_url { + MessagesUrl::Messages => self.url("/messages"), + MessagesUrl::RawPredict => { + let method = if stream { + "streamRawPredict" + } else { + "rawPredict" + }; + format!("{}/{model_id}:{method}", self.base_url) + } } } @@ -214,7 +240,7 @@ impl AnthropicModelConfig { /// Run the host's body preparation, when there is one. pub(crate) fn prepare_body(&self, body: Value) -> Value { - match &self.hooks.prepare_body { + match self.hooks.prepare_body { Some(prepare) => prepare(body), None => body, } diff --git a/aimux-providers/src/anthropic/mod.rs b/aimux-providers/src/anthropic/mod.rs index 4fe28093c..d69a125ee 100644 --- a/aimux-providers/src/anthropic/mod.rs +++ b/aimux-providers/src/anthropic/mod.rs @@ -48,7 +48,7 @@ use aimux_provider_utils::{ use crate::shared::{AuthScheme, Credential, credential_headers}; -use config::{AnthropicModelConfig, AnthropicModelHooks}; +use config::{AnthropicEndpoint, AnthropicModelConfig, AnthropicModelHooks}; pub(crate) fn anthropic_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -236,17 +236,15 @@ pub struct AnthropicProvider { impl AnthropicProvider { /// The model configuration reporting `provider` as its identity. fn model_config(&self, provider: String) -> AnthropicModelConfig { - let base = self.base_url.clone(); AnthropicModelConfig { provider, - url: Arc::new(move |path| format!("{base}{path}")), headers: self.headers.clone(), fetch: self.fetch.clone(), supported_urls: self.supported_urls.clone(), base_url: self.base_url.clone(), provider_options_name: options::options_name_of(&self.name), hooks: AnthropicModelHooks::default(), - resolve: None, + endpoint: AnthropicEndpoint::Fixed, } } diff --git a/aimux-providers/src/anthropic_aws/mod.rs b/aimux-providers/src/anthropic_aws/mod.rs index 1bb90aab9..9c88f86c1 100644 --- a/aimux-providers/src/anthropic_aws/mod.rs +++ b/aimux-providers/src/anthropic_aws/mod.rs @@ -29,7 +29,7 @@ use aimux_provider_utils::{ }; use crate::anthropic::AnthropicMessagesModel; -use crate::anthropic::config::{AnthropicModelConfig, AnthropicModelHooks}; +use crate::anthropic::config::{AnthropicEndpoint, AnthropicModelConfig, AnthropicModelHooks}; use crate::anthropic::options::{CANONICAL, options_name_of}; /// The region used when the settings name none and SigV4 credentials do not @@ -193,17 +193,15 @@ pub struct AnthropicAwsProvider { impl AnthropicAwsProvider { fn model_config(&self) -> AnthropicModelConfig { - let base = self.base_url.clone(); AnthropicModelConfig { provider: self.name.clone(), - url: Arc::new(move |path| format!("{base}{path}")), headers: self.headers.clone(), fetch: self.fetch.clone(), supported_urls: SupportedUrls::default(), base_url: self.base_url.clone(), provider_options_name: self.provider_options_name.clone(), hooks: AnthropicModelHooks::default(), - resolve: None, + endpoint: AnthropicEndpoint::Fixed, } } diff --git a/aimux-providers/src/azure/mod.rs b/aimux-providers/src/azure/mod.rs index cc66a90a8..a5d38eb92 100644 --- a/aimux-providers/src/azure/mod.rs +++ b/aimux-providers/src/azure/mod.rs @@ -44,7 +44,7 @@ use aimux_provider_utils::{ without_trailing_slash, }; -use crate::openai::config::OpenAIModelConfig; +use crate::openai::config::{OpenAIModelConfig, OpenAIUrl}; use crate::openai::responses::ResponsesProfile; use crate::openai::{ OpenAIEmbeddingModel, OpenAIImageModel, OpenAIModel, OpenAIResponsesModel, OpenAISpeechModel, @@ -306,9 +306,9 @@ pub struct AzureOpenAIProvider { fetch: Option, } -/// Everything the URL of a request depends on, owned so a closure can hold it. +/// Everything the URL of a request depends on. #[derive(Clone)] -struct UrlRules { +pub(crate) struct UrlRules { resource_name: Option, base_url: Option, info: BaseUrlInfo, @@ -340,7 +340,7 @@ impl UrlRules { } /// The AI SDK's `url({ path, modelId })`. - fn url(&self, path: &str, model_id: &str) -> Result { + pub(crate) fn url(&self, path: &str, model_id: &str) -> Result { let prefix = self.prefix()?; let info = self.info; let full = if self.use_deployment_based_urls { @@ -361,6 +361,25 @@ impl UrlRules { } Ok(url.to_string()) } + + /// The Azure AI Speech transcription URL: `speech_base_url`, or the + /// resource's Cognitive Services host. + pub(crate) fn speech_url(&self, speech_base_url: Option<&str>) -> Result { + let prefix = match speech_base_url { + Some(prefix) => without_trailing_slash(prefix), + None => { + // Reuse the resource-name validation and lazy environment lookup. + let mut rules = self.clone(); + rules.base_url = None; + rules + .prefix()? + .replace(".openai.azure.com/openai", ".cognitiveservices.azure.com") + } + }; + Ok(format!( + "{prefix}/speechtotext/transcriptions:transcribe?api-version=2025-10-15" + )) + } } /// `URLSearchParams.set`: replace the parameter, or append it. @@ -405,11 +424,12 @@ impl AzureOpenAIProvider { /// The model configuration of a deployment, reporting `provider` as its /// identity. fn model_config(&self, provider: &str, deployment: &str) -> OpenAIModelConfig { - let rules = self.url_rules(); - let deployment = deployment.to_string(); OpenAIModelConfig { provider: provider.to_string(), - url: Arc::new(move |path| rules.url(path, &deployment)), + url: OpenAIUrl::AzureDeployment { + rules: self.url_rules(), + deployment: deployment.to_string(), + }, headers: self.headers.clone(), token_provider: None, fetch: self.fetch.clone(), @@ -470,24 +490,10 @@ impl AzureOpenAIProvider { #[must_use] pub fn transcription(&self, deployment: &str) -> AzureTranscriptionModel { let mut speech = self.model_config("azure.transcription", deployment); - let rules = self.url_rules(); - let base_url = self.speech_base_url.clone(); - speech.url = Arc::new(move |_| { - let prefix = match &base_url { - Some(prefix) => without_trailing_slash(prefix), - None => { - // Reuse the resource-name validation and lazy environment lookup. - let mut rules = rules.clone(); - rules.base_url = None; - rules - .prefix()? - .replace(".openai.azure.com/openai", ".cognitiveservices.azure.com") - } - }; - Ok(format!( - "{prefix}/speechtotext/transcriptions:transcribe?api-version=2025-10-15" - )) - }); + speech.url = OpenAIUrl::AzureSpeech { + rules: self.url_rules(), + speech_base_url: self.speech_base_url.clone(), + }; speech.headers = self.speech_headers.clone(); AzureTranscriptionModel::new( deployment.to_string(), diff --git a/aimux-providers/src/openai/config.rs b/aimux-providers/src/openai/config.rs index c66ffea55..1346e2fed 100644 --- a/aimux-providers/src/openai/config.rs +++ b/aimux-providers/src/openai/config.rs @@ -9,7 +9,6 @@ //! refuses to send credentialed headers to any other origin after a redirect. use std::collections::HashMap; -use std::sync::Arc; use aimux_core::AiMuxError; use aimux_core::language_model::SupportedUrls; @@ -19,11 +18,43 @@ use aimux_provider_utils::{ }; use super::responses::ResponsesProfile; +use crate::azure::UrlRules; -/// Maps an endpoint path (`"/chat/completions"`) to the full request URL. It -/// can fail because a host may only know the address when a call is made (an -/// Azure resource name read from the environment). -pub(crate) type UrlFn = Arc Result + Send + Sync>; +/// How an endpoint path (`"/chat/completions"`) becomes the full request URL. +/// Building it can fail because a host may only know the address when a call +/// is made (an Azure resource name read from the environment). +#[derive(Clone)] +pub(crate) enum OpenAIUrl { + /// The base URL followed by the endpoint path. + Base(String), + /// Azure OpenAI's URL rules for a deployment. + AzureDeployment { rules: UrlRules, deployment: String }, + /// The Azure AI Speech transcription endpoint, whatever the path: the + /// speech base URL, or the resource's Cognitive Services host. + AzureSpeech { + rules: UrlRules, + speech_base_url: Option, + }, +} + +impl OpenAIUrl { + /// The full URL of an endpoint path. + /// + /// # Errors + /// + /// Returns the host's setting error (`AiMuxError::LoadSetting` for an + /// unset Azure resource name). + pub(crate) fn url(&self, path: &str) -> Result { + match self { + Self::Base(base_url) => Ok(format!("{base_url}{path}")), + Self::AzureDeployment { rules, deployment } => rules.url(path, deployment), + Self::AzureSpeech { + rules, + speech_base_url, + } => rules.speech_url(speech_base_url.as_deref()), + } + } +} /// What a model needs to talk to the API. Built by the provider that owns the /// model: `OpenAIProvider`, Azure, Codex and the Hugging Face chat extension. @@ -32,7 +63,7 @@ pub(crate) struct OpenAIModelConfig { /// The identity the model reports from `provider()`. pub(crate) provider: String, /// Endpoint path to full URL. - pub(crate) url: UrlFn, + pub(crate) url: OpenAIUrl, /// Provider headers (credential, organization, project, user headers), /// resolved on every request. pub(crate) headers: HeadersFn, @@ -57,7 +88,7 @@ impl OpenAIModelConfig { ) -> Self { Self { provider, - url: Arc::new(move |path| Ok(format!("{base_url}{path}"))), + url: OpenAIUrl::Base(base_url), headers, token_provider: None, fetch, @@ -74,7 +105,7 @@ impl OpenAIModelConfig { /// Returns the host's setting error (`AiMuxError::LoadSetting` for an /// unset Azure resource name). pub(crate) fn url(&self, path: &str) -> Result { - (self.url)(path) + self.url.url(path) } /// The WebSocket form of an endpoint path: `https` becomes `wss` and diff --git a/aimux-providers/src/vertex/anthropic_model.rs b/aimux-providers/src/vertex/anthropic_model.rs index d9da000a8..5800fc45f 100644 --- a/aimux-providers/src/vertex/anthropic_model.rs +++ b/aimux-providers/src/vertex/anthropic_model.rs @@ -27,8 +27,11 @@ use aimux_core::AiMuxError; use aimux_core::language_model::SupportedUrls; use aimux_provider_utils::{FetchFunction, Resolvable}; +use super::{Publisher, Resolver}; use crate::anthropic::AnthropicMessagesModel; -use crate::anthropic::config::{AnthropicModelConfig, AnthropicModelHooks}; +use crate::anthropic::config::{ + AnthropicEndpoint, AnthropicModelConfig, AnthropicModelHooks, MessagesUrl, ResolveEndpoint, +}; use crate::anthropic::options::options_name_of; use crate::shared::Endpoint; @@ -42,11 +45,6 @@ const PROVIDER: &str = "googleVertex.anthropic.messages"; /// An Anthropic Claude language model served via Vertex AI. pub type VertexAnthropicModel = AnthropicMessagesModel; -/// Resolves the endpoint of a request: the `.../publishers/anthropic/models` -/// base URL and the provider headers. -pub(super) type EndpointFn = - Arc BoxFuture<'static, Result> + Send + Sync>; - /// Wrap a standard Messages request body in the `rawPredict` envelope: drop /// `model` (the URL carries it) and add `anthropic_version`. fn raw_predict_envelope(body: Value) -> Value { @@ -61,58 +59,40 @@ fn raw_predict_envelope(body: Value) -> Value { Value::Object(envelope) } -/// The configuration of one request against `endpoint`. -fn request_config(endpoint: Endpoint, fetch: Option) -> AnthropicModelConfig { - let Endpoint { base_url, headers } = endpoint; - let models = base_url.clone(); - let url_base = base_url.clone(); - AnthropicModelConfig { - provider: PROVIDER.to_string(), - url: Arc::new(move |path| format!("{url_base}{path}")), - headers: Resolvable::Value(headers), - fetch, - supported_urls: SupportedUrls::default(), - base_url, - provider_options_name: options_name_of(PROVIDER), - hooks: AnthropicModelHooks { - request_url: Some(Arc::new(move |model_id, stream| { - let method = if stream { - "streamRawPredict" - } else { - "rawPredict" - }; - format!("{models}/{model_id}:{method}") - })), - prepare_body: Some(Arc::new(raw_predict_envelope)), - failed_response_handler: crate::google::google_failed_response_handler, - supports_native_structured_output: false, - supports_strict_tools: false, - }, - resolve: None, +/// The Vertex endpoint of Anthropic's models, resolved for each request: the +/// `.../publishers/anthropic/models` base URL and the provider headers. +struct VertexAnthropicEndpoint(Arc); + +impl ResolveEndpoint for VertexAnthropicEndpoint { + fn endpoint(&self) -> BoxFuture<'_, Result> { + Box::pin(self.0.endpoint(Publisher::Anthropic)) } } -/// The model `model_id` on the Vertex endpoint `endpoint` resolves to for each +/// The model `model_id` on the Vertex endpoint `resolver` resolves for each /// request. pub(super) fn model( model_id: &str, - endpoint: EndpointFn, + resolver: Arc, fetch: Option, ) -> VertexAnthropicModel { - let resolve_fetch = fetch.clone(); - let mut config = request_config( - Endpoint { - base_url: String::new(), - headers: Default::default(), - }, + // The model's own `base_url` and `headers` only describe its identity; + // every request uses the endpoint resolved for it. + let config = AnthropicModelConfig { + provider: PROVIDER.to_string(), + headers: Resolvable::Value(Default::default()), fetch, - ); - // The model's own fields only describe its identity; every request is - // built from the endpoint resolved for it. - config.resolve = Some(Arc::new(move || { - let endpoint = endpoint.clone(); - let fetch = resolve_fetch.clone(); - Box::pin(async move { Ok(request_config(endpoint().await?, fetch)) }) - })); + supported_urls: SupportedUrls::default(), + base_url: String::new(), + provider_options_name: options_name_of(PROVIDER), + hooks: AnthropicModelHooks { + messages_url: MessagesUrl::RawPredict, + prepare_body: Some(raw_predict_envelope), + failed_response_handler: crate::google::google_failed_response_handler, + supports_native_structured_output: false, + supports_strict_tools: false, + }, + endpoint: AnthropicEndpoint::PerRequest(Arc::new(VertexAnthropicEndpoint(resolver))), + }; AnthropicMessagesModel::with_config(model_id.to_string(), config) } diff --git a/aimux-providers/src/vertex/anthropic_provider.rs b/aimux-providers/src/vertex/anthropic_provider.rs index ed625bcd0..b3a1e0e29 100644 --- a/aimux-providers/src/vertex/anthropic_provider.rs +++ b/aimux-providers/src/vertex/anthropic_provider.rs @@ -82,13 +82,9 @@ pub struct VertexAnthropicProvider(VertexProvider); impl VertexAnthropicProvider { #[must_use] pub fn call(&self, model_id: &str) -> Arc { - let resolver = self.0.resolver.clone(); Arc::new(super::anthropic_model::model( model_id, - Arc::new(move || { - let resolver = resolver.clone(); - Box::pin(async move { resolver.endpoint(super::Publisher::Anthropic).await }) - }), + self.0.resolver.clone(), self.0.fetch.clone(), )) } From 70057b4aa4a1779d34a476582e3c1a54a726a414 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 12:14:19 +0800 Subject: [PATCH 05/16] refactor(providers)!: endpoint config is data, not closures EndpointConfig no longer holds EndpointFn / SupportedUrlsFn closures. Where a request's base URL and headers come from is now described by EndpointSource (Http { BaseUrl, EndpointHeaders } | Vertex | Bedrock | AwsPolly), resolved by EndpointConfig::endpoint() per request; the URL patterns a model fetches itself are SupportedUrlsSource (None | Fixed | Google { base_url }). - BaseUrl: Fixed | Env { var, description } (SearXNG instance URL). - EndpointHeaders: Provider(ProviderHeaders) | Basic { login, password, user } (DataForSEO) | WithUserHeaders { provider, user } (Open Responses' resolvable user headers). - Providers store ProviderHeaders instead of a HeadersFn; user-agent suffixes use ProviderHeaders::with_user_agent. - vertex: GeminiEndpoint and ProjectLocationSource replace the endpoint and ProjectLocationFn closures. - bedrock: BedrockEndpoint replaces the header/endpoint closures; the reranking model reads the region from the shared Auth (RegionFn gone). - aws_polly: PollyEndpoint replaces signing_headers and the dynamic endpoint closure. URLs, headers, error types and when the environment is read are unchanged. --- aimux-providers/src/assemblyai.rs | 8 +- aimux-providers/src/aws_polly.rs | 70 +++--- aimux-providers/src/bedrock/mod.rs | 149 +++++++------ aimux-providers/src/bedrock/reranking.rs | 13 +- aimux-providers/src/black_forest_labs.rs | 10 +- aimux-providers/src/cartesia.rs | 8 +- aimux-providers/src/cohere/mod.rs | 16 +- aimux-providers/src/dataforseo.rs | 52 ++--- aimux-providers/src/deepgram.rs | 8 +- aimux-providers/src/deepseek/mod.rs | 29 ++- aimux-providers/src/deepseek/model.rs | 2 +- aimux-providers/src/elevenlabs.rs | 10 +- aimux-providers/src/exa_ai.rs | 8 +- aimux-providers/src/fal.rs | 8 +- aimux-providers/src/firecrawl.rs | 8 +- aimux-providers/src/gladia.rs | 10 +- aimux-providers/src/google/mod.rs | 56 ++--- aimux-providers/src/google/model.rs | 2 +- aimux-providers/src/google/video.rs | 2 +- aimux-providers/src/google_pse.rs | 8 +- aimux-providers/src/groq/mod.rs | 27 ++- aimux-providers/src/huggingface.rs | 10 +- aimux-providers/src/hume.rs | 8 +- aimux-providers/src/jina_ai.rs | 8 +- aimux-providers/src/klingai.rs | 8 +- aimux-providers/src/linkup.rs | 8 +- aimux-providers/src/lmnt.rs | 8 +- aimux-providers/src/luma.rs | 10 +- aimux-providers/src/mistral/mod.rs | 29 ++- aimux-providers/src/mistral/model.rs | 2 +- aimux-providers/src/open_responses.rs | 26 +-- aimux-providers/src/parallel_ai.rs | 8 +- aimux-providers/src/prodia.rs | 6 +- aimux-providers/src/recraft.rs | 8 +- aimux-providers/src/replicate.rs | 6 +- aimux-providers/src/revai.rs | 10 +- aimux-providers/src/runwayml.rs | 8 +- aimux-providers/src/searxng.rs | 49 ++--- aimux-providers/src/serper.rs | 8 +- aimux-providers/src/shared/exchange.rs | 201 +++++++++++++++--- aimux-providers/src/shared/mod.rs | 65 ++++-- aimux-providers/src/stability.rs | 8 +- aimux-providers/src/tavily.rs | 8 +- aimux-providers/src/tinyfish.rs | 8 +- .../src/vertex/gemini_transcription.rs | 10 +- aimux-providers/src/vertex/mod.rs | 134 +++++++----- aimux-providers/src/vertex/model.rs | 2 +- aimux-providers/src/vertex/transcription.rs | 8 +- aimux-providers/src/voyage/mod.rs | 8 +- aimux-providers/src/xai/mod.rs | 27 ++- aimux-providers/src/you_com.rs | 8 +- 51 files changed, 668 insertions(+), 543 deletions(-) diff --git a/aimux-providers/src/assemblyai.rs b/aimux-providers/src/assemblyai.rs index 143c91925..0a9ca639b 100644 --- a/aimux-providers/src/assemblyai.rs +++ b/aimux-providers/src/assemblyai.rs @@ -29,10 +29,10 @@ use aimux_core::transcription_model::{ TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; use aimux_provider_utils::HttpBody; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use crate::shared::{ - AuthScheme, Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, credential_headers, + AuthScheme, Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, ProviderHeaders, is_poll_control_key, poll_interval_ms, poll_until, }; @@ -142,7 +142,7 @@ pub fn create_assemblyai( Ok(AssemblyAIProvider { name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), base_url, - headers: credential_headers( + headers: ProviderHeaders::new( Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "AssemblyAI"), AuthScheme::Header("Authorization"), Vec::new(), @@ -168,7 +168,7 @@ pub fn assemblyai() -> &'static AssemblyAIProvider { pub struct AssemblyAIProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/aws_polly.rs b/aimux-providers/src/aws_polly.rs index 39540ace3..34f6b7106 100644 --- a/aimux-providers/src/aws_polly.rs +++ b/aimux-providers/src/aws_polly.rs @@ -35,11 +35,11 @@ use aimux_core::speech_model::{ }; use aimux_provider_utils::HttpBody; use aimux_provider_utils::{ - AwsCredentials, FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, SigV4Fetch, default_fetch, + AwsCredentials, FetchFunction, HeaderMapOpt, Resolvable, SigV4Fetch, default_fetch, load_optional_setting, load_setting, validate_base_url, }; -use crate::shared::{Endpoint, EndpointConfig}; +use crate::shared::{Endpoint, EndpointConfig, EndpointSource}; /// AWS service name used for SigV4 signing. const SERVICE_NAME: &str = "polly"; @@ -235,17 +235,29 @@ impl Auth { } } -/// The provider headers: only the caller's, after checking that signing has -/// its credentials (the signature itself is made by the transport). -fn signing_headers(auth: Auth, user: Option) -> HeadersFn { - Resolvable::from_async_fn(move || { - let auth = auth.clone(); - let user = user.clone(); - async move { - auth.preflight()?; - Ok(user.unwrap_or_default()) - } - }) +/// The endpoint of a Polly request: the explicit base URL or the regional +/// host, and only the caller's headers, after checking that signing has its +/// credentials (the signature itself is made by the transport). +#[derive(Clone)] +pub(crate) struct PollyEndpoint { + auth: Auth, + base_url: Option, + user: Option, +} + +impl PollyEndpoint { + /// The base URL and provider headers of one request. + pub(crate) async fn resolve(&self) -> Result { + let base_url = self + .base_url + .clone() + .unwrap_or_else(|| format!("https://polly.{}.amazonaws.com", self.auth.region())); + self.auth.preflight()?; + Ok(Endpoint { + base_url, + headers: self.user.clone().unwrap_or_default(), + }) + } } /// Create an Amazon Polly provider. @@ -271,7 +283,6 @@ pub fn create_aws_polly( credential_provider: settings.credential_provider, }; let signing_auth = auth.clone(); - let headers_auth = auth.clone(); let fetch: FetchFunction = Arc::new(SigV4Fetch::new( settings.fetch.unwrap_or_else(default_fetch), Resolvable::from_async_fn(move || { @@ -282,10 +293,12 @@ pub fn create_aws_polly( )); Ok(AwsPollyProvider { name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), - auth, - base_url, // SigV4 signs in the transport; the provider headers carry no credential. - headers: signing_headers(headers_auth, settings.headers), + endpoint: PollyEndpoint { + auth, + base_url, + user: settings.headers, + }, fetch, }) } @@ -304,32 +317,15 @@ pub fn aws_polly() -> &'static AwsPollyProvider { /// An Amazon Polly provider. Speech only; it holds no HTTP client. pub struct AwsPollyProvider { name: String, - auth: Auth, - base_url: Option, - headers: HeadersFn, + endpoint: PollyEndpoint, fetch: FetchFunction, } impl AwsPollyProvider { fn model_config(&self, method: &str) -> EndpointConfig { - let auth = self.auth.clone(); - let base_url = self.base_url.clone(); - let headers = self.headers.clone(); - EndpointConfig::dynamic( + EndpointConfig::new( format!("{}.{method}", self.name), - Arc::new(move || { - let auth = auth.clone(); - let base_url = base_url.clone(); - let headers = headers.clone(); - Box::pin(async move { - Ok(Endpoint { - base_url: base_url.unwrap_or_else(|| { - format!("https://polly.{}.amazonaws.com", auth.region()) - }), - headers: headers.resolve().await?, - }) - }) - }), + EndpointSource::AwsPolly(self.endpoint.clone()), Some(self.fetch.clone()), ) } diff --git a/aimux-providers/src/bedrock/mod.rs b/aimux-providers/src/bedrock/mod.rs index 713055e6c..d51c0bd7c 100644 --- a/aimux-providers/src/bedrock/mod.rs +++ b/aimux-providers/src/bedrock/mod.rs @@ -46,26 +46,23 @@ 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, SupportedUrls}; +use aimux_core::language_model::LanguageModel; use aimux_core::model_catalogue::RuntimeModel; use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_core::reranking_model::RerankingModel; use aimux_provider_utils::{ AwsCredentials, Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, - HeadersFn, Resolvable, SigV4Fetch, combine_headers, default_fetch, load_optional_setting, - load_setting, validate_base_url, without_trailing_slash, + Resolvable, SigV4Fetch, combine_headers, default_fetch, load_optional_setting, load_setting, + normalize_headers, validate_base_url, with_user_agent_suffix, without_trailing_slash, }; -use crate::shared::{Endpoint, EndpointConfig, is_valid_hostname_part}; +use crate::shared::{Endpoint, EndpointConfig, EndpointSource, is_valid_hostname_part}; pub use embedding::BedrockEmbeddingModel; pub use image::BedrockImageModel; pub use model::BedrockModel; pub use reranking::BedrockRerankingModel; -/// The region of a request, resolved when the request is made. -pub(crate) type RegionFn = Arc Result + Send + Sync>; - const PROVIDER: &str = "amazon-bedrock"; const SIGV4_SERVICE: &str = "bedrock"; @@ -223,7 +220,7 @@ impl std::fmt::Debug for AmazonBedrockProviderSettings { /// Where the region, the API key and the AWS credentials of a request come /// from. -struct Auth { +pub(crate) struct Auth { region: Option, api_key: Option>, access_key_id: Option, @@ -246,7 +243,7 @@ impl Auth { /// The region: the setting, `AWS_REGION`, or the region of a plain-value /// credential provider. - fn region(&self) -> Result { + pub(crate) fn region(&self) -> Result { let region = if let Some(region) = load_optional_setting(self.region.as_deref(), REGION_ENV_VAR) { region @@ -329,7 +326,7 @@ impl Auth { explicit: Option<&str>, service: &str, endpoint_env_var: &str, - region: &dyn Fn() -> Result, + auth: &Self, ) -> Result { if let Some(url) = explicit { return Ok(url.to_string()); @@ -339,7 +336,7 @@ impl Auth { { return validate_base_url(&url); } - let region = region()?; + let region = auth.region()?; let suffix = [ ("cn-", "amazonaws.com.cn"), ("us-iso-", "c2s.ic.gov"), @@ -397,32 +394,6 @@ pub fn create_amazon_bedrock( credential_provider: settings.credential_provider, }); - let user_headers = settings.headers; - let header_auth = auth.clone(); - let headers: HeadersFn = Resolvable::from_async_fn(move || { - let auth = header_auth.clone(); - let mut user = combine_headers(&[&user_headers.clone().unwrap_or_default()]); - let suffix = concat!("ai-sdk-amazon-bedrock/", env!("CARGO_PKG_VERSION")); - let agent = user - .get("user-agent") - .and_then(|value| value.as_deref()) - .map_or_else(|| suffix.to_string(), |value| format!("{value} {suffix}")); - user.insert("user-agent".into(), Some(agent)); - async move { - match auth.bearer().await? { - Some(key) => { - let mut bearer = HeaderMapOpt::new(); - bearer.insert("Authorization".to_string(), Some(format!("Bearer {key}"))); - Ok(combine_headers(&[&user, &bearer])) - } - None => { - auth.preflight()?; - Ok(user) - } - } - } - }); - let credentials_auth = auth.clone(); let inner = settings.fetch.unwrap_or_else(default_fetch); let fetch: FetchFunction = Arc::new(BedrockAuthFetch { @@ -440,11 +411,7 @@ pub fn create_amazon_bedrock( Ok(AmazonBedrockProvider { auth, base_url: base_url.map(|url| without_trailing_slash(&url)), - headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( - headers, - "amazon-bedrock", - "5.0.100", - ), + user_headers: settings.headers, fetch, }) } @@ -466,10 +433,68 @@ pub fn amazon_bedrock() -> &'static AmazonBedrockProvider { pub struct AmazonBedrockProvider { auth: Arc, base_url: Option, - headers: HeadersFn, + user_headers: Option, fetch: FetchFunction, } +/// The endpoint of a Bedrock request: the explicit base URL, the service's +/// endpoint environment variable or the regional host; the caller's headers +/// with the user-agent suffixes, and an API key as `Authorization: Bearer` +/// when there is one (otherwise the transport signs with SigV4, after +/// checking here that signing has its inputs). +#[derive(Clone)] +pub(crate) struct BedrockEndpoint { + auth: Arc, + base_url: Option, + /// `bedrock-runtime` or `bedrock-agent-runtime`. + service: &'static str, + endpoint_env_var: &'static str, + user: Option, +} + +impl BedrockEndpoint { + /// The base URL and provider headers of one request. + pub(crate) async fn resolve(&self) -> Result { + let base_url = Auth::base_url( + self.base_url.as_deref(), + self.service, + self.endpoint_env_var, + &self.auth, + )?; + Ok(Endpoint { + base_url, + headers: self.headers().await?, + }) + } + + async fn headers(&self) -> Result { + let mut user = combine_headers(&[&self.user.clone().unwrap_or_default()]); + let suffix = concat!("ai-sdk-amazon-bedrock/", env!("CARGO_PKG_VERSION")); + let agent = user + .get("user-agent") + .and_then(|value| value.as_deref()) + .map_or_else(|| suffix.to_string(), |value| format!("{value} {suffix}")); + user.insert("user-agent".into(), Some(agent)); + let headers = match self.auth.bearer().await? { + Some(key) => { + let mut bearer = HeaderMapOpt::new(); + bearer.insert("Authorization".to_string(), Some(format!("Bearer {key}"))); + combine_headers(&[&user, &bearer]) + } + None => { + self.auth.preflight()?; + user + } + }; + let mut headers = normalize_headers(headers).into_iter().collect(); + with_user_agent_suffix(&mut headers, "ai-sdk-amazon-bedrock/5.0.100"); + Ok(headers + .into_iter() + .map(|(name, value)| (name, Some(value))) + .collect()) + } +} + impl AmazonBedrockProvider { /// The configuration of a model that calls `service` (`bedrock-runtime` /// or `bedrock-agent-runtime`). @@ -478,32 +503,17 @@ impl AmazonBedrockProvider { service: &'static str, endpoint_env_var: &'static str, ) -> EndpointConfig { - let auth = self.auth.clone(); - let base_url = self.base_url.clone(); - let headers = self.headers.clone(); - EndpointConfig { - provider: PROVIDER.to_string(), - endpoint: Arc::new(move || { - let auth = auth.clone(); - let base_url = base_url.clone(); - let headers = headers.clone(); - Box::pin(async move { - let region_auth = auth.clone(); - let base_url = Auth::base_url( - base_url.as_deref(), - service, - endpoint_env_var, - &move || region_auth.region(), - )?; - Ok(Endpoint { - base_url, - headers: headers.resolve().await?, - }) - }) + EndpointConfig::new( + PROVIDER.to_string(), + EndpointSource::Bedrock(BedrockEndpoint { + auth: self.auth.clone(), + base_url: self.base_url.clone(), + service, + endpoint_env_var, + user: self.user_headers.clone(), }), - fetch: Some(self.fetch.clone()), - supported_urls: Arc::new(|_| SupportedUrls::default()), - } + Some(self.fetch.clone()), + ) } fn runtime_config(&self) -> EndpointConfig { @@ -568,11 +578,10 @@ impl AmazonBedrockProvider { /// `rerank` endpoint; `provider()` is `"amazon-bedrock"`. #[must_use] pub fn reranking(&self, model_id: &str) -> BedrockRerankingModel { - let auth = self.auth.clone(); BedrockRerankingModel::from_config( model_id.to_string(), self.model_config("bedrock-agent-runtime", AGENT_RUNTIME_ENDPOINT_ENV_VAR), - Arc::new(move || auth.region()), + self.auth.clone(), ) } diff --git a/aimux-providers/src/bedrock/reranking.rs b/aimux-providers/src/bedrock/reranking.rs index 49ad3f157..82b07aa3f 100644 --- a/aimux-providers/src/bedrock/reranking.rs +++ b/aimux-providers/src/bedrock/reranking.rs @@ -5,6 +5,8 @@ //! //! Uses the Bedrock Agent Runtime `/rerank` endpoint. +use std::sync::Arc; + use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; @@ -18,7 +20,7 @@ use aimux_core::shared::SharedProviderOptions; use aimux_provider_utils::HttpBody; -use super::RegionFn; +use super::Auth; use super::options; use crate::shared::EndpointConfig; @@ -76,15 +78,16 @@ struct BedrockRerankingResult { pub struct BedrockRerankingModel { model_id: String, config: EndpointConfig, - region: RegionFn, + /// The region of a request (the model ARN names it), resolved per request. + auth: Arc, } impl BedrockRerankingModel { - pub(crate) fn from_config(model_id: String, config: EndpointConfig, region: RegionFn) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig, auth: Arc) -> Self { Self { model_id, config, - region, + auth, } } } @@ -105,7 +108,7 @@ impl RerankingModel for BedrockRerankingModel { ) -> Result { let bedrock_options = parse_bedrock_reranking_options(options.provider_options.as_ref())?; - let region = (self.region)()?; + let region = self.auth.region()?; let model_arn = format!( "arn:aws:bedrock:{region}::foundation-model/{}", self.model_id diff --git a/aimux-providers/src/black_forest_labs.rs b/aimux-providers/src/black_forest_labs.rs index e397277f1..40e5ac0f7 100644 --- a/aimux-providers/src/black_forest_labs.rs +++ b/aimux-providers/src/black_forest_labs.rs @@ -26,11 +26,11 @@ use aimux_core::image_model::{ }; use aimux_core::shared::Warning; use aimux_provider_utils::HttpRequest; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use crate::shared::{ - Credential, EndpointConfig, POLL_INTERVAL_MILLIS_KEY, PollStep, is_poll_control_key, - poll_interval_ms, poll_until, provider_headers, retry_download, + Credential, EndpointConfig, POLL_INTERVAL_MILLIS_KEY, PollStep, ProviderHeaders, + is_poll_control_key, poll_interval_ms, poll_until, retry_download, }; /// AI SDK's `isTrustedUrl` (black-forest-labs-api.ts): credentials may go to @@ -156,7 +156,7 @@ pub fn create_black_forest_labs( Ok(BlackForestLabsProvider { name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), base_url, - headers: provider_headers( + headers: ProviderHeaders::bearer( Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Black Forest Labs"), Vec::new(), settings.headers, @@ -181,7 +181,7 @@ pub fn black_forest_labs() -> &'static BlackForestLabsProvider { pub struct BlackForestLabsProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/cartesia.rs b/aimux-providers/src/cartesia.rs index 57ae2d463..ff3424071 100644 --- a/aimux-providers/src/cartesia.rs +++ b/aimux-providers/src/cartesia.rs @@ -34,9 +34,9 @@ use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; use aimux_provider_utils::HttpBody; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; fn cartesia_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -133,7 +133,7 @@ pub fn create_cartesia(settings: CartesiaProviderSettings) -> Result &'static CartesiaProvider { pub struct CartesiaProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/cohere/mod.rs b/aimux-providers/src/cohere/mod.rs index 5fe4c451a..18f5d0daa 100644 --- a/aimux-providers/src/cohere/mod.rs +++ b/aimux-providers/src/cohere/mod.rs @@ -33,11 +33,9 @@ use aimux_core::language_model::LanguageModel; use aimux_core::model_catalogue::RuntimeModel; use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_core::reranking_model::RerankingModel; -use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, without_trailing_slash, -}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, without_trailing_slash}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; pub(crate) fn cohere_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -111,7 +109,7 @@ pub fn create_cohere(settings: CohereProviderSettings) -> Result without_trailing_slash(url), None => DEFAULT_BASE_URL.to_string(), }; - let headers = provider_headers( + let headers = ProviderHeaders::bearer( Credential::explicit_or_env( settings.api_key.map(Resolvable::Value), API_KEY_ENV_VAR, @@ -120,11 +118,7 @@ pub fn create_cohere(settings: CohereProviderSettings) -> Result &'static CohereProvider { pub struct CohereProvider { generate_id: Option, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/dataforseo.rs b/aimux-providers/src/dataforseo.rs index eb84838bf..405a7f354 100644 --- a/aimux-providers/src/dataforseo.rs +++ b/aimux-providers/src/dataforseo.rs @@ -16,7 +16,6 @@ use std::sync::OnceLock; use async_trait::async_trait; -use base64::Engine; use serde::Deserialize; use serde_json::{Value, json}; @@ -24,11 +23,9 @@ use aimux_core::error::AiMuxError; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, combine_headers, validate_base_url, -}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig}; +use crate::shared::{Credential, EndpointConfig, EndpointHeaders}; /// Fixed model id for the DataForSEO search model. const MODEL_ID: &str = "dataforseo-search"; @@ -110,35 +107,6 @@ impl std::fmt::Debug for DataforseoProviderSettings { } } -/// The provider headers: `Authorization: Basic base64(login:password)` with -/// both parts resolved for every request, then the caller's headers. -fn basic_auth_headers( - login: Credential, - password: Credential, - user: Option, -) -> HeadersFn { - Resolvable::from_async_fn(move || { - let login = login.clone(); - let password = password.clone(); - let user = user.clone(); - async move { - let login = login.secret().await?.unwrap_or_default(); - let password = password.secret().await?.unwrap_or_default(); - let encoded = base64::engine::general_purpose::STANDARD - .encode(format!("{login}:{password}").as_bytes()); - let mut layer = HeaderMapOpt::new(); - layer.insert( - "Authorization".to_string(), - Some(format!("Basic {encoded}")), - ); - Ok(match &user { - Some(user) => combine_headers(&[&layer, user]), - None => layer, - }) - } - }) -} - /// Create a DataForSEO provider. /// /// # Errors @@ -156,11 +124,15 @@ pub fn create_dataforseo( Ok(DataforseoProvider { name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), base_url, - headers: basic_auth_headers( - Credential::explicit_or_env(settings.login, LOGIN_ENV_VAR, "DataForSEO login"), - Credential::explicit_or_env(settings.password, PASSWORD_ENV_VAR, "DataForSEO password"), - settings.headers, - ), + headers: EndpointHeaders::Basic { + login: Credential::explicit_or_env(settings.login, LOGIN_ENV_VAR, "DataForSEO login"), + password: Credential::explicit_or_env( + settings.password, + PASSWORD_ENV_VAR, + "DataForSEO password", + ), + user: settings.headers, + }, fetch: settings.fetch, }) } @@ -180,7 +152,7 @@ pub fn dataforseo() -> &'static DataforseoProvider { pub struct DataforseoProvider { name: String, base_url: String, - headers: HeadersFn, + headers: EndpointHeaders, fetch: Option, } diff --git a/aimux-providers/src/deepgram.rs b/aimux-providers/src/deepgram.rs index 9eef9ec96..12c4ce992 100644 --- a/aimux-providers/src/deepgram.rs +++ b/aimux-providers/src/deepgram.rs @@ -28,9 +28,9 @@ use aimux_core::transcription_model::{ TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; use aimux_provider_utils::HttpBody; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders}; /// Deepgram errors: `{"err_code": "...", "err_msg": "...", "request_id": ...}` /// on most endpoints; some return `{"category": "...", "message": "...", @@ -131,7 +131,7 @@ pub fn create_deepgram(settings: DeepgramProviderSettings) -> Result &'static DeepgramProvider { pub struct DeepgramProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/deepseek/mod.rs b/aimux-providers/src/deepseek/mod.rs index 5706df4f6..5efb2f3e6 100644 --- a/aimux-providers/src/deepseek/mod.rs +++ b/aimux-providers/src/deepseek/mod.rs @@ -31,9 +31,9 @@ use aimux_core::image_model::ImageModel; use aimux_core::language_model::{LanguageModel, SupportedUrls}; use aimux_core::model_catalogue::RuntimeModel; use aimux_core::provider::{Provider, ProviderDiscovery}; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders, SupportedUrlsSource}; use model::DeepSeekChatConfig; const DEFAULT_BASE_URL: &str = "https://api.deepseek.com"; @@ -91,19 +91,16 @@ pub fn create_deepseek(settings: DeepSeekProviderSettings) -> Result &'static DeepSeekProvider { pub struct DeepSeekProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } @@ -145,7 +142,7 @@ impl DeepSeekProvider { DeepSeekChatConfig { endpoint: self .endpoint_config("chat") - .with_supported_urls(Arc::new(|_| chat_supported_urls())), + .with_supported_urls(SupportedUrlsSource::Fixed(chat_supported_urls())), supports_assistant_prefix_completion: self.base_url.ends_with("/beta"), supports_strict_tool_calls: self.base_url.ends_with("/beta"), }, diff --git a/aimux-providers/src/deepseek/model.rs b/aimux-providers/src/deepseek/model.rs index 0ffbf96aa..8d9bd469d 100644 --- a/aimux-providers/src/deepseek/model.rs +++ b/aimux-providers/src/deepseek/model.rs @@ -358,7 +358,7 @@ impl LanguageModel for DeepSeekChatLanguageModel { } fn supported_urls(&self) -> SupportedUrls { - (self.config.endpoint.supported_urls)(&self.model_id) + self.config.endpoint.supported_urls(&self.model_id) } async fn do_generate(&self, options: &CallOptions) -> Result { diff --git a/aimux-providers/src/elevenlabs.rs b/aimux-providers/src/elevenlabs.rs index 6688ba1af..1850a6dbd 100644 --- a/aimux-providers/src/elevenlabs.rs +++ b/aimux-providers/src/elevenlabs.rs @@ -28,11 +28,9 @@ use aimux_core::embedding_model::EmbeddingModel; use aimux_core::image_model::ImageModel; use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; -use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, HttpBody, Resolvable, validate_base_url, -}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HttpBody, Resolvable, validate_base_url}; -use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders}; pub(crate) mod options; @@ -150,7 +148,7 @@ pub fn create_elevenlabs( Ok(ElevenLabsProvider { name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), base_url, - headers: credential_headers( + headers: ProviderHeaders::new( Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "ElevenLabs"), AuthScheme::Header("xi-api-key"), Vec::new(), @@ -181,7 +179,7 @@ pub fn elevenlabs() -> &'static ElevenLabsProvider { pub struct ElevenLabsProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, #[cfg(feature = "realtime")] web_socket: Option>, diff --git a/aimux-providers/src/exa_ai.rs b/aimux-providers/src/exa_ai.rs index c1cd99cbe..3fddfcaa7 100644 --- a/aimux-providers/src/exa_ai.rs +++ b/aimux-providers/src/exa_ai.rs @@ -19,9 +19,9 @@ use aimux_core::error::AiMuxError; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders}; fn exa_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -104,7 +104,7 @@ pub fn create_exa_ai(settings: ExaAiProviderSettings) -> Result &'static ExaAiProvider { pub struct ExaAiProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/fal.rs b/aimux-providers/src/fal.rs index c38820160..d0a731191 100644 --- a/aimux-providers/src/fal.rs +++ b/aimux-providers/src/fal.rs @@ -26,10 +26,10 @@ use aimux_core::transcription_model::{ TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; use aimux_provider_utils::HttpRequest; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use crate::shared::{ - AuthScheme, Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, credential_headers, + AuthScheme, Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, ProviderHeaders, is_poll_control_key, poll_interval_ms, poll_until, retry_download, }; @@ -146,7 +146,7 @@ pub fn create_fal(settings: FalProviderSettings) -> Result &'static FalProvider { pub struct FalProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/firecrawl.rs b/aimux-providers/src/firecrawl.rs index 408cd831d..ecc684ad5 100644 --- a/aimux-providers/src/firecrawl.rs +++ b/aimux-providers/src/firecrawl.rs @@ -19,9 +19,9 @@ use aimux_core::error::AiMuxError; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; const MODEL_ID: &str = "firecrawl-search"; @@ -103,7 +103,7 @@ pub fn create_firecrawl( Ok(FirecrawlProvider { name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), base_url, - headers: provider_headers( + headers: ProviderHeaders::bearer( Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Firecrawl"), Vec::new(), settings.headers, @@ -128,7 +128,7 @@ pub fn firecrawl() -> &'static FirecrawlProvider { pub struct FirecrawlProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/gladia.rs b/aimux-providers/src/gladia.rs index 410a28ea3..76362a441 100644 --- a/aimux-providers/src/gladia.rs +++ b/aimux-providers/src/gladia.rs @@ -26,12 +26,12 @@ use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use aimux_provider_utils::{HttpBody, MultipartForm, media_type_to_extension}; use crate::shared::{ - Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, is_poll_control_key, - poll_interval_ms, poll_until, provider_headers, + Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, ProviderHeaders, + is_poll_control_key, poll_interval_ms, poll_until, }; fn gladia_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -127,7 +127,7 @@ pub fn create_gladia(settings: GladiaProviderSettings) -> Result &'static GladiaProvider { pub struct GladiaProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/google/mod.rs b/aimux-providers/src/google/mod.rs index 489755717..96220765d 100644 --- a/aimux-providers/src/google/mod.rs +++ b/aimux-providers/src/google/mod.rs @@ -58,10 +58,10 @@ use aimux_core::speech_model::SpeechModel; use aimux_core::transcription_model::TranscriptionModel; use aimux_core::video_model::VideoModel; use aimux_provider_utils::{ - EvaluationLanguageModel, FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url, + EvaluationLanguageModel, FetchFunction, HeaderMapOpt, Resolvable, validate_base_url, }; -use crate::shared::{AuthScheme, Credential, Endpoint, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders, SupportedUrlsSource}; pub(crate) fn google_stream_error( error: &types::GoogleError, @@ -151,7 +151,7 @@ fn supports_external_file_urls(model_id: &str) -> bool { /// The URL patterns the API fetches itself (`getSupportedUrls` in /// `createGoogle`): Files API URLs and YouTube links for every media type, /// plus external `https` URLs for the models that accept them. -fn supported_urls(base_url: &str, model_id: Option<&str>) -> SupportedUrls { +pub(crate) fn supported_urls(base_url: &str, model_id: Option<&str>) -> SupportedUrls { let include_external = model_id.is_none_or(supports_external_file_urls); let pattern = |source: &str| regex::Regex::new(source).expect("static pattern"); let mut urls = std::collections::HashMap::new(); @@ -243,20 +243,17 @@ pub fn create_google(settings: GoogleProviderSettings) -> Result &'static GoogleProvider { pub struct GoogleProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, generate_id: Option, #[cfg(feature = "realtime")] @@ -290,24 +287,15 @@ pub struct GoogleProvider { impl GoogleProvider { /// The model configuration reporting `provider` as its identity. fn model_config(&self, provider: String) -> EndpointConfig { - let base_url = self.base_url.clone(); - let headers = self.headers.clone(); - let urls_base = self.base_url.clone(); - EndpointConfig { + EndpointConfig::fixed( provider, - endpoint: Arc::new(move || { - let base_url = base_url.clone(); - let headers = headers.clone(); - Box::pin(async move { - Ok(Endpoint { - base_url, - headers: headers.resolve().await?, - }) - }) - }), - fetch: self.fetch.clone(), - supported_urls: Arc::new(move |model_id| supported_urls(&urls_base, Some(model_id))), - } + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + ) + .with_supported_urls(SupportedUrlsSource::Google { + base_url: self.base_url.clone(), + }) } /// A language model; `provider()` is the provider name diff --git a/aimux-providers/src/google/model.rs b/aimux-providers/src/google/model.rs index f5cb9947d..daadf682f 100644 --- a/aimux-providers/src/google/model.rs +++ b/aimux-providers/src/google/model.rs @@ -79,7 +79,7 @@ impl LanguageModel for GoogleModel { } fn supported_urls(&self) -> SupportedUrls { - (self.config.supported_urls)(&self.model_id) + self.config.supported_urls(&self.model_id) } async fn do_generate(&self, options: &CallOptions) -> Result { diff --git a/aimux-providers/src/google/video.rs b/aimux-providers/src/google/video.rs index c4438c7d5..c29c71331 100644 --- a/aimux-providers/src/google/video.rs +++ b/aimux-providers/src/google/video.rs @@ -260,7 +260,7 @@ impl VideoModel for GoogleVideoModel { "No videos in response. Response: {raw_body}" )) })?; - let endpoint = (self.config.endpoint)().await?; + let endpoint = self.config.endpoint().await?; let api_key = endpoint .headers .iter() diff --git a/aimux-providers/src/google_pse.rs b/aimux-providers/src/google_pse.rs index 60d82ece5..32a723767 100644 --- a/aimux-providers/src/google_pse.rs +++ b/aimux-providers/src/google_pse.rs @@ -25,9 +25,9 @@ use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders}; /// Fixed model ID for the Google PSE search model. const MODEL_ID: &str = "google-pse-search"; @@ -129,7 +129,7 @@ pub fn create_google_pse( credential: Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Google PSE"), cx: settings.cx, // The key travels in the query string, so the headers carry none. - headers: credential_headers( + headers: ProviderHeaders::new( Credential::None, AuthScheme::Bearer, Vec::new(), @@ -156,7 +156,7 @@ pub struct GooglePseProvider { base_url: String, credential: Credential, cx: Option, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/groq/mod.rs b/aimux-providers/src/groq/mod.rs index 207101a10..8291c52d0 100644 --- a/aimux-providers/src/groq/mod.rs +++ b/aimux-providers/src/groq/mod.rs @@ -42,9 +42,9 @@ use aimux_core::language_model::LanguageModel; use aimux_core::model_catalogue::RuntimeModel; use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_core::transcription_model::TranscriptionModel; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; const DEFAULT_BASE_URL: &str = "https://api.groq.com/openai/v1"; const API_KEY_ENV_VAR: &str = "GROQ_API_KEY"; @@ -95,19 +95,16 @@ pub fn create_groq(settings: GroqProviderSettings) -> Result &'static GroqProvider { pub struct GroqProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/huggingface.rs b/aimux-providers/src/huggingface.rs index 8343e4ccc..e65ef8194 100644 --- a/aimux-providers/src/huggingface.rs +++ b/aimux-providers/src/huggingface.rs @@ -29,10 +29,10 @@ 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_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use crate::openai::config::OpenAIModelConfig; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; const DEFAULT_BASE_URL: &str = "https://router.huggingface.co/v1"; const API_KEY_ENV_VAR: &str = "HUGGINGFACE_API_KEY"; @@ -100,7 +100,7 @@ pub fn create_huggingface( Ok(HuggingFaceProvider { name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), base_url, - headers: provider_headers( + headers: ProviderHeaders::bearer( Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Hugging Face"), Vec::new(), settings.headers, @@ -125,7 +125,7 @@ pub fn huggingface() -> &'static HuggingFaceProvider { pub struct HuggingFaceProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } @@ -143,7 +143,7 @@ impl HuggingFaceProvider { OpenAIModelConfig::fixed( format!("{}.{method}", self.name), self.base_url.clone(), - self.headers.clone(), + self.headers.clone().into_headers_fn(), self.fetch.clone(), ) } diff --git a/aimux-providers/src/hume.rs b/aimux-providers/src/hume.rs index 4292a5f20..b2d804ca7 100644 --- a/aimux-providers/src/hume.rs +++ b/aimux-providers/src/hume.rs @@ -30,9 +30,9 @@ use aimux_core::shared::{SharedProviderOptions, Warning}; use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders}; /// Hume errors carry a top-level `message` with a `code` such as `"E0101"` /// (https://dev.hume.ai/docs/resources/errors); the nested `{error: @@ -133,7 +133,7 @@ pub fn create_hume(settings: HumeProviderSettings) -> Result &'static HumeProvider { pub struct HumeProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/jina_ai.rs b/aimux-providers/src/jina_ai.rs index 784f75fa8..c4e9552c8 100644 --- a/aimux-providers/src/jina_ai.rs +++ b/aimux-providers/src/jina_ai.rs @@ -25,9 +25,9 @@ use aimux_core::reranking_model::{ }; use aimux_core::shared::SharedProviderOptions; use aimux_core::types::Warning; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; /// Jina AI error response structure: `{ "detail": "...", "code": "..." }`. /// @@ -113,7 +113,7 @@ pub fn create_jina_ai(settings: JinaAiProviderSettings) -> Result &'static JinaAiProvider { pub struct JinaAiProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/klingai.rs b/aimux-providers/src/klingai.rs index 624ac12c3..fa16cfbc2 100644 --- a/aimux-providers/src/klingai.rs +++ b/aimux-providers/src/klingai.rs @@ -25,9 +25,9 @@ use aimux_core::video_model::{ VideoCallOptions, VideoData, VideoModel, VideoOperationStart, VideoOperationStatus, VideoResponse, VideoResult, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; fn klingai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -112,7 +112,7 @@ pub fn create_klingai(settings: KlingAIProviderSettings) -> Result &'static KlingAIProvider { pub struct KlingAIProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/linkup.rs b/aimux-providers/src/linkup.rs index 102c95b13..e9ed5eb58 100644 --- a/aimux-providers/src/linkup.rs +++ b/aimux-providers/src/linkup.rs @@ -24,9 +24,9 @@ use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; /// Fixed model ID for the Linkup search model. const MODEL_ID: &str = "linkup-search"; @@ -118,7 +118,7 @@ pub fn create_linkup(settings: LinkupProviderSettings) -> Result &'static LinkupProvider { pub struct LinkupProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/lmnt.rs b/aimux-providers/src/lmnt.rs index aa514c22b..6accf247f 100644 --- a/aimux-providers/src/lmnt.rs +++ b/aimux-providers/src/lmnt.rs @@ -28,9 +28,9 @@ use aimux_core::shared::{SharedProviderOptions, Warning}; use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders}; fn lmnt_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -118,7 +118,7 @@ pub fn create_lmnt(settings: LMNTProviderSettings) -> Result &'static LMNTProvider { pub struct LMNTProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/luma.rs b/aimux-providers/src/luma.rs index 063a6822f..8c819d492 100644 --- a/aimux-providers/src/luma.rs +++ b/aimux-providers/src/luma.rs @@ -24,11 +24,11 @@ use aimux_core::image_model::{ }; use aimux_core::shared::Warning; use aimux_provider_utils::HttpRequest; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use crate::shared::{ - Credential, EndpointConfig, POLL_INTERVAL_MILLIS_KEY, PollStep, is_poll_control_key, - poll_interval_ms, poll_until, provider_headers, retry_download, + Credential, EndpointConfig, POLL_INTERVAL_MILLIS_KEY, PollStep, ProviderHeaders, + is_poll_control_key, poll_interval_ms, poll_until, retry_download, }; fn luma_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -118,7 +118,7 @@ pub fn create_luma(settings: LumaProviderSettings) -> Result &'static LumaProvider { pub struct LumaProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/mistral/mod.rs b/aimux-providers/src/mistral/mod.rs index 2e1cefb27..64a64f517 100644 --- a/aimux-providers/src/mistral/mod.rs +++ b/aimux-providers/src/mistral/mod.rs @@ -38,9 +38,9 @@ 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_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders, SupportedUrlsSource}; pub(crate) fn mistral_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -161,19 +161,16 @@ pub fn create_mistral(settings: MistralProviderSettings) -> Result validate_base_url(url)?, None => DEFAULT_BASE_URL.to_string(), }; - let headers = aimux_provider_utils::headers::with_user_agent_suffix_fn( - provider_headers( - Credential::explicit_or_env( - settings.api_key.map(Resolvable::Value), - API_KEY_ENV_VAR, - "Mistral", - ), - Vec::new(), - settings.headers, + let headers = ProviderHeaders::bearer( + Credential::explicit_or_env( + settings.api_key.map(Resolvable::Value), + API_KEY_ENV_VAR, + "Mistral", ), - options::NAMESPACE, - "4.0.54", - ); + Vec::new(), + settings.headers, + ) + .with_user_agent(options::NAMESPACE, "4.0.54"); Ok(MistralProvider { base_url, headers, @@ -197,7 +194,7 @@ pub fn mistral() -> &'static MistralProvider { /// models out of; it holds no HTTP client. pub struct MistralProvider { base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, generate_id: Option, } @@ -218,7 +215,7 @@ impl MistralProvider { MistralModel::from_config( model_id.to_string(), self.model_config("chat") - .with_supported_urls(Arc::new(|_| chat_supported_urls())), + .with_supported_urls(SupportedUrlsSource::Fixed(chat_supported_urls())), self.generate_id, ) } diff --git a/aimux-providers/src/mistral/model.rs b/aimux-providers/src/mistral/model.rs index c2a70dba3..de38ef565 100644 --- a/aimux-providers/src/mistral/model.rs +++ b/aimux-providers/src/mistral/model.rs @@ -163,7 +163,7 @@ impl LanguageModel for MistralModel { } fn supported_urls(&self) -> aimux_core::language_model::SupportedUrls { - (self.config.supported_urls)(&self.model_id) + self.config.supported_urls(&self.model_id) } fn model_id(&self) -> &str { diff --git a/aimux-providers/src/open_responses.rs b/aimux-providers/src/open_responses.rs index f75a5fd4b..d89dd6b04 100644 --- a/aimux-providers/src/open_responses.rs +++ b/aimux-providers/src/open_responses.rs @@ -33,11 +33,9 @@ use aimux_core::types::{ use aimux_core::embedding_model::EmbeddingModel; use aimux_core::image_model::ImageModel; use aimux_core::provider::Provider; -use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, combine_headers, validate_base_url, -}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, EndpointHeaders, ProviderHeaders}; fn open_responses_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -173,19 +171,11 @@ pub fn create_open_responses( Some(key) => Credential::Explicit(key), None => Credential::None, }; - let provider_layer = provider_headers(credential, Vec::new(), None); - let user = settings.headers; - let headers: HeadersFn = Resolvable::from_async_fn(move || { - let provider_layer = provider_layer.clone(); - let user = user.clone(); - async move { - let layer = provider_layer.resolve().await?; - match &user { - Some(user) => Ok(combine_headers(&[&layer, &user.resolve().await?])), - None => Ok(layer), - } - } - }); + let provider = ProviderHeaders::bearer(credential, Vec::new(), None); + let headers = match settings.headers { + Some(user) => EndpointHeaders::WithUserHeaders { provider, user }, + None => EndpointHeaders::Provider(provider), + }; Ok(OpenResponsesProvider { name: settings.name, base_url, @@ -200,7 +190,7 @@ pub fn create_open_responses( pub struct OpenResponsesProvider { name: String, base_url: String, - headers: HeadersFn, + headers: EndpointHeaders, fetch: Option, } diff --git a/aimux-providers/src/parallel_ai.rs b/aimux-providers/src/parallel_ai.rs index 63d4bd10c..7f6b5cf02 100644 --- a/aimux-providers/src/parallel_ai.rs +++ b/aimux-providers/src/parallel_ai.rs @@ -23,9 +23,9 @@ use aimux_core::error::AiMuxError; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders}; /// Fixed model ID for the Parallel AI search model. const MODEL_ID: &str = "parallel-search"; @@ -94,7 +94,7 @@ pub fn create_parallel_ai( Ok(ParallelAiProvider { name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), base_url, - headers: credential_headers( + headers: ProviderHeaders::new( Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Parallel AI"), AuthScheme::Header("x-api-key"), Vec::new(), @@ -120,7 +120,7 @@ pub fn parallel_ai() -> &'static ParallelAiProvider { pub struct ParallelAiProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/prodia.rs b/aimux-providers/src/prodia.rs index c69ceb313..4ef0f33f7 100644 --- a/aimux-providers/src/prodia.rs +++ b/aimux-providers/src/prodia.rs @@ -24,7 +24,7 @@ use aimux_core::image_model::{ use aimux_core::shared::Warning; use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders}; fn prodia_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -152,7 +152,7 @@ impl ProdiaProvider { format!("{}.{method}", self.name), self.base_url.clone(), match method { - "image" => credential_headers( + "image" => ProviderHeaders::new( self.credential.clone(), AuthScheme::Header("X-Prodia-Key"), vec![( @@ -161,7 +161,7 @@ impl ProdiaProvider { )], self.user_headers.clone(), ), - _ => credential_headers( + _ => ProviderHeaders::new( self.credential.clone(), AuthScheme::Header("X-Prodia-Key"), Vec::new(), diff --git a/aimux-providers/src/recraft.rs b/aimux-providers/src/recraft.rs index ffcfbe14e..6b8aea9fb 100644 --- a/aimux-providers/src/recraft.rs +++ b/aimux-providers/src/recraft.rs @@ -26,10 +26,10 @@ use aimux_core::image_model::{ }; use aimux_core::shared::Warning; use aimux_provider_utils::HttpRequest; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use crate::shared::{ - Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, poll_interval_ms, provider_headers, + Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, ProviderHeaders, poll_interval_ms, retry_download, }; @@ -119,7 +119,7 @@ pub fn create_recraft(settings: RecraftProviderSettings) -> Result &'static RecraftProvider { pub struct RecraftProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/replicate.rs b/aimux-providers/src/replicate.rs index 0c0a29bdf..a33ebbf47 100644 --- a/aimux-providers/src/replicate.rs +++ b/aimux-providers/src/replicate.rs @@ -24,7 +24,7 @@ use aimux_provider_utils::HttpRequest; use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use crate::shared::{ - AuthScheme, Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, credential_headers, + AuthScheme, Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, ProviderHeaders, is_poll_control_key, poll_interval_ms, retry_download, }; @@ -153,13 +153,13 @@ impl ReplicateProvider { format!("{}.{method}", self.name), self.base_url.clone(), match method { - "image" => credential_headers( + "image" => ProviderHeaders::new( self.credential.clone(), AuthScheme::Bearer, Vec::new(), self.user_headers.clone(), ), - _ => credential_headers( + _ => ProviderHeaders::new( self.credential.clone(), AuthScheme::Scheme("Token"), Vec::new(), diff --git a/aimux-providers/src/revai.rs b/aimux-providers/src/revai.rs index 5bcbdaeca..3038e2fde 100644 --- a/aimux-providers/src/revai.rs +++ b/aimux-providers/src/revai.rs @@ -27,12 +27,12 @@ use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use aimux_provider_utils::{MultipartForm, media_type_to_extension}; use crate::shared::{ - Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, poll_interval_ms, poll_until, - provider_headers, retry_download, + Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, ProviderHeaders, poll_interval_ms, + poll_until, retry_download, }; fn revai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -128,7 +128,7 @@ pub fn create_revai(settings: RevaiProviderSettings) -> Result &'static RevaiProvider { pub struct RevaiProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/runwayml.rs b/aimux-providers/src/runwayml.rs index 452fed5d1..1fcc70260 100644 --- a/aimux-providers/src/runwayml.rs +++ b/aimux-providers/src/runwayml.rs @@ -26,9 +26,9 @@ use aimux_core::video_model::{ VideoCallOptions, VideoData, VideoFile, VideoFileData, VideoFrameType, VideoModel, VideoOperationStart, VideoOperationStatus, VideoPollConfig, VideoResponse, VideoResult, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; const PROVIDER_NAME: &str = "runwayml"; /// Milliseconds between two status checks of a task. The Core poll loop reads @@ -117,7 +117,7 @@ pub fn create_runwayml(settings: RunwaymlProviderSettings) -> Result &'static RunwaymlProvider { pub struct RunwaymlProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/searxng.rs b/aimux-providers/src/searxng.rs index 9e27cf271..d1ff79f0f 100644 --- a/aimux-providers/src/searxng.rs +++ b/aimux-providers/src/searxng.rs @@ -15,7 +15,7 @@ //! there: the instance URL is resolved for every request. //! [`searxng()`] is the default instance; it reads nothing and cannot fail. -use std::sync::{Arc, OnceLock}; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; @@ -25,11 +25,9 @@ use aimux_core::error::AiMuxError; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, load_setting, validate_base_url, -}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, Endpoint, EndpointConfig, provider_headers}; +use crate::shared::{BaseUrl, Credential, EndpointConfig, EndpointSource, ProviderHeaders}; /// Fixed model ID for the SearXNG search model. const MODEL_ID: &str = "searxng-search"; @@ -91,7 +89,14 @@ pub fn create_searxng(settings: SearxngProviderSettings) -> Result Credential::Explicit(key), None => Credential::None, @@ -99,7 +104,7 @@ pub fn create_searxng(settings: SearxngProviderSettings) -> Result &'static SearxngProvider { /// A SearXNG provider. Search only; it holds no HTTP client. pub struct SearxngProvider { name: String, - base_url: Option, - headers: HeadersFn, + base_url: BaseUrl, + headers: ProviderHeaders, fetch: Option, } impl SearxngProvider { fn model_config(&self, method: &str) -> EndpointConfig { - let base_url = self.base_url.clone(); - let headers = self.headers.clone(); - EndpointConfig::dynamic( + EndpointConfig::new( format!("{}.{method}", self.name), - Arc::new(move || { - let base_url = base_url.clone(); - let headers = headers.clone(); - Box::pin(async move { - let base_url = match base_url { - Some(url) => url, - None => validate_base_url(&load_setting( - None, - BASE_URL_ENV_VAR, - "SearXNG instance URL", - )?)?, - }; - Ok(Endpoint { - base_url, - headers: headers.resolve().await?, - }) - }) - }), + EndpointSource::Http { + base_url: self.base_url.clone(), + headers: self.headers.clone().into(), + }, self.fetch.clone(), ) } diff --git a/aimux-providers/src/serper.rs b/aimux-providers/src/serper.rs index 18aceae9e..901909071 100644 --- a/aimux-providers/src/serper.rs +++ b/aimux-providers/src/serper.rs @@ -19,9 +19,9 @@ use aimux_core::error::AiMuxError; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders}; const MODEL_ID: &str = "serper-search"; @@ -87,7 +87,7 @@ pub fn create_serper(settings: SerperProviderSettings) -> Result &'static SerperProvider { pub struct SerperProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/shared/exchange.rs b/aimux-providers/src/shared/exchange.rs index 5dfebcc60..9af2b983b 100644 --- a/aimux-providers/src/shared/exchange.rs +++ b/aimux-providers/src/shared/exchange.rs @@ -9,17 +9,18 @@ //! the credential or the settings that produced it. use std::collections::HashMap; -use std::sync::Arc; -use futures::future::BoxFuture; +use base64::Engine as _; use aimux_core::AiMuxError; use aimux_core::language_model::SupportedUrls; use aimux_provider_utils::{ ExchangeContext, FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, combine_headers, - normalize_headers, + load_setting, normalize_headers, validate_base_url, }; +use super::{Credential, ProviderHeaders}; + /// What one request is addressed to: the base URL and the provider headers, /// both resolved for this request. pub(crate) struct Endpoint { @@ -27,12 +28,144 @@ pub(crate) struct Endpoint { pub(crate) headers: HeaderMapOpt, } -/// Resolves the endpoint of one request. -pub(crate) type EndpointFn = - Arc BoxFuture<'static, Result> + Send + Sync>; +/// Where the base URL of a request comes from. +#[derive(Clone, Debug)] +pub(crate) enum BaseUrl { + /// Fixed when the provider is created. + Fixed(String), + /// Read from the environment variable on every request and validated; an + /// unset variable fails that request with `AiMuxError::LoadSetting` + /// (a self-hosted instance). + Env { + var: &'static str, + description: &'static str, + }, +} + +impl BaseUrl { + fn resolve(&self) -> Result { + match self { + Self::Fixed(url) => Ok(url.clone()), + Self::Env { var, description } => { + validate_base_url(&load_setting(None, var, description)?) + } + } + } +} + +/// Where the provider headers of a request come from. +#[derive(Clone, Debug)] +pub(crate) enum EndpointHeaders { + /// Credential, fixed and user headers (the common case). + Provider(ProviderHeaders), + /// `Authorization: Basic base64(login:password)`, both parts resolved for + /// every request, then the user's headers. + Basic { + login: Credential, + password: Credential, + user: Option, + }, + /// The provider headers, then the user's headers as the caller gave them + /// (a value or a producer called per request). + WithUserHeaders { + provider: ProviderHeaders, + user: HeadersFn, + }, +} + +impl From for EndpointHeaders { + fn from(headers: ProviderHeaders) -> Self { + Self::Provider(headers) + } +} + +impl EndpointHeaders { + async fn resolve(&self) -> Result { + match self { + Self::Provider(headers) => headers.resolve().await, + Self::Basic { + login, + password, + user, + } => { + let login = login.secret().await?.unwrap_or_default(); + let password = password.secret().await?.unwrap_or_default(); + let encoded = base64::engine::general_purpose::STANDARD + .encode(format!("{login}:{password}").as_bytes()); + let mut layer = HeaderMapOpt::new(); + layer.insert( + "Authorization".to_string(), + Some(format!("Basic {encoded}")), + ); + Ok(match user { + Some(user) => combine_headers(&[&layer, user]), + None => layer, + }) + } + Self::WithUserHeaders { provider, user } => { + let layer = provider.resolve().await?; + Ok(combine_headers(&[&layer, &user.resolve().await?])) + } + } + } +} + +/// Where the endpoint (base URL and provider headers) of a request comes +/// from. Every variant is data; the work happens in [`EndpointConfig::endpoint`] +/// when a request is made. +#[derive(Clone)] +pub(crate) enum EndpointSource { + /// A base URL and headers, each resolved on its own. + Http { + base_url: BaseUrl, + headers: EndpointHeaders, + }, + /// A Gemini-publisher model of Google Vertex AI: the mode, key, token, + /// project and location decide the URL and the headers together. + Vertex(crate::vertex::GeminiEndpoint), + /// An Amazon Bedrock service: explicit URL, endpoint variables or the + /// regional host; API key or SigV4. + Bedrock(crate::bedrock::BedrockEndpoint), + /// Amazon Polly: explicit URL or the regional host; SigV4. + AwsPolly(crate::aws_polly::PollyEndpoint), +} + +impl EndpointSource { + async fn resolve(&self) -> Result { + match self { + Self::Http { base_url, headers } => Ok(Endpoint { + base_url: base_url.resolve()?, + headers: headers.resolve().await?, + }), + Self::Vertex(endpoint) => endpoint.resolve().await, + Self::Bedrock(endpoint) => endpoint.resolve().await, + Self::AwsPolly(endpoint) => endpoint.resolve().await, + } + } +} + +/// The URL patterns a model fetches itself. +#[derive(Clone, Debug, Default)] +pub(crate) enum SupportedUrlsSource { + /// No patterns: the caller downloads every URL. + #[default] + None, + /// The same patterns for every model. + Fixed(SupportedUrls), + /// The Gemini API's patterns, which depend on the base URL and the model + /// id. + Google { base_url: String }, +} -/// The URL patterns a model fetches itself, for a model id. -pub(crate) type SupportedUrlsFn = Arc SupportedUrls + Send + Sync>; +impl SupportedUrlsSource { + fn for_model(&self, model_id: &str) -> SupportedUrls { + match self { + Self::None => SupportedUrls::default(), + Self::Fixed(urls) => urls.clone(), + Self::Google { base_url } => crate::google::supported_urls(base_url, Some(model_id)), + } + } +} /// What a model needs to talk to the API. Built by the provider that owns the /// model. @@ -41,12 +174,12 @@ pub(crate) struct EndpointConfig { /// The identity the model reports from `provider()`. pub(crate) provider: String, /// Base URL and provider headers, resolved on every request. - pub(crate) endpoint: EndpointFn, + pub(crate) endpoint: EndpointSource, /// Transport; `None` uses the process default, resolved per request. pub(crate) fetch: Option, /// URLs the model fetches itself, by model id. Empty: the caller /// downloads them. - pub(crate) supported_urls: SupportedUrlsFn, + pub(crate) supported_urls: SupportedUrlsSource, } impl EndpointConfig { @@ -55,48 +188,54 @@ impl EndpointConfig { pub(crate) fn fixed( provider: String, base_url: String, - headers: HeadersFn, + headers: impl Into, fetch: Option, ) -> Self { - Self { + Self::new( provider, - endpoint: Arc::new(move || { - let base_url = base_url.clone(); - let headers = headers.clone(); - Box::pin(async move { - Ok(Endpoint { - base_url, - headers: headers.resolve().await?, - }) - }) - }), + EndpointSource::Http { + base_url: BaseUrl::Fixed(base_url), + headers: headers.into(), + }, fetch, - supported_urls: Arc::new(|_| SupportedUrls::default()), - } + ) } - /// A config whose base URL is itself resolved on every request (a - /// self-hosted instance named by an environment variable). - pub(crate) fn dynamic( + /// A config whose endpoint comes from `endpoint`. + pub(crate) fn new( provider: String, - endpoint: EndpointFn, + endpoint: EndpointSource, fetch: Option, ) -> Self { Self { provider, endpoint, fetch, - supported_urls: Arc::new(|_| SupportedUrls::default()), + supported_urls: SupportedUrlsSource::None, } } /// The same config with the URLs the model fetches itself. #[must_use] - pub(crate) fn with_supported_urls(mut self, urls: SupportedUrlsFn) -> Self { + pub(crate) fn with_supported_urls(mut self, urls: SupportedUrlsSource) -> Self { self.supported_urls = urls; self } + /// The URL patterns `model_id` fetches itself. + pub(crate) fn supported_urls(&self, model_id: &str) -> SupportedUrls { + self.supported_urls.for_model(model_id) + } + + /// The base URL and provider headers of one request. + /// + /// # Errors + /// + /// Returns the error of the credential, header or setting lookup. + pub(crate) async fn endpoint(&self) -> Result { + self.endpoint.resolve().await + } + /// Resolve the endpoint and layer the per-call headers over the provider /// headers (case-insensitive, later wins, `None` removes). /// @@ -109,7 +248,7 @@ impl EndpointConfig { &self, call_headers: Option<&HashMap>, ) -> Result { - let endpoint = (self.endpoint)().await?; + let endpoint = self.endpoint().await?; let call: HeaderMapOpt = call_headers .map(|headers| { headers diff --git a/aimux-providers/src/shared/mod.rs b/aimux-providers/src/shared/mod.rs index 442f797c2..6f51ec523 100644 --- a/aimux-providers/src/shared/mod.rs +++ b/aimux-providers/src/shared/mod.rs @@ -1,14 +1,18 @@ //! Helpers the native packages and the presets share. //! -//! Only transport plumbing lives here: how the provider-level request headers -//! are produced. Nothing in this module knows a vendor. +//! Only transport plumbing lives here: where the base URL and the +//! provider-level request headers of a request come from, as data. The few +//! vendors whose endpoint is decided by more than a URL and a credential +//! (Vertex, Bedrock, Polly) keep that logic in their own modules. mod discovery; mod exchange; mod poll; pub(crate) use discovery::list_data_models; -pub(crate) use exchange::{Endpoint, EndpointConfig}; +pub(crate) use exchange::{ + BaseUrl, Endpoint, EndpointConfig, EndpointHeaders, EndpointSource, SupportedUrlsSource, +}; pub(crate) use poll::{ POLL_INTERVAL_MILLIS_KEY, POLL_INTERVAL_MS_KEY, PollStep, is_poll_control_key, poll_interval_ms, poll_until, retry_download, @@ -113,6 +117,49 @@ pub(crate) struct ProviderHeaders { } impl ProviderHeaders { + /// Headers with `credential` put on the wire by `scheme`, no user-agent + /// suffix. + pub(crate) fn new( + credential: Credential, + scheme: AuthScheme, + fixed: Vec<(String, String)>, + user: Option, + ) -> Self { + Self { + credential, + scheme, + fixed, + user, + user_agent: None, + } + } + + /// Headers with a bearer credential, no user-agent suffix. + pub(crate) fn bearer( + credential: Credential, + fixed: Vec<(String, String)>, + user: Option, + ) -> Self { + Self::new(credential, AuthScheme::Bearer, fixed, user) + } + + /// The same headers with the `ai-sdk-{package}/{version}` user-agent + /// suffix. + #[must_use] + pub(crate) fn with_user_agent(mut self, package: &'static str, version: &'static str) -> Self { + self.user_agent = Some((package, version)); + self + } + + /// These headers as a [`HeadersFn`], for the configs that still take + /// one; it resolves them on every request. + pub(crate) fn into_headers_fn(self) -> HeadersFn { + Resolvable::from_async_fn(move || { + let headers = self.clone(); + async move { headers.resolve().await } + }) + } + pub(crate) async fn resolve(&self) -> Result { let mut layer = HeaderMapOpt::new(); if let Some(secret) = self.credential.secret().await? { @@ -150,15 +197,5 @@ pub(crate) fn credential_headers( fixed: Vec<(String, String)>, user: Option, ) -> HeadersFn { - let headers = ProviderHeaders { - credential, - scheme, - fixed, - user, - user_agent: None, - }; - Resolvable::from_async_fn(move || { - let headers = headers.clone(); - async move { headers.resolve().await } - }) + ProviderHeaders::new(credential, scheme, fixed, user).into_headers_fn() } diff --git a/aimux-providers/src/stability.rs b/aimux-providers/src/stability.rs index 9cde9610b..6b377b51f 100644 --- a/aimux-providers/src/stability.rs +++ b/aimux-providers/src/stability.rs @@ -25,10 +25,10 @@ use aimux_core::image_model::{ ImageCallOptions, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; use aimux_core::shared::Warning; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use aimux_provider_utils::{HttpBody, MultipartForm}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; /// Stability error response structure: `{ "id": "...", "name": "...", "errors": ["..."] }`. /// @@ -180,7 +180,7 @@ pub fn create_stability( Ok(StabilityProvider { name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), base_url, - headers: provider_headers( + headers: ProviderHeaders::bearer( Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Stability"), vec![("Accept".to_string(), "image/*".to_string())], settings.headers, @@ -205,7 +205,7 @@ pub fn stability() -> &'static StabilityProvider { pub struct StabilityProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/tavily.rs b/aimux-providers/src/tavily.rs index ab5af75ec..d2ae4b6f6 100644 --- a/aimux-providers/src/tavily.rs +++ b/aimux-providers/src/tavily.rs @@ -18,9 +18,9 @@ use aimux_core::error::AiMuxError; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; const MODEL_ID: &str = "tavily-search"; @@ -101,7 +101,7 @@ pub fn create_tavily(settings: TavilyProviderSettings) -> Result &'static TavilyProvider { pub struct TavilyProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/tinyfish.rs b/aimux-providers/src/tinyfish.rs index b414b3f53..0de6cacb7 100644 --- a/aimux-providers/src/tinyfish.rs +++ b/aimux-providers/src/tinyfish.rs @@ -22,9 +22,9 @@ use aimux_core::error::AiMuxError; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders}; /// Fixed model id for the TinyFish search model. const MODEL_ID: &str = "tinyfish-search"; @@ -96,7 +96,7 @@ pub fn create_tinyfish(settings: TinyfishProviderSettings) -> Result &'static TinyfishProvider { pub struct TinyfishProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/vertex/gemini_transcription.rs b/aimux-providers/src/vertex/gemini_transcription.rs index 0abb78a3d..764d3c1d8 100644 --- a/aimux-providers/src/vertex/gemini_transcription.rs +++ b/aimux-providers/src/vertex/gemini_transcription.rs @@ -4,7 +4,7 @@ use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; -use super::ProjectLocationFn; +use super::ProjectLocationSource; use crate::google::options::{GOOGLE, Namespace}; use crate::google::transcription::parse_offset_seconds; use crate::shared::EndpointConfig; @@ -57,7 +57,7 @@ struct Response { pub struct VertexGeminiTranscriptionModel { model_id: String, config: EndpointConfig, - project_location: ProjectLocationFn, + project_location: ProjectLocationSource, #[cfg(feature = "realtime")] web_socket: Option>, } @@ -66,7 +66,7 @@ impl VertexGeminiTranscriptionModel { pub(crate) fn from_config( model_id: String, config: EndpointConfig, - project_location: ProjectLocationFn, + project_location: ProjectLocationSource, #[cfg(feature = "realtime")] web_socket: Option< std::sync::Arc, >, @@ -100,7 +100,7 @@ impl TranscriptionModel for VertexGeminiTranscriptionModel { self.model_id ))); } - (self.project_location)().await?; + self.project_location.resolve().await?; let timestamp = chrono::Utc::now().to_rfc3339(); let config = crate::google::transcription::TranscriptionOptions::parse( Namespace::Vertex.read(options.provider_options.as_ref()), @@ -223,7 +223,7 @@ impl VertexGeminiTranscriptionModel { use aimux_provider_utils::ws::{WebSocketRequest, ws_connect}; crate::google::transcription::validate_live_input_audio_format(&options)?; - let target = (self.project_location)().await?; + let target = self.project_location.resolve().await?; let config = crate::google::transcription::TranscriptionOptions::parse( Namespace::Vertex.read(options.provider_options.as_ref()), )? diff --git a/aimux-providers/src/vertex/mod.rs b/aimux-providers/src/vertex/mod.rs index ebaa12cf0..3f1cc1fe2 100644 --- a/aimux-providers/src/vertex/mod.rs +++ b/aimux-providers/src/vertex/mod.rs @@ -17,12 +17,13 @@ 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, combine_headers, load_setting, + FetchFunction, HeaderMapOpt, Resolvable, combine_headers, load_setting, normalize_headers, + with_user_agent_suffix, }; use std::sync::OnceLock; use crate::google::GoogleSpeechModel; -use crate::shared::{Endpoint, EndpointConfig}; +use crate::shared::{Endpoint, EndpointConfig, EndpointSource, SupportedUrlsSource}; mod anthropic_model; mod anthropic_provider; @@ -140,10 +141,62 @@ pub(crate) struct ProjectLocation { pub(crate) location: String, } -/// Resolves the project and location of a request (Speech-to-Text and the -/// operations that name them), failing for Express mode. -pub(crate) type ProjectLocationFn = - Arc BoxFuture<'static, Result> + Send + Sync>; +/// Where the project and location of a request come from (Speech-to-Text and +/// the operations that name them); Express mode fails the request. +#[derive(Clone)] +pub(crate) struct ProjectLocationSource { + resolver: Arc, + /// The model kind named in the Express-mode error. + what: &'static str, +} + +impl ProjectLocationSource { + /// The project and location of one request. + pub(crate) async fn resolve(&self) -> Result { + if self.resolver.express_key().await?.is_some() { + return Err(AiMuxError::InvalidArgument(format!( + "Google Vertex {} models do not support Express Mode API keys. \ + Use standard Google Cloud credentials instead.", + self.what + ))); + } + Ok(ProjectLocation { + location: self.resolver.location()?, + project: self.resolver.project()?, + }) + } +} + +/// The endpoint of a Gemini-publisher model. With `tuned`, the request fails +/// in Express mode (tuned models are served from a deployed endpoint, which +/// an API key cannot reach) and the URL is the project/location root. +#[derive(Clone)] +pub(crate) struct GeminiEndpoint { + resolver: Arc, + tuned: bool, +} + +impl GeminiEndpoint { + /// The base URL and provider headers of one request. + pub(crate) async fn resolve(&self) -> Result { + if self.tuned && self.resolver.express_key().await?.is_some() { + return Err(AiMuxError::InvalidArgument( + "Google Vertex tuned models do not support Express Mode API keys. \ + Use standard Google Cloud credentials instead." + .to_string(), + )); + } + let mut endpoint = self.resolver.endpoint(Publisher::Google).await?; + if self.tuned && self.resolver.base_url.is_none() { + endpoint.base_url = endpoint + .base_url + .strip_suffix("/publishers/google") + .unwrap_or(&endpoint.base_url) + .to_string(); + } + Ok(endpoint) + } +} /// The settings that decide, per request, which mode applies, where requests /// go and what authenticates them. @@ -247,13 +300,12 @@ impl Resolver { }; let base_url = self.base_url(express_key.is_some(), publisher)?; let headers = self.headers(express_key.as_deref()).await?; - let headers = aimux_provider_utils::headers::with_user_agent_suffix_fn( - Resolvable::Value(headers), - "google-vertex", - "5.0.98", - ) - .resolve() - .await?; + let mut headers = normalize_headers(headers).into_iter().collect(); + with_user_agent_suffix(&mut headers, "ai-sdk-google-vertex/5.0.98"); + let headers = headers + .into_iter() + .map(|(name, value)| (name, Some(value))) + .collect(); Ok(Endpoint { base_url, headers }) } } @@ -339,53 +391,23 @@ impl VertexProvider { /// With `tuned`, the request fails in Express mode (tuned models are /// served from a deployed endpoint, which an API key cannot reach). fn model_config(&self, provider: &str, tuned: bool) -> EndpointConfig { - let resolver = self.resolver.clone(); - EndpointConfig { - provider: provider.to_string(), - endpoint: Arc::new(move || { - let resolver = resolver.clone(); - Box::pin(async move { - if tuned && resolver.express_key().await?.is_some() { - return Err(AiMuxError::InvalidArgument( - "Google Vertex tuned models do not support Express Mode API keys. \ - Use standard Google Cloud credentials instead." - .to_string(), - )); - } - let mut endpoint = resolver.endpoint(Publisher::Google).await?; - if tuned && resolver.base_url.is_none() { - endpoint.base_url = endpoint - .base_url - .strip_suffix("/publishers/google") - .unwrap_or(&endpoint.base_url) - .to_string(); - } - Ok(endpoint) - }) + EndpointConfig::new( + provider.to_string(), + EndpointSource::Vertex(GeminiEndpoint { + resolver: self.resolver.clone(), + tuned, }), - fetch: self.fetch.clone(), - supported_urls: Arc::new(|_| supported_urls()), - } + self.fetch.clone(), + ) + .with_supported_urls(SupportedUrlsSource::Fixed(supported_urls())) } /// The project and location of a request, for the Speech-to-Text models. - fn project_location(&self, what: &'static str) -> ProjectLocationFn { - let resolver = self.resolver.clone(); - Arc::new(move || { - let resolver = resolver.clone(); - Box::pin(async move { - if resolver.express_key().await?.is_some() { - return Err(AiMuxError::InvalidArgument(format!( - "Google Vertex {what} models do not support Express Mode API keys. \ - Use standard Google Cloud credentials instead." - ))); - } - Ok(ProjectLocation { - location: resolver.location()?, - project: resolver.project()?, - }) - }) - }) + fn project_location(&self, what: &'static str) -> ProjectLocationSource { + ProjectLocationSource { + resolver: self.resolver.clone(), + what, + } } /// A Gemini model (e.g. `"gemini-2.0-flash"`, or a tuned diff --git a/aimux-providers/src/vertex/model.rs b/aimux-providers/src/vertex/model.rs index 3c3e40315..50d9da4b8 100644 --- a/aimux-providers/src/vertex/model.rs +++ b/aimux-providers/src/vertex/model.rs @@ -81,7 +81,7 @@ impl LanguageModel for VertexModel { } fn supported_urls(&self) -> SupportedUrls { - (self.config.supported_urls)(&self.model_id) + self.config.supported_urls(&self.model_id) } async fn do_generate(&self, options: &CallOptions) -> Result { diff --git a/aimux-providers/src/vertex/transcription.rs b/aimux-providers/src/vertex/transcription.rs index 41f10eae5..3687e4e36 100644 --- a/aimux-providers/src/vertex/transcription.rs +++ b/aimux-providers/src/vertex/transcription.rs @@ -23,7 +23,7 @@ use aimux_core::transcription_model::{ TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use super::ProjectLocationFn; +use super::ProjectLocationSource; // ── Response schema ───────────────────────────────────────────────────────── @@ -109,14 +109,14 @@ fn audio_input_to_base64(audio: &AudioInput) -> Result { pub struct VertexTranscriptionModel { model_id: String, config: EndpointConfig, - project_location: ProjectLocationFn, + project_location: ProjectLocationSource, } impl VertexTranscriptionModel { pub(crate) fn from_config( model_id: String, config: EndpointConfig, - project_location: ProjectLocationFn, + project_location: ProjectLocationSource, ) -> Self { Self { model_id, @@ -160,7 +160,7 @@ impl TranscriptionModel for VertexTranscriptionModel { let warnings: Vec = Vec::new(); // Parse provider options (`googleVertex`, then `google`). - let target = (self.project_location)().await?; + let target = self.project_location.resolve().await?; let mut region = target.location.clone(); let mut language_codes: Vec = vec!["auto".to_string()]; let mut enable_word_time_offsets = true; diff --git a/aimux-providers/src/voyage/mod.rs b/aimux-providers/src/voyage/mod.rs index 699dc047f..c2905f5c4 100644 --- a/aimux-providers/src/voyage/mod.rs +++ b/aimux-providers/src/voyage/mod.rs @@ -25,9 +25,9 @@ use aimux_core::image_model::ImageModel; use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; use aimux_core::reranking_model::RerankingModel; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; pub(crate) fn voyage_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -106,7 +106,7 @@ pub fn create_voyage(settings: VoyageProviderSettings) -> Result &'static VoyageProvider { pub struct VoyageProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/xai/mod.rs b/aimux-providers/src/xai/mod.rs index 49c42d490..9702dd71e 100644 --- a/aimux-providers/src/xai/mod.rs +++ b/aimux-providers/src/xai/mod.rs @@ -43,9 +43,9 @@ 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, HeadersFn, Resolvable}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable}; -use crate::shared::{Credential, EndpointConfig, provider_headers}; +use crate::shared::{Credential, EndpointConfig, ProviderHeaders}; const DEFAULT_BASE_URL: &str = "https://api.x.ai/v1"; const API_KEY_ENV_VAR: &str = "XAI_API_KEY"; @@ -323,19 +323,16 @@ pub fn create_xai(settings: XAIProviderSettings) -> Result &'static XAIProvider { /// of; it holds no HTTP client. pub struct XAIProvider { base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, #[cfg(feature = "realtime")] web_socket: Option>, diff --git a/aimux-providers/src/you_com.rs b/aimux-providers/src/you_com.rs index e69db571d..ddf4594c8 100644 --- a/aimux-providers/src/you_com.rs +++ b/aimux-providers/src/you_com.rs @@ -23,9 +23,9 @@ use aimux_core::error::AiMuxError; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, ProviderHeaders}; /// Fixed model id for the You.com search model. const MODEL_ID: &str = "youcom-search"; @@ -97,7 +97,7 @@ pub fn create_you_com(settings: YouComProviderSettings) -> Result &'static YouComProvider { pub struct YouComProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } From 78fa018dfd9a3ba189e1e2ed6846d7b227607348 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 12:26:25 +0800 Subject: [PATCH 06/16] refactor(providers)!: OpenAI and Anthropic configs hold ProviderHeaders data OpenAIModelConfig and AnthropicModelConfig take the ProviderHeaders data struct instead of a HeadersFn closure; the credential_headers / provider_headers / into_headers_fn closure builders are gone. Vertex Anthropic resolves its per-request endpoint through AnthropicEndpoint::Vertex holding the Vertex resolver, so the ResolveEndpoint trait object is gone. Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-providers/src/anthropic/config.rs | 28 ++++++------- aimux-providers/src/anthropic/mod.rs | 24 +++++------ aimux-providers/src/anthropic_aws/mod.rs | 5 ++- aimux-providers/src/anthropic_aws/model.rs | 13 +++--- aimux-providers/src/azure/mod.rs | 32 ++++++-------- aimux-providers/src/codex.rs | 9 ++-- aimux-providers/src/huggingface.rs | 2 +- aimux-providers/src/openai/config.rs | 9 ++-- aimux-providers/src/openai/mod.rs | 42 +++++++++---------- aimux-providers/src/shared/mod.rs | 32 +------------- aimux-providers/src/vertex/anthropic_model.rs | 25 ++++------- aimux-providers/src/vertex/mod.rs | 6 +-- 12 files changed, 86 insertions(+), 141 deletions(-) diff --git a/aimux-providers/src/anthropic/config.rs b/aimux-providers/src/anthropic/config.rs index d42f1305d..8557e1c73 100644 --- a/aimux-providers/src/anthropic/config.rs +++ b/aimux-providers/src/anthropic/config.rs @@ -12,14 +12,14 @@ use std::collections::{BTreeSet, HashMap}; use std::sync::Arc; -use futures::future::BoxFuture; use serde_json::Value; +use crate::shared::{Credential, ProviderHeaders}; use aimux_core::AiMuxError; use aimux_core::language_model::SupportedUrls; use aimux_provider_utils::{ - ExchangeContext, FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, Resolvable, - ResponseHandler, combine_headers, normalize_headers, + ExchangeContext, FetchFunction, HeaderMapOpt, HttpRequest, ResponseHandler, combine_headers, + normalize_headers, }; use super::convert::RequestProfile; @@ -42,17 +42,11 @@ pub(crate) enum AnthropicEndpoint { /// The config's own `base_url` and `headers`. #[default] Fixed, - /// Resolved for every request, for hosts whose endpoint or credentials (a - /// Vertex project and location, say) can only be known when a call is + /// The Vertex endpoint of Anthropic's models, resolved for every request: + /// the project, location and credentials are only known when a call is /// made. The config's own `base_url` and `headers` then only describe the /// model's static identity. - PerRequest(Arc), -} - -/// A host's per-request endpoint lookup (the Vertex resolver, for instance). -pub(crate) trait ResolveEndpoint: Send + Sync { - /// The base URL and provider headers of one request. - fn endpoint(&self) -> BoxFuture<'_, Result>; + Vertex(Arc), } /// Builds the host's error handler for non-2xx responses. @@ -97,7 +91,7 @@ pub(crate) struct AnthropicModelConfig { pub(crate) provider: String, /// Provider headers (credential, `anthropic-version`, user headers), /// resolved on every request. - pub(crate) headers: HeadersFn, + pub(crate) headers: ProviderHeaders, /// Transport; `None` uses the process default, resolved per request. pub(crate) fetch: Option, /// URLs the model fetches itself. Empty: the caller downloads them. @@ -125,11 +119,13 @@ impl AnthropicModelConfig { pub(crate) async fn resolved(&self) -> Result { match &self.endpoint { AnthropicEndpoint::Fixed => Ok(self.clone()), - AnthropicEndpoint::PerRequest(resolver) => { - let Endpoint { base_url, headers } = resolver.endpoint().await?; + AnthropicEndpoint::Vertex(resolver) => { + let Endpoint { base_url, headers } = resolver + .endpoint(crate::vertex::Publisher::Anthropic) + .await?; Ok(Self { base_url, - headers: Resolvable::Value(headers), + headers: ProviderHeaders::bearer(Credential::None, Vec::new(), Some(headers)), endpoint: AnthropicEndpoint::Fixed, ..self.clone() }) diff --git a/aimux-providers/src/anthropic/mod.rs b/aimux-providers/src/anthropic/mod.rs index d69a125ee..344ddd395 100644 --- a/aimux-providers/src/anthropic/mod.rs +++ b/aimux-providers/src/anthropic/mod.rs @@ -35,6 +35,7 @@ use std::sync::{Arc, OnceLock}; use futures::future::BoxFuture; use serde_json::Value; +use crate::shared::ProviderHeaders; use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; use aimux_core::files_model::Files; @@ -43,10 +44,10 @@ use aimux_core::language_model::{LanguageModel, SupportedUrls}; use aimux_core::model_catalogue::RuntimeModel; use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, load_optional_setting, validate_base_url, + FetchFunction, HeaderMapOpt, Resolvable, load_optional_setting, validate_base_url, }; -use crate::shared::{AuthScheme, Credential, credential_headers}; +use crate::shared::{AuthScheme, Credential}; use config::{AnthropicEndpoint, AnthropicModelConfig, AnthropicModelHooks}; @@ -194,16 +195,13 @@ pub fn create_anthropic( Ok(AnthropicProvider { name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), base_url, - headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( - credential_headers( - credential, - scheme, - vec![("anthropic-version".to_string(), API_VERSION.to_string())], - settings.headers, - ), - options::CANONICAL, - "4.0.68", - ), + headers: ProviderHeaders::new( + credential, + scheme, + vec![("anthropic-version".to_string(), API_VERSION.to_string())], + settings.headers, + ) + .with_user_agent(options::CANONICAL, "4.0.68"), fetch: settings.fetch, supported_urls: supported_urls(), generate_id: settings.generate_id, @@ -227,7 +225,7 @@ pub fn anthropic() -> &'static AnthropicProvider { pub struct AnthropicProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, supported_urls: SupportedUrls, generate_id: Option, diff --git a/aimux-providers/src/anthropic_aws/mod.rs b/aimux-providers/src/anthropic_aws/mod.rs index 9c88f86c1..b15d1c812 100644 --- a/aimux-providers/src/anthropic_aws/mod.rs +++ b/aimux-providers/src/anthropic_aws/mod.rs @@ -18,6 +18,7 @@ use std::sync::{Arc, OnceLock}; use futures::future::BoxFuture; +use crate::shared::ProviderHeaders; use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; use aimux_core::image_model::ImageModel; @@ -25,7 +26,7 @@ use aimux_core::language_model::{LanguageModel, SupportedUrls}; use aimux_core::model_catalogue::RuntimeModel; use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_provider_utils::{ - AwsCredentials, FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url, + AwsCredentials, FetchFunction, HeaderMapOpt, Resolvable, validate_base_url, }; use crate::anthropic::AnthropicMessagesModel; @@ -187,7 +188,7 @@ pub struct AnthropicAwsProvider { name: String, provider_options_name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/anthropic_aws/model.rs b/aimux-providers/src/anthropic_aws/model.rs index 524870370..ea20152dd 100644 --- a/aimux-providers/src/anthropic_aws/model.rs +++ b/aimux-providers/src/anthropic_aws/model.rs @@ -9,12 +9,11 @@ use std::sync::Arc; -use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, SigV4Fetch, default_fetch, -}; +use crate::shared::ProviderHeaders; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, SigV4Fetch, default_fetch}; use crate::anthropic::AnthropicMessagesModel; -use crate::shared::{AuthScheme, Credential, credential_headers}; +use crate::shared::{AuthScheme, Credential}; use super::AnthropicAwsAuth; @@ -40,7 +39,7 @@ pub(super) fn authenticated( workspace_id: Option, user_headers: Option, fetch: Option, -) -> (HeadersFn, Option) { +) -> (ProviderHeaders, Option) { let mut fixed = vec![("anthropic-version".to_string(), API_VERSION.to_string())]; if let Some(workspace_id) = workspace_id { fixed.push(("anthropic-workspace-id".to_string(), workspace_id)); @@ -51,7 +50,7 @@ pub(super) fn authenticated( let signing: FetchFunction = Arc::new(SigV4Fetch::new(inner, credentials, SIGV4_SERVICE)); ( - credential_headers(Credential::None, AuthScheme::Bearer, fixed, user_headers), + ProviderHeaders::new(Credential::None, AuthScheme::Bearer, fixed, user_headers), Some(signing), ) } @@ -61,7 +60,7 @@ pub(super) fn authenticated( _ => None, }; ( - credential_headers( + ProviderHeaders::new( Credential::explicit_or_env(key, API_KEY_ENV_VAR, "Anthropic AWS"), AuthScheme::Header("x-api-key"), fixed, diff --git a/aimux-providers/src/azure/mod.rs b/aimux-providers/src/azure/mod.rs index a5d38eb92..42ba2c4e6 100644 --- a/aimux-providers/src/azure/mod.rs +++ b/aimux-providers/src/azure/mod.rs @@ -30,6 +30,7 @@ use std::sync::{Arc, OnceLock}; use futures::future::BoxFuture; +use crate::shared::ProviderHeaders; use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; use aimux_core::image_model::ImageModel; @@ -39,9 +40,8 @@ use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_core::speech_model::SpeechModel; use aimux_core::transcription_model::TranscriptionModel; use aimux_provider_utils::{ - Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, HeadersFn, - HttpRequest, Resolvable, default_fetch, load_setting, validate_base_url, - without_trailing_slash, + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, HttpRequest, + Resolvable, default_fetch, load_setting, validate_base_url, without_trailing_slash, }; use crate::openai::config::{OpenAIModelConfig, OpenAIUrl}; @@ -50,7 +50,7 @@ use crate::openai::{ OpenAIEmbeddingModel, OpenAIImageModel, OpenAIModel, OpenAIResponsesModel, OpenAISpeechModel, OpenAITranscriptionModel, }; -use crate::shared::{AuthScheme, Credential, credential_headers, is_valid_hostname_part}; +use crate::shared::{AuthScheme, Credential, is_valid_hostname_part}; /// The chat-completions model of the Azure package (`provider.chat(id)`): the /// OpenAI one, configured for Azure. @@ -210,8 +210,8 @@ pub fn create_azure( }) .or(settings.fetch); let speech_headers = match &settings.token_provider { - Some(_) => Resolvable::Value(settings.headers.clone().unwrap_or_default()), - None => credential_headers( + Some(_) => ProviderHeaders::bearer(Credential::None, Vec::new(), settings.headers.clone()), + None => ProviderHeaders::new( Credential::explicit_or_env(settings.api_key.clone(), API_KEY_ENV_VAR, "Azure Speech"), AuthScheme::Header("Ocp-Apim-Subscription-Key"), Vec::new(), @@ -219,8 +219,8 @@ pub fn create_azure( ), }; let headers = match settings.token_provider { - Some(_) => Resolvable::Value(settings.headers.unwrap_or_default()), - None => credential_headers( + Some(_) => ProviderHeaders::bearer(Credential::None, Vec::new(), settings.headers), + None => ProviderHeaders::new( Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Azure OpenAI"), AuthScheme::Header("api-key"), Vec::new(), @@ -231,20 +231,12 @@ pub fn create_azure( resource_name: settings.resource_name, base_url, speech_base_url, - speech_headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( - speech_headers, - options::NAMESPACE, - "4.0.84", - ), + speech_headers: speech_headers.with_user_agent(options::NAMESPACE, "4.0.84"), info, api_version: settings.api_version, use_deployment_based_urls: settings.use_deployment_based_urls, fetch, - headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( - headers, - options::NAMESPACE, - "4.0.84", - ), + headers: headers.with_user_agent(options::NAMESPACE, "4.0.84"), }) } @@ -298,11 +290,11 @@ pub struct AzureOpenAIProvider { resource_name: Option, base_url: Option, speech_base_url: Option, - speech_headers: HeadersFn, + speech_headers: ProviderHeaders, info: BaseUrlInfo, api_version: Option, use_deployment_based_urls: bool, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/codex.rs b/aimux-providers/src/codex.rs index 02a6b89ac..062fbaedf 100644 --- a/aimux-providers/src/codex.rs +++ b/aimux-providers/src/codex.rs @@ -48,8 +48,9 @@ use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_core::result::{GenerateResult, StreamResult}; use aimux_core::types::Warning; +use crate::shared::ProviderHeaders; use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, Resolvable, validate_base_url, + FetchFunction, HeaderMapOpt, HttpRequest, Resolvable, validate_base_url, }; use crate::openai::config::OpenAIModelConfig; @@ -59,7 +60,7 @@ use crate::openai::responses::responses_convert::{ use crate::openai::responses::{ OpenAIResponsesModel, ResponsesNamespace, build_responses_request_body, }; -use crate::shared::{AuthScheme, Credential, credential_headers}; +use crate::shared::{AuthScheme, Credential}; /// Default API-key base URL (official OpenAI Responses endpoint). pub const CODEX_API_BASE_URL: &str = "https://api.openai.com/v1"; @@ -200,7 +201,7 @@ pub fn create_codex(settings: CodexProviderSettings) -> Result, } diff --git a/aimux-providers/src/huggingface.rs b/aimux-providers/src/huggingface.rs index e65ef8194..005e5b0ba 100644 --- a/aimux-providers/src/huggingface.rs +++ b/aimux-providers/src/huggingface.rs @@ -143,7 +143,7 @@ impl HuggingFaceProvider { OpenAIModelConfig::fixed( format!("{}.{method}", self.name), self.base_url.clone(), - self.headers.clone().into_headers_fn(), + self.headers.clone(), self.fetch.clone(), ) } diff --git a/aimux-providers/src/openai/config.rs b/aimux-providers/src/openai/config.rs index 1346e2fed..527096c01 100644 --- a/aimux-providers/src/openai/config.rs +++ b/aimux-providers/src/openai/config.rs @@ -10,11 +10,12 @@ use std::collections::HashMap; +use crate::shared::ProviderHeaders; use aimux_core::AiMuxError; use aimux_core::language_model::SupportedUrls; use aimux_provider_utils::{ - ExchangeContext, FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, Resolvable, - combine_headers, normalize_headers, + ExchangeContext, FetchFunction, HeaderMapOpt, HttpRequest, Resolvable, combine_headers, + normalize_headers, }; use super::responses::ResponsesProfile; @@ -66,7 +67,7 @@ pub(crate) struct OpenAIModelConfig { pub(crate) url: OpenAIUrl, /// Provider headers (credential, organization, project, user headers), /// resolved on every request. - pub(crate) headers: HeadersFn, + pub(crate) headers: ProviderHeaders, pub(crate) token_provider: Option>, /// Transport; `None` uses the process default, resolved per request. pub(crate) fetch: Option, @@ -83,7 +84,7 @@ impl OpenAIModelConfig { pub(crate) fn fixed( provider: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, ) -> Self { Self { diff --git a/aimux-providers/src/openai/mod.rs b/aimux-providers/src/openai/mod.rs index 8e8aebd21..f69be7d3a 100644 --- a/aimux-providers/src/openai/mod.rs +++ b/aimux-providers/src/openai/mod.rs @@ -40,6 +40,7 @@ use std::sync::{Arc, OnceLock}; use futures::future::BoxFuture; use serde_json::Value; +use crate::shared::ProviderHeaders; use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; use aimux_core::evaluation_model::EvaluationModel; @@ -51,11 +52,11 @@ use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_core::speech_model::SpeechModel; use aimux_core::transcription_model::TranscriptionModel; use aimux_provider_utils::{ - EvaluationLanguageModel, FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, - load_optional_setting, validate_base_url, + EvaluationLanguageModel, FetchFunction, HeaderMapOpt, Resolvable, load_optional_setting, + validate_base_url, }; -use crate::shared::{Credential, provider_headers}; +use crate::shared::Credential; use config::OpenAIModelConfig; @@ -243,25 +244,22 @@ pub fn create_openai(settings: OpenAIProviderSettings) -> Result &'static OpenAIProvider { pub struct OpenAIProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, #[cfg(feature = "realtime")] web_socket: Option>, diff --git a/aimux-providers/src/shared/mod.rs b/aimux-providers/src/shared/mod.rs index 6f51ec523..3de9b6dac 100644 --- a/aimux-providers/src/shared/mod.rs +++ b/aimux-providers/src/shared/mod.rs @@ -20,7 +20,7 @@ pub(crate) use poll::{ use aimux_core::AiMuxError; use aimux_provider_utils::{ - HeaderMapOpt, HeadersFn, Resolvable, combine_headers, load_api_key, normalize_headers, + HeaderMapOpt, Resolvable, combine_headers, load_api_key, normalize_headers, with_user_agent_suffix, }; @@ -91,16 +91,6 @@ pub(crate) enum AuthScheme { Scheme(&'static str), } -/// Provider headers with a bearer credential: [`credential_headers`] with -/// [`AuthScheme::Bearer`]. -pub(crate) fn provider_headers( - credential: Credential, - fixed: Vec<(String, String)>, - user: Option, -) -> HeadersFn { - credential_headers(credential, AuthScheme::Bearer, fixed, user) -} - /// The provider-level request headers, as data. They are produced on every /// request: the credential first, then the fixed headers (organization, /// project, version, ...), then the user's headers, which may override or @@ -151,15 +141,6 @@ impl ProviderHeaders { self } - /// These headers as a [`HeadersFn`], for the configs that still take - /// one; it resolves them on every request. - pub(crate) fn into_headers_fn(self) -> HeadersFn { - Resolvable::from_async_fn(move || { - let headers = self.clone(); - async move { headers.resolve().await } - }) - } - pub(crate) async fn resolve(&self) -> Result { let mut layer = HeaderMapOpt::new(); if let Some(secret) = self.credential.secret().await? { @@ -188,14 +169,3 @@ impl ProviderHeaders { .collect()) } } - -/// [`ProviderHeaders`] without a user-agent suffix, as a [`HeadersFn`] for -/// the packages that still take one. -pub(crate) fn credential_headers( - credential: Credential, - scheme: AuthScheme, - fixed: Vec<(String, String)>, - user: Option, -) -> HeadersFn { - ProviderHeaders::new(credential, scheme, fixed, user).into_headers_fn() -} diff --git a/aimux-providers/src/vertex/anthropic_model.rs b/aimux-providers/src/vertex/anthropic_model.rs index 5800fc45f..0aa99fb6c 100644 --- a/aimux-providers/src/vertex/anthropic_model.rs +++ b/aimux-providers/src/vertex/anthropic_model.rs @@ -20,20 +20,19 @@ use std::sync::Arc; -use futures::future::BoxFuture; use serde_json::{Value, json}; -use aimux_core::AiMuxError; use aimux_core::language_model::SupportedUrls; -use aimux_provider_utils::{FetchFunction, Resolvable}; +use aimux_provider_utils::FetchFunction; -use super::{Publisher, Resolver}; +use crate::shared::{Credential, ProviderHeaders}; + +use super::Resolver; use crate::anthropic::AnthropicMessagesModel; use crate::anthropic::config::{ - AnthropicEndpoint, AnthropicModelConfig, AnthropicModelHooks, MessagesUrl, ResolveEndpoint, + AnthropicEndpoint, AnthropicModelConfig, AnthropicModelHooks, MessagesUrl, }; use crate::anthropic::options::options_name_of; -use crate::shared::Endpoint; /// `anthropic_version` envelope value required by the Vertex AI `rawPredict` / /// `streamRawPredict` endpoints. @@ -59,16 +58,6 @@ fn raw_predict_envelope(body: Value) -> Value { Value::Object(envelope) } -/// The Vertex endpoint of Anthropic's models, resolved for each request: the -/// `.../publishers/anthropic/models` base URL and the provider headers. -struct VertexAnthropicEndpoint(Arc); - -impl ResolveEndpoint for VertexAnthropicEndpoint { - fn endpoint(&self) -> BoxFuture<'_, Result> { - Box::pin(self.0.endpoint(Publisher::Anthropic)) - } -} - /// The model `model_id` on the Vertex endpoint `resolver` resolves for each /// request. pub(super) fn model( @@ -80,7 +69,7 @@ pub(super) fn model( // every request uses the endpoint resolved for it. let config = AnthropicModelConfig { provider: PROVIDER.to_string(), - headers: Resolvable::Value(Default::default()), + headers: ProviderHeaders::bearer(Credential::None, Vec::new(), None), fetch, supported_urls: SupportedUrls::default(), base_url: String::new(), @@ -92,7 +81,7 @@ pub(super) fn model( supports_native_structured_output: false, supports_strict_tools: false, }, - endpoint: AnthropicEndpoint::PerRequest(Arc::new(VertexAnthropicEndpoint(resolver))), + endpoint: AnthropicEndpoint::Vertex(resolver), }; AnthropicMessagesModel::with_config(model_id.to_string(), config) } diff --git a/aimux-providers/src/vertex/mod.rs b/aimux-providers/src/vertex/mod.rs index 3f1cc1fe2..14bbaa84d 100644 --- a/aimux-providers/src/vertex/mod.rs +++ b/aimux-providers/src/vertex/mod.rs @@ -130,7 +130,7 @@ fn location_host(location: &str) -> String { /// Which publisher's models a base URL addresses. #[derive(Clone, Copy)] -enum Publisher { +pub(crate) enum Publisher { Google, Anthropic, } @@ -200,7 +200,7 @@ impl GeminiEndpoint { /// The settings that decide, per request, which mode applies, where requests /// go and what authenticates them. -struct Resolver { +pub(crate) struct Resolver { api_key: Option>, express_mode: bool, location: Option, @@ -293,7 +293,7 @@ impl Resolver { } /// The endpoint of `publisher`'s models for one request. - async fn endpoint(&self, publisher: Publisher) -> Result { + pub(crate) async fn endpoint(&self, publisher: Publisher) -> Result { let express_key = match publisher { Publisher::Google => self.express_key().await?, Publisher::Anthropic => None, From c18edcd3af8749741f0867c48f64388c995824d1 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 12:27:29 +0800 Subject: [PATCH 07/16] feat(provider-utils): TransformStream transformer and pipe_through The Rust form of the WHATWG TransformStream the AI SDK models build doStream with: a Transformer value holds the state the TypeScript closures capture, with start / transform / flush methods writing to a controller (enqueue, error, terminate), and pipe_through drives it over an input stream. Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-provider-utils/src/lib.rs | 2 + aimux-provider-utils/src/transform_stream.rs | 160 +++++++++++++++++++ 2 files changed, 162 insertions(+) create mode 100644 aimux-provider-utils/src/transform_stream.rs diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs index f5b003eae..a13ce4c02 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -30,6 +30,7 @@ pub mod response_handler; pub mod sigv4_fetch; pub mod streaming_tool_call_argument_state; pub mod streaming_tool_call_tracker; +pub mod transform_stream; pub mod url; /// WebSocket client for realtime provider APIs (RFC-0028). Empty unless the /// `ws` feature is enabled. @@ -71,4 +72,5 @@ pub use streaming_tool_call_argument_state::{ pub use streaming_tool_call_tracker::{ StreamingToolCallDelta, StreamingToolCallTracker, TrackerError, TypeValidation, }; +pub use transform_stream::{TransformStreamController, Transformer, pipe_through}; pub use url::{validate_base_url, without_trailing_slash, without_trailing_slash_opt}; diff --git a/aimux-provider-utils/src/transform_stream.rs b/aimux-provider-utils/src/transform_stream.rs new file mode 100644 index 000000000..eca71fa2f --- /dev/null +++ b/aimux-provider-utils/src/transform_stream.rs @@ -0,0 +1,160 @@ +//! The Rust form of the WHATWG `TransformStream` that the AI SDK's models +//! build their `doStream` output with (`response.pipeThrough(new +//! TransformStream({ start, transform, flush }))`). +//! +//! The state the TypeScript closures capture lives in a [`Transformer`] +//! value instead, so a model's stream logic is a plain struct with methods +//! that can be driven and tested without any I/O. + +use futures::{Stream, StreamExt}; + +use aimux_core::error::AiMuxError; + +/// What a transformer writes to: the `TransformStreamDefaultController`. +pub struct TransformStreamController { + queue: Vec>, + closed: bool, +} + +impl TransformStreamController { + fn new() -> Self { + Self { + queue: Vec::new(), + closed: false, + } + } + + /// Emit one value downstream. + pub fn enqueue(&mut self, value: O) { + if !self.closed { + self.queue.push(Ok(value)); + } + } + + /// Emit `error` and end the stream: no further chunk is transformed and + /// `flush` is not called. + pub fn error(&mut self, error: AiMuxError) { + if !self.closed { + self.queue.push(Err(error)); + self.closed = true; + } + } + + /// End the stream without an error: no further chunk is transformed and + /// `flush` is not called. + pub fn terminate(&mut self) { + self.closed = true; + } +} + +/// The `{ start, transform, flush }` object of a `TransformStream`. +pub trait Transformer: Send + 'static { + type Input: Send + 'static; + type Output: Send + 'static; + + /// Called once, before the first chunk. + fn start(&mut self, controller: &mut TransformStreamController) { + let _ = controller; + } + + /// Called for every input chunk. + fn transform( + &mut self, + chunk: Self::Input, + controller: &mut TransformStreamController, + ); + + /// Called once when the input ends, unless the stream was ended earlier + /// with [`TransformStreamController::error`] or + /// [`TransformStreamController::terminate`]. + fn flush(self, controller: &mut TransformStreamController); +} + +/// `input.pipeThrough(new TransformStream(transformer))`. +pub fn pipe_through( + input: S, + mut transformer: T, +) -> impl Stream> + Send +where + S: Stream + Send + 'static, + T: Transformer, +{ + async_stream::stream! { + let mut controller = TransformStreamController::new(); + transformer.start(&mut controller); + for item in std::mem::take(&mut controller.queue) { + yield item; + } + if controller.closed { + return; + } + let mut input = Box::pin(input); + while let Some(chunk) = input.next().await { + transformer.transform(chunk, &mut controller); + for item in std::mem::take(&mut controller.queue) { + yield item; + } + if controller.closed { + return; + } + } + transformer.flush(&mut controller); + for item in std::mem::take(&mut controller.queue) { + yield item; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Doubles each number, stops with an error on a negative one, and + /// emits the running total on flush. + #[derive(Default)] + struct Doubler { + total: i32, + } + + impl Transformer for Doubler { + type Input = i32; + type Output = i32; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(0); + } + + fn transform(&mut self, chunk: i32, controller: &mut TransformStreamController) { + if chunk < 0 { + controller.error(AiMuxError::InvalidResponseData("negative".into())); + return; + } + self.total += chunk; + controller.enqueue(chunk * 2); + } + + fn flush(self, controller: &mut TransformStreamController) { + controller.enqueue(self.total); + } + } + + async fn run(input: Vec) -> Vec> { + pipe_through(futures::stream::iter(input), Doubler::default()) + .map(|item| item.map_err(|error| error.to_string())) + .collect() + .await + } + + #[tokio::test] + async fn start_transform_then_flush_in_order() { + assert_eq!(run(vec![1, 2]).await, vec![Ok(0), Ok(2), Ok(4), Ok(3)]); + } + + #[tokio::test] + async fn error_ends_the_stream_without_flush() { + let out = run(vec![1, -1, 5]).await; + assert_eq!(out.len(), 3); + assert_eq!(out[..2], [Ok(0), Ok(2)]); + assert!(out[2].is_err()); + } +} From d49a1b89b4b1c747749dbd6e5d50fb714a7ff7bd Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 12:39:59 +0800 Subject: [PATCH 08/16] refactor(providers): Groq and Mistral chat streams as Transformer structs GroqChatStream and MistralChatStream hold the doStream state and are driven by pipe_through. DeepSeek keeps its generator: it yields a recoverable Err (then continues) that the controller cannot express yet. Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-providers/src/groq/model.rs | 383 +++++++++++++----------- aimux-providers/src/mistral/model.rs | 431 +++++++++++++++------------ 2 files changed, 453 insertions(+), 361 deletions(-) diff --git a/aimux-providers/src/groq/model.rs b/aimux-providers/src/groq/model.rs index 4271f34c3..01e820208 100644 --- a/aimux-providers/src/groq/model.rs +++ b/aimux-providers/src/groq/model.rs @@ -8,7 +8,6 @@ use aimux_core::tool::RawToolCall; use std::collections::HashMap; use async_trait::async_trait; -use futures::StreamExt; use serde_json::{Map, Value, json}; use aimux_core::error::AiMuxError; @@ -21,7 +20,8 @@ use aimux_core::types::{ FinishReason, FinishReasonUnified, ReasoningEffort, ResponseMetadata, Warning, }; use aimux_provider_utils::{ - StreamingToolCallDelta, StreamingToolCallTracker, TypeValidation, generate_id, + StreamingToolCallDelta, StreamingToolCallTracker, TransformStreamController, Transformer, + TypeValidation, generate_id, pipe_through, }; use crate::shared::EndpointConfig; @@ -31,7 +31,7 @@ use super::error::{GroqErrorData, groq_failed_response_handler}; use super::finish_reason::map_groq_finish_reason; use super::options::{self, GroqLanguageModelChatOptions, parse_groq_options}; use super::prepare_tools::prepare_tools; -use super::types::{GroqChatChunk, GroqChatResponse}; +use super::types::{GroqChatChunk, GroqChatResponse, GroqUsage}; use super::usage::convert_groq_usage; /// A Groq chat language model. @@ -363,209 +363,242 @@ impl LanguageModel for GroqChatLanguageModel { ) .await?; let response_headers = resp.response_headers; - let mut sse_stream = resp.value; - let emit_raw_chunks = options.include_raw_chunks == Some(true); - let stream_error_url = endpoint; - let stream_error_body = body.clone(); - let stream_response_headers = response_headers.clone(); + let stream = pipe_through( + resp.value, + GroqChatStream { + warnings, + emit_raw_chunks: options.include_raw_chunks == Some(true), + stream_error_url: endpoint, + stream_error_body: body.clone(), + stream_response_headers: response_headers.clone(), + tool_calls: StreamingToolCallTracker::new() + .with_generate_id(generate_id) + .with_type_validation(TypeValidation::Required), + tool_parts: Vec::new(), + finish_reason: FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }, + usage: None, + is_first_chunk: true, + is_active_text: false, + is_active_reasoning: false, + }, + ); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(RequestInfo { body: Some(body) }), + response: Some(StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} - let mut tool_calls = StreamingToolCallTracker::new() - .with_generate_id(generate_id) - .with_type_validation(TypeValidation::Required); - let mut tool_parts = Vec::new(); +/// The `TransformStream` of `GroqChatLanguageModel.doStream`. +struct GroqChatStream { + warnings: Vec, + emit_raw_chunks: bool, + stream_error_url: String, + stream_error_body: Value, + stream_response_headers: HashMap, + tool_calls: StreamingToolCallTracker, + tool_parts: Vec, + finish_reason: FinishReason, + usage: Option, + is_first_chunk: bool, + is_active_text: bool, + is_active_reasoning: bool, +} - let mut finish_reason = FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }; - let mut usage = None; - let mut is_first_chunk = true; - let mut is_active_text = false; - let mut is_active_reasoning = false; - - while let Some(event) = sse_stream.next().await { - let parsed = match event { - Ok(parsed) => parsed, - // a chunk that fails to parse is reported as an error part - // (`chunk.success === false`); a transport failure ends the stream. - Err(error) => { - if !error.is_recoverable_stream_error() { - yield Err(error); - return; - } - if emit_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: Value::Null }); - } - finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - }; - yield Ok(StreamPart::Error { error }); - continue; - } - }; +impl GroqChatStream { + fn end_active_reasoning(&mut self, controller: &mut TransformStreamController) { + if self.is_active_reasoning { + controller.enqueue(StreamPart::ReasoningEnd { + id: "reasoning-0".to_string(), + provider_metadata: None, + }); + self.is_active_reasoning = false; + } + } +} - // Emit the raw chunk if requested (before anything else). - if emit_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); - } +impl Transformer for GroqChatStream { + type Input = Result; + type Output = StreamPart; - // handle error chunks: - if parsed.get("error").is_some() { - finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - }; - yield Ok(StreamPart::Error { - error: create_groq_stream_error( - &parsed, - &stream_error_url, - stream_error_body.clone(), - stream_response_headers.clone(), - ), + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + let parsed = match event { + Ok(parsed) => parsed, + // a chunk that fails to parse is reported as an error part + // (`chunk.success === false`); a transport failure ends the stream. + Err(error) => { + if !error.is_recoverable_stream_error() { + controller.error(error); + return; + } + if self.emit_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: Value::Null, }); - continue; } - - // handle failed chunk parsing / validation: - let value: GroqChatChunk = match serde_json::from_value(parsed) { - Ok(value) => value, - Err(error) => { - finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - }; - yield Ok(StreamPart::Error { error: error.into() }); - continue; - } + self.finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, }; + controller.enqueue(StreamPart::Error { error }); + return; + } + }; - if is_first_chunk { - is_first_chunk = false; - let metadata = response_metadata(value.id, value.created, value.model); - yield Ok(StreamPart::ResponseMetadata(metadata)); - } + // Emit the raw chunk if requested (before anything else). + if self.emit_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: parsed.clone(), + }); + } - if let Some(chunk_usage) = value.x_groq.and_then(|x_groq| x_groq.usage) { - usage = Some(chunk_usage); - } + // handle error chunks: + if parsed.get("error").is_some() { + self.finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }; + controller.enqueue(StreamPart::Error { + error: create_groq_stream_error( + &parsed, + &self.stream_error_url, + self.stream_error_body.clone(), + self.stream_response_headers.clone(), + ), + }); + return; + } - let Some(choice) = value.choices.into_iter().next() else { - continue; + // handle failed chunk parsing / validation: + let value: GroqChatChunk = match serde_json::from_value(parsed) { + Ok(value) => value, + Err(error) => { + self.finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, }; + controller.enqueue(StreamPart::Error { + error: error.into(), + }); + return; + } + }; - if let Some(reason) = choice.finish_reason { - finish_reason = FinishReason { - unified: map_groq_finish_reason(Some(&reason)), - raw: Some(reason), - }; - } + if self.is_first_chunk { + self.is_first_chunk = false; + let metadata = response_metadata(value.id, value.created, value.model); + controller.enqueue(StreamPart::ResponseMetadata(metadata)); + } - let Some(delta) = choice.delta else { - continue; - }; + if let Some(chunk_usage) = value.x_groq.and_then(|x_groq| x_groq.usage) { + self.usage = Some(chunk_usage); + } - if let Some(reasoning) = delta.reasoning.filter(|text| !text.is_empty()) { - if !is_active_reasoning { - yield Ok(StreamPart::ReasoningStart { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = true; - } - yield Ok(StreamPart::ReasoningDelta { - id: "reasoning-0".to_string(), - delta: reasoning, - provider_metadata: None, - }); - } + let Some(choice) = value.choices.into_iter().next() else { + return; + }; - if let Some(text) = delta.content.filter(|text| !text.is_empty()) { - // end the active reasoning block before text starts - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = false; - } - if !is_active_text { - yield Ok(StreamPart::TextStart { - id: "txt-0".to_string(), - provider_metadata: None, - }); - is_active_text = true; - } - yield Ok(StreamPart::TextDelta { - id: "txt-0".to_string(), - delta: text, - provider_metadata: None, - }); - } + if let Some(reason) = choice.finish_reason { + self.finish_reason = FinishReason { + unified: map_groq_finish_reason(Some(&reason)), + raw: Some(reason), + }; + } - if let Some(deltas) = delta.tool_calls.filter(|deltas| !deltas.is_empty()) { - // end the active reasoning block before tool calls start - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = false; - } - for tool_call in &deltas { - let delta = StreamingToolCallDelta { - index: Some(tool_call.index), - id: tool_call.id.as_deref(), - r#type: tool_call.r#type.as_deref(), - name: tool_call.function.name.as_deref(), - arguments: tool_call.function.arguments.as_deref(), - provider_metadata: None, - }; - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; - } - for part in tool_parts.drain(..) { - yield Ok(part); - } - } - } - } + let Some(delta) = choice.delta else { + return; + }; - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { + if let Some(reasoning) = delta.reasoning.filter(|text| !text.is_empty()) { + if !self.is_active_reasoning { + controller.enqueue(StreamPart::ReasoningStart { id: "reasoning-0".to_string(), provider_metadata: None, }); + self.is_active_reasoning = true; } - if is_active_text { - yield Ok(StreamPart::TextEnd { + controller.enqueue(StreamPart::ReasoningDelta { + id: "reasoning-0".to_string(), + delta: reasoning, + provider_metadata: None, + }); + } + + if let Some(text) = delta.content.filter(|text| !text.is_empty()) { + // end the active reasoning block before text starts + self.end_active_reasoning(controller); + if !self.is_active_text { + controller.enqueue(StreamPart::TextStart { id: "txt-0".to_string(), provider_metadata: None, }); + self.is_active_text = true; } - tool_calls.finish(&mut tool_parts); - for part in tool_parts.drain(..) { - yield Ok(part); + controller.enqueue(StreamPart::TextDelta { + id: "txt-0".to_string(), + delta: text, + provider_metadata: None, + }); + } + + if let Some(deltas) = delta.tool_calls.filter(|deltas| !deltas.is_empty()) { + // end the active reasoning block before tool calls start + self.end_active_reasoning(controller); + for tool_call in &deltas { + let delta = StreamingToolCallDelta { + index: Some(tool_call.index), + id: tool_call.id.as_deref(), + r#type: tool_call.r#type.as_deref(), + name: tool_call.function.name.as_deref(), + arguments: tool_call.function.arguments.as_deref(), + provider_metadata: None, + }; + if let Err(error) = self.tool_calls.process(delta, &mut self.tool_parts) { + controller.error(error.into()); + return; + } + for part in self.tool_parts.drain(..) { + controller.enqueue(part); + } } + } + } - yield Ok(StreamPart::Finish { - finish_reason, - usage: convert_groq_usage(usage.as_ref()), + fn flush(mut self, controller: &mut TransformStreamController) { + self.end_active_reasoning(controller); + if self.is_active_text { + controller.enqueue(StreamPart::TextEnd { + id: "txt-0".to_string(), provider_metadata: None, }); - }; + } + self.tool_calls.finish(&mut self.tool_parts); + for part in self.tool_parts.drain(..) { + controller.enqueue(part); + } - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(RequestInfo { body: Some(body) }), - response: Some(StreamResponseInfo { - headers: Some(response_headers), - }), - }) + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason, + usage: convert_groq_usage(self.usage.as_ref()), + provider_metadata: None, + }); } } diff --git a/aimux-providers/src/mistral/model.rs b/aimux-providers/src/mistral/model.rs index de38ef565..b39fa8664 100644 --- a/aimux-providers/src/mistral/model.rs +++ b/aimux-providers/src/mistral/model.rs @@ -10,8 +10,9 @@ //! - Finish reasons include `model_length`. use aimux_core::tool::RawToolCall; +use std::collections::HashMap; + use async_trait::async_trait; -use futures::StreamExt; use serde_json::Value; use aimux_core::error::AiMuxError; @@ -19,9 +20,12 @@ use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; use aimux_core::stream_part::StreamPart; -use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; +use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage, Warning}; -use aimux_provider_utils::{StreamingToolCallDelta, StreamingToolCallTracker}; +use aimux_provider_utils::{ + StreamingToolCallDelta, StreamingToolCallTracker, TransformStreamController, Transformer, + pipe_through, +}; use crate::shared::EndpointConfig; @@ -285,210 +289,265 @@ impl LanguageModel for MistralModel { .await?; let response_headers = resp.response_headers; - let mut sse_stream = resp.value; - - let warnings = request_warnings(options, &self.model_id); - let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); - - let stream_error_url = endpoint; - let stream_error_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let generate_id_fn = self.generate_id; - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); - - let text_id = 0usize; - let mut text_started = false; - let mut reasoning_started = false; - let mut tool_calls = StreamingToolCallTracker::new().with_generate_id(generate_id_fn); - let mut tool_parts = Vec::new(); - let mut reasoning_id: Option = None; - let mut final_usage = Usage::default(); - let mut final_finish_reason: Option = None; - let mut response_metadata_emitted = false; - - while let Some(event) = sse_stream.next().await { - - match event { - Ok(parsed) => { - if include_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); - } - if let Some(err_obj) = parsed.get("error") { - yield Ok(StreamPart::Error { - error: super::mistral_stream_error( - err_obj, - &stream_error_url, - stream_error_body.clone(), - stream_response_headers.clone(), - ), - }); - continue; - } + let stream = pipe_through( + resp.value, + MistralChatStream { + warnings: request_warnings(options, &self.model_id), + include_raw_chunks: options.include_raw_chunks.unwrap_or(false), + stream_error_url: endpoint, + stream_error_body: body.clone(), + stream_response_headers: response_headers.clone(), + generate_id: self.generate_id, + text_started: false, + reasoning_started: false, + tool_calls: StreamingToolCallTracker::new().with_generate_id(self.generate_id), + tool_parts: Vec::new(), + reasoning_id: None, + final_usage: Usage::default(), + final_finish_reason: None, + response_metadata_emitted: false, + }, + ); - let chunk: StreamChunk = match serde_json::from_value(parsed) { - Ok(c) => c, - Err(e) => { - yield Ok(StreamPart::Error { error: e.into() }); - continue; - } - }; - - // Emit ResponseMetadata from the first valid chunk. - if !response_metadata_emitted { - response_metadata_emitted = true; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: chunk.id.clone(), - timestamp: chunk - .created - .and_then(|secs| { - chrono::DateTime::from_timestamp(secs as i64, 0) - }) - .map(|dt| dt.to_rfc3339()), - model_id: chunk.model.clone(), - })); - } + 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), + }), + }) + } +} - // Update usage. - if let Some(usage) = &chunk.usage { - final_usage = convert_usage(usage); - } +/// The id of the (single) text block of a streamed response. +const TEXT_ID: &str = "0"; - // Upstream processes only the first choice. - if let Some(choice) = chunk.choices.into_iter().next() { - // Reasoning content (from thinking parts in array content). - for reasoning_delta in choice.delta.content.as_ref().and_then(Value::as_array).into_iter().flatten() - .filter_map(extract_reasoning_part) { - if !reasoning_started { - // End any active text before starting reasoning. - if text_started { - yield Ok(StreamPart::TextEnd { - id: format!("{text_id}"), - provider_metadata: None, - }); - text_started = false; - } - let rid = generate_id_fn(); - reasoning_id = Some(rid.clone()); - reasoning_started = true; - yield Ok(StreamPart::ReasoningStart { id: rid, - provider_metadata: None, - }); - } - if let Some(ref rid) = reasoning_id { - yield Ok(StreamPart::ReasoningDelta { - id: rid.clone(), - delta: reasoning_delta, - provider_metadata: None, +/// The `TransformStream` of the Mistral chat model's `doStream`. +struct MistralChatStream { + warnings: Vec, + include_raw_chunks: bool, + stream_error_url: String, + stream_error_body: Value, + stream_response_headers: HashMap, + generate_id: aimux_provider_utils::IdGenerator, + text_started: bool, + reasoning_started: bool, + tool_calls: StreamingToolCallTracker, + tool_parts: Vec, + reasoning_id: Option, + final_usage: Usage, + final_finish_reason: Option, + response_metadata_emitted: bool, +} + +impl MistralChatStream { + /// Handle one parsed SSE `data:` payload. + fn transform_chunk( + &mut self, + parsed: Value, + controller: &mut TransformStreamController, + ) { + if self.include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: parsed.clone(), }); - } - } - - // Text content. - if let Some(text_delta) = - extract_text_content(&choice.delta.content) - && !text_delta.is_empty() { - if !text_started { - // End reasoning before starting text. - if reasoning_started { - if let Some(ref rid) = reasoning_id { - yield Ok(StreamPart::ReasoningEnd { - id: rid.clone(), - provider_metadata: None, + } + + if let Some(err_obj) = parsed.get("error") { + controller.enqueue(StreamPart::Error { + error: super::mistral_stream_error( + err_obj, + &self.stream_error_url, + self.stream_error_body.clone(), + self.stream_response_headers.clone(), + ), }); - } - reasoning_started = false; - reasoning_id = None; - } - yield Ok(StreamPart::TextStart { - id: format!("{text_id}"), - provider_metadata: None, - }); - text_started = true; - } - yield Ok(StreamPart::TextDelta { - id: format!("{text_id}"), - delta: text_delta, - provider_metadata: None, - }); - } - - // Tool calls: assembled by the shared tracker. - if let Some(tool_call_deltas) = choice.delta.tool_calls { - for dtc in &tool_call_deltas { - let delta = StreamingToolCallDelta { - index: dtc.index, - id: dtc.id.as_deref(), - name: dtc.function.name.as_deref(), - arguments: dtc.function.arguments.as_deref(), - ..Default::default() - }; - // A new call without a function name is - // invalid response data, as in the AI SDK; - // the stream ends. - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; - } - for part in tool_parts.drain(..) { - yield Ok(part); - } - } - } - - // Finish reason. - if let Some(reason) = choice.finish_reason { - final_finish_reason = Some(parse_finish_reason(&reason)); - } - } - } - Err(error) => { - if include_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: Value::Null }); - } - yield Ok(StreamPart::Error { error }); - } + return; + } + + let chunk: StreamChunk = match serde_json::from_value(parsed) { + Ok(c) => c, + Err(e) => { + controller.enqueue(StreamPart::Error { error: e.into() }); + return; + } + }; + + // Emit ResponseMetadata from the first valid chunk. + if !self.response_metadata_emitted { + self.response_metadata_emitted = true; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: chunk.id.clone(), + timestamp: chunk + .created + .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) + .map(|dt| dt.to_rfc3339()), + model_id: chunk.model.clone(), + })); + } + + // Update usage. + if let Some(usage) = &chunk.usage { + self.final_usage = convert_usage(usage); + } + + // Upstream processes only the first choice. + let Some(choice) = chunk.choices.into_iter().next() else { + return; + }; + + // Reasoning content (from thinking parts in array content). + for reasoning_delta in choice + .delta + .content + .as_ref() + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(extract_reasoning_part) + { + if !self.reasoning_started { + // End any active text before starting reasoning. + if self.text_started { + controller.enqueue(StreamPart::TextEnd { + id: TEXT_ID.to_string(), + provider_metadata: None, + }); + self.text_started = false; } + let rid = (self.generate_id)(); + self.reasoning_id = Some(rid.clone()); + self.reasoning_started = true; + controller.enqueue(StreamPart::ReasoningStart { + id: rid, + provider_metadata: None, + }); } + if let Some(ref rid) = self.reasoning_id { + controller.enqueue(StreamPart::ReasoningDelta { + id: rid.clone(), + delta: reasoning_delta, + provider_metadata: None, + }); + } + } - // Close any remaining open segments. - if text_started { - yield Ok(StreamPart::TextEnd { - id: format!("{text_id}"), + // Text content. + if let Some(text_delta) = extract_text_content(&choice.delta.content) + && !text_delta.is_empty() + { + if !self.text_started { + // End reasoning before starting text. + if self.reasoning_started { + if let Some(ref rid) = self.reasoning_id { + controller.enqueue(StreamPart::ReasoningEnd { + id: rid.clone(), + provider_metadata: None, + }); + } + self.reasoning_started = false; + self.reasoning_id = None; + } + controller.enqueue(StreamPart::TextStart { + id: TEXT_ID.to_string(), provider_metadata: None, }); + self.text_started = true; } - if reasoning_started - && let Some(ref rid) = reasoning_id { - yield Ok(StreamPart::ReasoningEnd { - id: rid.clone(), + controller.enqueue(StreamPart::TextDelta { + id: TEXT_ID.to_string(), + delta: text_delta, provider_metadata: None, }); + } + + // Tool calls: assembled by the shared tracker. + if let Some(tool_call_deltas) = choice.delta.tool_calls { + for dtc in &tool_call_deltas { + let delta = StreamingToolCallDelta { + index: dtc.index, + id: dtc.id.as_deref(), + name: dtc.function.name.as_deref(), + arguments: dtc.function.arguments.as_deref(), + ..Default::default() + }; + // A new call without a function name is invalid response + // data, as in the AI SDK; the stream ends. + if let Err(error) = self.tool_calls.process(delta, &mut self.tool_parts) { + controller.error(error.into()); + return; + } + for part in self.tool_parts.drain(..) { + controller.enqueue(part); } + } + } + + // Finish reason. + if let Some(reason) = choice.finish_reason { + self.final_finish_reason = Some(parse_finish_reason(&reason)); + } + } +} - tool_calls.finish(&mut tool_parts); - for part in tool_parts.drain(..) { - yield Ok(part); +impl Transformer for MistralChatStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + match event { + Ok(parsed) => self.transform_chunk(parsed, controller), + Err(error) => { + if self.include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: Value::Null, + }); + } + controller.enqueue(StreamPart::Error { error }); } + } + } - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }), - usage: final_usage, + fn flush(mut self, controller: &mut TransformStreamController) { + // Close any remaining open segments. + if self.text_started { + controller.enqueue(StreamPart::TextEnd { + id: TEXT_ID.to_string(), provider_metadata: None, }); - }; + } + if self.reasoning_started + && let Some(rid) = self.reasoning_id.take() + { + controller.enqueue(StreamPart::ReasoningEnd { + id: rid, + provider_metadata: None, + }); + } - 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), + self.tool_calls.finish(&mut self.tool_parts); + for part in self.tool_parts.drain(..) { + controller.enqueue(part); + } + + controller.enqueue(StreamPart::Finish { + finish_reason: self.final_finish_reason.unwrap_or(FinishReason { + unified: FinishReasonUnified::Other, + raw: None, }), - }) + usage: self.final_usage, + provider_metadata: None, + }); } } From 65fdc35fb353033fcda01fa28d1ea043d8c87f7a Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 12:42:20 +0800 Subject: [PATCH 09/16] refactor(providers): OpenAI and OpenAI-compatible chat streams as Transformer structs OpenAICompatibleChatStream and OpenAIChatStream hold the doStream state that lived in async_stream locals and are driven by pipe_through. The emitted part sequence is unchanged. Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-providers/src/openai/model.rs | 551 +++++++-------- aimux-providers/src/openai_compatible/chat.rs | 625 ++++++++++-------- 2 files changed, 638 insertions(+), 538 deletions(-) diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs index 8e41b8c27..fa079cd69 100644 --- a/aimux-providers/src/openai/model.rs +++ b/aimux-providers/src/openai/model.rs @@ -8,6 +8,8 @@ //! the conversion + streaming logic while supplying their own URL and auth. use aimux_core::tool::RawToolCall; +use std::collections::HashMap; + use async_trait::async_trait; use futures::StreamExt; use serde_json::{Value, json}; @@ -16,12 +18,13 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; -use aimux_core::shared::provider_namespace; +use aimux_core::shared::{Warning, provider_namespace}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; use aimux_provider_utils::{ - HttpRequest, StreamingToolCallDelta, StreamingToolCallTracker, TypeValidation, generate_id, + HttpRequest, StreamingToolCallDelta, StreamingToolCallTracker, TransformStreamController, + Transformer, TypeValidation, generate_id, pipe_through, }; use super::config::OpenAIModelConfig; @@ -371,313 +374,327 @@ pub(crate) async fn execute_stream( )); } - // M2 (RFC-0016): capture whether raw chunks should be emitted — the - // borrowed `options` cannot be moved into the generator. - let emit_raw_chunks = options.include_raw_chunks == Some(true); - let stream_error_url = endpoint; - let stream_error_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let stream = async_stream::stream! { - // First part: StreamStart. - yield Ok(StreamPart::StreamStart { warnings }); - - let text_id = 0usize; - let mut text_started = false; - let reasoning_id = "reasoning-0".to_string(); - let mut reasoning_started = false; - let mut final_usage = Usage::default(); - let mut final_usage_raw: Option = None; - let mut final_finish_reason: Option = None; - let mut response_metadata_emitted = false; - let mut final_logprobs: Option = None; - - // Streamed tool calls, correlated by wire id, index and function name - // (the AI SDK's StreamingToolCallTracker) and finalized on flush. + let transformer = OpenAIChatStream { + warnings, + // M2 (RFC-0016): whether raw chunks should be emitted. + include_raw_chunks: options.include_raw_chunks == Some(true), + url: endpoint, + request_body: body.clone(), + response_headers: response_headers.clone(), + text_started: false, + reasoning_started: false, + final_usage: Usage::default(), + final_usage_raw: None, + final_finish_reason: None, + response_metadata_emitted: false, + final_logprobs: None, // Same options as `@ai-sdk/openai`'s chat model. - let mut tool_calls = StreamingToolCallTracker::new() + tool_calls: StreamingToolCallTracker::new() .with_generate_id(generate_id) - .with_type_validation(TypeValidation::IfPresent); - let mut tool_parts = Vec::new(); - - // Process the first event (already peeked) then the rest. - let mut event_iter = - futures::stream::iter(first_event.into_iter()).chain(sse_stream); - - let mut stream_errored = false; - - while let Some(event) = event_iter.next().await { - match event { - Ok(parsed) => { - - // M2 (RFC-0016): emit the raw provider chunk for debugging - // before it is consumed below. JSON payloads only — the - // "[DONE]" sentinel is skipped by the early break above. - if emit_raw_chunks { - yield Ok(StreamPart::Raw { - raw_value: parsed.clone(), - }); - } - - // Check for mid-stream error. - if let Some(err_obj) = parsed.get("error") { - yield Ok(StreamPart::Error { - error: super::openai_stream_error( - err_obj, - &stream_error_url, - stream_error_body.clone(), - stream_response_headers.clone(), - ), - }); - stream_errored = true; - continue; - } - - // The chunk's top-level `usage`, taken from the raw JSON before - // the chunk is consumed below. `chunk_usage_raw` keeps the - // provider's original object for `Usage.raw` (M10). - let chunk_usage_raw: Option = parsed.get("usage").cloned(); - let chunk_usage: Option = chunk_usage_raw - .as_ref() - .and_then(|u| serde_json::from_value(u.clone()).ok()); - - // Parse as StreamChunk. - let chunk: StreamChunk = match serde_json::from_value(parsed) { - Ok(c) => c, - Err(e) => { - stream_errored = true; - yield Ok(StreamPart::Error { error: e.into() }); - continue; - } - }; + .with_type_validation(TypeValidation::IfPresent), + tool_parts: Vec::new(), + stream_errored: false, + }; + // Process the first event (already peeked) then the rest. + let events = futures::stream::iter(first_event).chain(sse_stream); + let stream = pipe_through(events, transformer); - // Emit ResponseMetadata from the first valid chunk. - if !response_metadata_emitted - && (chunk.id.as_ref().is_some_and(|id| !id.is_empty()) || chunk.model.as_ref().is_some_and(|model| !model.is_empty()) || chunk.created.is_some_and(|created| created != 0)) - { - response_metadata_emitted = true; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: chunk.id.clone(), - timestamp: chunk - .created - .filter(|created| *created != 0) - .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) - .map(|dt| dt.to_rfc3339()), - model_id: chunk.model.clone(), - })); - } + 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), + }), + }) +} - // Update usage from the chunk that carries it. - if let Some(usage) = &chunk_usage { - final_usage = convert_usage(usage, chunk_usage_raw.as_ref()); - final_usage_raw = Some(usage.clone()); - } +const TEXT_ID: &str = "0"; +const REASONING_ID: &str = "reasoning-0"; + +/// The `TransformStream` of the Chat Completions stream ([`execute_stream`]): +/// turns the parsed SSE events into stream parts. +struct OpenAIChatStream { + warnings: Vec, + include_raw_chunks: bool, + /// Request context for in-stream `error` payloads. + url: String, + request_body: Value, + response_headers: HashMap, + + text_started: bool, + reasoning_started: bool, + final_usage: Usage, + final_usage_raw: Option, + final_finish_reason: Option, + response_metadata_emitted: bool, + final_logprobs: Option, + /// Streamed tool calls, correlated by wire id, index and function name + /// (the AI SDK's StreamingToolCallTracker) and finalized on flush. + tool_calls: StreamingToolCallTracker, + tool_parts: Vec, + stream_errored: bool, +} - // Process choices. - for choice in chunk.choices { - // Capture logprobs from the finish_reason chunk. - if let Some(lp) = &choice.logprobs - && let Some(content) = lp.get("content") { - final_logprobs = Some(content.clone()); - } - - // Reasoning delta: prefer reasoning_content over reasoning. - let reasoning_delta = choice - .delta - .reasoning_content - .clone() - .or_else(|| choice.delta.reasoning.clone()); - if let Some(reasoning) = reasoning_delta - && !reasoning.is_empty() - { - if !reasoning_started { - reasoning_started = true; - yield Ok(StreamPart::ReasoningStart { - id: reasoning_id.clone(), - provider_metadata: None, - }); - } - yield Ok(StreamPart::ReasoningDelta { - id: reasoning_id.clone(), - delta: reasoning, +impl OpenAIChatStream { + fn end_reasoning(&mut self, controller: &mut TransformStreamController) { + if self.reasoning_started { + controller.enqueue(StreamPart::ReasoningEnd { + id: REASONING_ID.to_string(), provider_metadata: None, }); - } - - // Text delta. - if let Some(content) = choice.delta.content { - // End active reasoning block before text starts. - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, - }); - reasoning_started = false; - } - if !text_started { - text_started = true; - yield Ok(StreamPart::TextStart { - id: format!("{text_id}"), - provider_metadata: None, - }); - } - yield Ok(StreamPart::TextDelta { - id: format!("{text_id}"), - delta: content, - provider_metadata: None, - }); - } - - // Tool-call deltas. - if let Some(tool_call_deltas) = choice.delta.tool_calls { - // End active reasoning block before tool calls start. - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, + self.reasoning_started = false; + } + } +} + +impl Transformer for OpenAIChatStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + let parsed = match event { + Ok(parsed) => parsed, + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + self.stream_errored = true; + controller.enqueue(StreamPart::Error { error }); + if !recoverable { + controller.terminate(); + } + return; + } + }; + + // M2 (RFC-0016): emit the raw provider chunk for debugging before it + // is consumed below. + if self.include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: parsed.clone(), }); - reasoning_started = false; - } - for dtc in &tool_call_deltas { - let function = dtc.function.as_ref(); - let delta = StreamingToolCallDelta { - index: Some(dtc.index), - id: dtc.id.as_deref(), - r#type: dtc.r#type.as_deref(), - name: function.and_then(|f| f.name.as_deref()), - arguments: function.and_then(|f| f.arguments.as_deref()), - provider_metadata: None, - }; - // A malformed delta (new call without a - // function name, or a non-`function` type) - // is invalid response data, as in the AI SDK; - // the stream ends. - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; - } - for part in tool_parts.drain(..) { - yield Ok(part); - } - } - } - - // Annotations / citations (URL citations → Source). - if let Some(annotations) = choice.delta.annotations { - for (i, ann) in annotations.iter().enumerate() { - if ann.get("type").and_then(|v| v.as_str()) - == Some("url_citation") - && let Some(uc) = ann.get("url_citation") - { - yield Ok(StreamPart::Source(Source::Url { - id: format!("annotation-{i}"), - url: uc - .get("url") - .and_then(|v| v.as_str()) - .unwrap_or_default() - .to_string(), - title: uc - .get("title") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - provider_metadata: None, - })); - } - } - } - - // Finish reason. - if let Some(reason) = choice.finish_reason { - // Close any open reasoning segment. - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, + } + + // Check for mid-stream error. + if let Some(err_obj) = parsed.get("error") { + controller.enqueue(StreamPart::Error { + error: super::openai_stream_error( + err_obj, + &self.url, + self.request_body.clone(), + self.response_headers.clone(), + ), }); - reasoning_started = false; - } + self.stream_errored = true; + return; + } - stream_errored = false; - final_finish_reason = Some(parse_finish_reason(&reason)); - } - } + // The chunk's top-level `usage`, taken from the raw JSON before the + // chunk is consumed below. `chunk_usage_raw` keeps the provider's + // original object for `Usage.raw` (M10). + let chunk_usage_raw: Option = parsed.get("usage").cloned(); + let chunk_usage: Option = chunk_usage_raw + .as_ref() + .and_then(|u| serde_json::from_value(u.clone()).ok()); + + let chunk: StreamChunk = match serde_json::from_value(parsed) { + Ok(c) => c, + Err(e) => { + self.stream_errored = true; + controller.enqueue(StreamPart::Error { error: e.into() }); + return; + } + }; + + // Emit ResponseMetadata from the first valid chunk. + if !self.response_metadata_emitted + && (chunk.id.as_ref().is_some_and(|id| !id.is_empty()) + || chunk.model.as_ref().is_some_and(|model| !model.is_empty()) + || chunk.created.is_some_and(|created| created != 0)) + { + self.response_metadata_emitted = true; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: chunk.id.clone(), + timestamp: chunk + .created + .filter(|created| *created != 0) + .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) + .map(|dt| dt.to_rfc3339()), + model_id: chunk.model.clone(), + })); + } + + // Update usage from the chunk that carries it. + if let Some(usage) = &chunk_usage { + self.final_usage = convert_usage(usage, chunk_usage_raw.as_ref()); + self.final_usage_raw = Some(usage.clone()); + } + + for choice in chunk.choices { + // Capture logprobs from the finish_reason chunk. + if let Some(lp) = &choice.logprobs + && let Some(content) = lp.get("content") + { + self.final_logprobs = Some(content.clone()); + } + + // Reasoning delta: prefer reasoning_content over reasoning. + let reasoning_delta = choice + .delta + .reasoning_content + .clone() + .or_else(|| choice.delta.reasoning.clone()); + if let Some(reasoning) = reasoning_delta + && !reasoning.is_empty() + { + if !self.reasoning_started { + self.reasoning_started = true; + controller.enqueue(StreamPart::ReasoningStart { + id: REASONING_ID.to_string(), + provider_metadata: None, + }); } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - stream_errored = true; - yield Ok(StreamPart::Error { error }); - if !recoverable { + controller.enqueue(StreamPart::ReasoningDelta { + id: REASONING_ID.to_string(), + delta: reasoning, + provider_metadata: None, + }); + } + + // Text delta; an active reasoning block ends before text starts. + if let Some(content) = choice.delta.content { + self.end_reasoning(controller); + if !self.text_started { + self.text_started = true; + controller.enqueue(StreamPart::TextStart { + id: TEXT_ID.to_string(), + provider_metadata: None, + }); + } + controller.enqueue(StreamPart::TextDelta { + id: TEXT_ID.to_string(), + delta: content, + provider_metadata: None, + }); + } + + // Tool-call deltas; an active reasoning block ends first. + if let Some(tool_call_deltas) = choice.delta.tool_calls { + self.end_reasoning(controller); + for dtc in &tool_call_deltas { + let function = dtc.function.as_ref(); + let delta = StreamingToolCallDelta { + index: Some(dtc.index), + id: dtc.id.as_deref(), + r#type: dtc.r#type.as_deref(), + name: function.and_then(|f| f.name.as_deref()), + arguments: function.and_then(|f| f.arguments.as_deref()), + provider_metadata: None, + }; + // A malformed delta (new call without a function name, or + // a non-`function` type) is invalid response data, as in + // the AI SDK; the stream ends. + if let Err(error) = self.tool_calls.process(delta, &mut self.tool_parts) { + controller.error(error.into()); return; } + for part in self.tool_parts.drain(..) { + controller.enqueue(part); + } } } - } - // Close any remaining open reasoning segment. - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, - }); + // Annotations / citations (URL citations → Source). + if let Some(annotations) = choice.delta.annotations { + for (i, ann) in annotations.iter().enumerate() { + if ann.get("type").and_then(|v| v.as_str()) == Some("url_citation") + && let Some(uc) = ann.get("url_citation") + { + controller.enqueue(StreamPart::Source(Source::Url { + id: format!("annotation-{i}"), + url: uc + .get("url") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(), + title: uc + .get("title") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string), + provider_metadata: None, + })); + } + } + } + + // Finish reason; closes any open reasoning segment. + if let Some(reason) = choice.finish_reason { + self.end_reasoning(controller); + self.stream_errored = false; + self.final_finish_reason = Some(parse_finish_reason(&reason)); + } } + } - // Close any remaining open text segment. - if text_started { - yield Ok(StreamPart::TextEnd { - id: format!("{text_id}"), + fn flush(mut self, controller: &mut TransformStreamController) { + // Close any remaining open reasoning and text segments. + self.end_reasoning(controller); + if self.text_started { + controller.enqueue(StreamPart::TextEnd { + id: TEXT_ID.to_string(), provider_metadata: None, }); } // A parsable argument buffer can still be a prefix of a longer input: // like the AI SDK's tracker, finalize only when the stream flushes. - tool_calls.finish(&mut tool_parts); - for part in tool_parts.drain(..) { - yield Ok(part); + self.tool_calls.finish(&mut self.tool_parts); + for part in self.tool_parts.drain(..) { + controller.enqueue(part); } - // Build provider metadata for the Finish part. + // Provider metadata for the Finish part: logprobs + prediction tokens. let mut pm_openai = serde_json::json!({}); - if let Some(ref logprobs) = final_logprobs { + if let Some(ref logprobs) = self.final_logprobs { pm_openai["logprobs"] = json!(logprobs); } - // Prediction tokens from raw usage. - if let Some(ref raw_usage) = final_usage_raw - && let Some(ref details) = raw_usage.completion_tokens_details { - if let Some(apt) = details.accepted_prediction_tokens { - pm_openai["acceptedPredictionTokens"] = json!(apt); - } - if let Some(rpt) = details.rejected_prediction_tokens { - pm_openai["rejectedPredictionTokens"] = json!(rpt); - } + if let Some(ref raw_usage) = self.final_usage_raw + && let Some(ref details) = raw_usage.completion_tokens_details + { + if let Some(apt) = details.accepted_prediction_tokens { + pm_openai["acceptedPredictionTokens"] = json!(apt); } - let provider_metadata = provider_namespace("openai", pm_openai).expect("provider metadata must be an object"); + if let Some(rpt) = details.rejected_prediction_tokens { + pm_openai["rejectedPredictionTokens"] = json!(rpt); + } + } + let provider_metadata = + provider_namespace("openai", pm_openai).expect("provider metadata must be an object"); - // Final part: Finish. - yield Ok(StreamPart::Finish { - finish_reason: if stream_errored { + controller.enqueue(StreamPart::Finish { + finish_reason: if self.stream_errored { FinishReason { unified: FinishReasonUnified::Error, raw: None, } } else { - final_finish_reason.unwrap_or(FinishReason { + self.final_finish_reason.unwrap_or(FinishReason { unified: FinishReasonUnified::Other, raw: None, }) }, - usage: final_usage, + usage: self.final_usage, provider_metadata: Some(provider_metadata), }); - }; - - 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), - }), - }) + } } // ── Model listing (RFC-0027) ───────────────────────────────────────────────── diff --git a/aimux-providers/src/openai_compatible/chat.rs b/aimux-providers/src/openai_compatible/chat.rs index 08991b4a8..079d26ad2 100644 --- a/aimux-providers/src/openai_compatible/chat.rs +++ b/aimux-providers/src/openai_compatible/chat.rs @@ -8,21 +8,26 @@ use aimux_core::tool::RawToolCall; use std::collections::{HashMap, HashSet}; use async_trait::async_trait; -use futures::StreamExt; use serde_json::{Map, 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::shared::{RequestInfo, ResponseInfo, StreamResponseInfo, provider_namespace}; +use aimux_core::shared::{ + RequestInfo, ResponseInfo, SharedProviderMetadata, StreamResponseInfo, Warning, + provider_namespace, +}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{ FinishReason, FinishReasonUnified, InputTokenUsage, OutputTokenUsage, ResponseMetadata, Usage, }; -use aimux_provider_utils::{StreamingToolCallDelta, StreamingToolCallTracker, generate_id}; +use aimux_provider_utils::{ + StreamingToolCallDelta, StreamingToolCallTracker, TransformStreamController, Transformer, + generate_id, pipe_through, +}; -use super::config::CompatModelConfig; +use super::config::{CompatModelConfig, ConvertUsage}; use super::convert::{ChatBodySpec, RequestBodyResult, build_request_body, parse_finish_reason}; use super::types::{ChatCompletionResponse, StreamChunk, UsageResponse}; @@ -329,297 +334,375 @@ impl LanguageModel for OpenAICompatibleChatModel { ) .await?; let response_headers = resp.response_headers; - let mut sse_stream = resp.value; - - let convert_usage = self.config.chat.convert_usage; - let stream_usage_key = self.config.chat.stream_usage_key.clone(); - let emit_raw_chunks = options.include_raw_chunks == Some(true); - let stream_error_url = endpoint; - let stream_error_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); - - let text_id = "txt-0"; - let mut text_started = false; - let reasoning_id = "reasoning-0".to_string(); - let mut reasoning_started = false; - let mut final_usage_raw = None; - let mut final_usage_parsed: Option = None; - let mut final_finish_reason: Option = None; - let mut response_metadata_emitted = false; - - let mut tool_calls = StreamingToolCallTracker::new().with_generate_id(generate_id); - let mut tool_parts = Vec::new(); - let signature_metadata = |extra: Option<&Value>| { - thought_signature(extra).map(|signature| { - provider_namespace(&metadata_key, json!({ "thoughtSignature": signature })) - .expect("provider metadata must be an object") - }) - }; + let sse_stream = resp.value; + + let transformer = OpenAICompatibleChatStream { + warnings, + metadata_key, + convert_usage: self.config.chat.convert_usage, + stream_usage_key: self.config.chat.stream_usage_key.clone(), + include_raw_chunks: options.include_raw_chunks == Some(true), + url: endpoint, + request_body: body.clone(), + response_headers: response_headers.clone(), + text_started: false, + reasoning_started: false, + final_usage_raw: None, + final_usage_parsed: None, + final_finish_reason: None, + response_metadata_emitted: false, + tool_calls: StreamingToolCallTracker::new().with_generate_id(generate_id), + tool_parts: Vec::new(), + pending: HashMap::new(), + forwarded: HashSet::new(), + }; + let stream = pipe_through(sse_stream, transformer); + + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(RequestInfo { body: Some(body) }), + response: Some(StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} - // Some compatible servers send the first delta of a call without - // `function.name`; buffer by index until the name is known. - let mut pending: HashMap = HashMap::new(); - let mut forwarded: HashSet = HashSet::new(); +const TEXT_ID: &str = "txt-0"; +const REASONING_ID: &str = "reasoning-0"; + +/// The `TransformStream` of [`OpenAICompatibleChatModel::do_stream`]: turns +/// the parsed SSE events into stream parts. +struct OpenAICompatibleChatStream { + warnings: Vec, + metadata_key: String, + convert_usage: ConvertUsage, + stream_usage_key: Option, + include_raw_chunks: bool, + /// Request context for in-stream `error` payloads. + url: String, + request_body: Value, + response_headers: HashMap, - while let Some(event) = sse_stream.next().await { - match event { - Ok(parsed) => { - if emit_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); - } - if let Some(error) = parsed.get("error") { - yield Ok(StreamPart::Error { - error: stream_error( - &metadata_key, - error, - &stream_error_url, - stream_error_body.clone(), - stream_response_headers.clone(), - ), + text_started: bool, + reasoning_started: bool, + final_usage_raw: Option, + final_usage_parsed: Option, + final_finish_reason: Option, + response_metadata_emitted: bool, + tool_calls: StreamingToolCallTracker, + tool_parts: Vec, + /// Some compatible servers send the first delta of a call without + /// `function.name`; buffer by index until the name is known. + pending: HashMap, + forwarded: HashSet, +} + +impl OpenAICompatibleChatStream { + fn signature_metadata(&self, extra: Option<&Value>) -> Option { + thought_signature(extra).map(|signature| { + provider_namespace(&self.metadata_key, json!({ "thoughtSignature": signature })) + .expect("provider metadata must be an object") + }) + } + + fn enqueue_tool_parts(&mut self, controller: &mut TransformStreamController) { + for part in self.tool_parts.drain(..) { + controller.enqueue(part); + } + } + + fn error_finish_reason() -> FinishReason { + FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + } + } +} + +impl Transformer for OpenAICompatibleChatStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + let parsed = match event { + Ok(parsed) => parsed, + Err(error) => { + self.final_finish_reason = Some(Self::error_finish_reason()); + controller.enqueue(StreamPart::Error { error }); + return; + } + }; + if self.include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: parsed.clone(), + }); + } + if let Some(error) = parsed.get("error") { + controller.enqueue(StreamPart::Error { + error: stream_error( + &self.metadata_key, + error, + &self.url, + self.request_body.clone(), + self.response_headers.clone(), + ), + }); + self.final_finish_reason = Some(Self::error_finish_reason()); + return; + } + + let chunk_usage_raw: Option = match &self.stream_usage_key { + Some(key) => parsed.get(key).and_then(|v| v.get("usage")).cloned(), + None => parsed.get("usage").cloned(), + } + .filter(|usage| !usage.is_null()); + + let chunk: StreamChunk = match serde_json::from_value(parsed) { + Ok(chunk) => chunk, + Err(e) => { + self.final_finish_reason = Some(Self::error_finish_reason()); + controller.enqueue(StreamPart::Error { error: e.into() }); + return; + } + }; + + if !self.response_metadata_emitted + && (chunk.id.as_ref().is_some_and(|id| !id.is_empty()) + || chunk.model.as_ref().is_some_and(|model| !model.is_empty()) + || chunk.created.is_some_and(|created| created != 0)) + { + self.response_metadata_emitted = true; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: chunk.id.clone(), + timestamp: chunk + .created + .filter(|created| *created != 0) + .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) + .map(|dt| dt.to_rfc3339()), + model_id: chunk.model.clone(), + })); + } + + if let Some(raw_usage) = &chunk_usage_raw { + self.final_usage_raw = Some(raw_usage.clone()); + self.final_usage_parsed = serde_json::from_value(raw_usage.clone()).ok(); + } + + for choice in chunk.choices.into_iter().take(1) { + let reasoning_delta = choice + .delta + .reasoning_content + .clone() + .or_else(|| choice.delta.reasoning.clone()); + let mut delta_content = Vec::new(); + if let Some(text) = reasoning_delta.filter(|text| !text.is_empty()) { + delta_content.push(GenerateContent::Reasoning(ReasoningOutput { + text, + provider_metadata: None, + })); + } + delta_content.extend(content_parts(choice.delta.content.as_ref())); + for part in delta_content { + match part { + GenerateContent::Reasoning(ReasoningOutput { text, .. }) => { + if self.text_started { + controller.enqueue(StreamPart::TextEnd { + id: TEXT_ID.to_string(), + provider_metadata: None, }); - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - continue; + self.text_started = false; } - - let chunk_usage_raw: Option = match &stream_usage_key { - Some(key) => parsed.get(key).and_then(|v| v.get("usage")).cloned(), - None => parsed.get("usage").cloned(), + if !self.reasoning_started { + self.reasoning_started = true; + controller.enqueue(StreamPart::ReasoningStart { + id: REASONING_ID.to_string(), + provider_metadata: None, + }); } - .filter(|usage| !usage.is_null()); - - let chunk: StreamChunk = match serde_json::from_value(parsed) { - Ok(chunk) => chunk, - Err(e) => { - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - yield Ok(StreamPart::Error { error: e.into() }); - continue; - } - }; - - if !response_metadata_emitted - && (chunk.id.as_ref().is_some_and(|id| !id.is_empty()) - || chunk.model.as_ref().is_some_and(|model| !model.is_empty()) - || chunk.created.is_some_and(|created| created != 0)) - { - response_metadata_emitted = true; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: chunk.id.clone(), - timestamp: chunk - .created - .filter(|created| *created != 0) - .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) - .map(|dt| dt.to_rfc3339()), - model_id: chunk.model.clone(), - })); + controller.enqueue(StreamPart::ReasoningDelta { + id: REASONING_ID.to_string(), + delta: text, + provider_metadata: None, + }); + } + GenerateContent::Text { text, .. } => { + if self.reasoning_started { + controller.enqueue(StreamPart::ReasoningEnd { + id: REASONING_ID.to_string(), + provider_metadata: None, + }); + self.reasoning_started = false; } - - if let Some(raw_usage) = &chunk_usage_raw { - final_usage_raw = Some(raw_usage.clone()); - final_usage_parsed = serde_json::from_value(raw_usage.clone()).ok(); + if !self.text_started { + self.text_started = true; + controller.enqueue(StreamPart::TextStart { + id: TEXT_ID.to_string(), + provider_metadata: None, + }); } + controller.enqueue(StreamPart::TextDelta { + id: TEXT_ID.to_string(), + delta: text, + provider_metadata: None, + }); + } + _ => unreachable!(), + } + } - for choice in chunk.choices.into_iter().take(1) { - let reasoning_delta = choice - .delta - .reasoning_content - .clone() - .or_else(|| choice.delta.reasoning.clone()); - let mut delta_content = Vec::new(); - if let Some(text) = reasoning_delta.filter(|text| !text.is_empty()) { - delta_content.push(GenerateContent::Reasoning(ReasoningOutput { - text, - provider_metadata: None, - })); - } - delta_content.extend(content_parts(choice.delta.content.as_ref())); - for part in delta_content { - match part { - GenerateContent::Reasoning(ReasoningOutput { text, .. }) => { - if text_started { - yield Ok(StreamPart::TextEnd { id: text_id.to_string(), provider_metadata: None }); - text_started = false; - } - if !reasoning_started { - reasoning_started = true; - yield Ok(StreamPart::ReasoningStart { id: reasoning_id.clone(), provider_metadata: None }); - } - yield Ok(StreamPart::ReasoningDelta { id: reasoning_id.clone(), delta: text, provider_metadata: None }); - } - GenerateContent::Text { text, .. } => { - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { id: reasoning_id.clone(), provider_metadata: None }); - reasoning_started = false; - } - if !text_started { - text_started = true; - yield Ok(StreamPart::TextStart { id: text_id.to_string(), provider_metadata: None }); - } - yield Ok(StreamPart::TextDelta { id: text_id.to_string(), delta: text, provider_metadata: None }); - } - _ => unreachable!(), - } + if let Some(deltas) = choice.delta.tool_calls.filter(|deltas| !deltas.is_empty()) { + if self.reasoning_started { + controller.enqueue(StreamPart::ReasoningEnd { + id: REASONING_ID.to_string(), + provider_metadata: None, + }); + self.reasoning_started = false; + } + for dtc in &deltas { + let function = dtc.function.as_ref(); + let name = function + .and_then(|f| f.name.as_deref()) + .filter(|n| !n.trim().is_empty()); + let arguments = function.and_then(|f| f.arguments.as_deref()); + let buffered; + let delta = match dtc.index { + Some(index) if !self.forwarded.contains(&index) => { + let entry = self.pending.entry(index).or_default(); + if entry.id.is_none() { + entry.id.clone_from(&dtc.id); } - - if let Some(deltas) = choice.delta.tool_calls.filter(|deltas| !deltas.is_empty()) { - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, - }); - reasoning_started = false; - } - for dtc in &deltas { - let function = dtc.function.as_ref(); - let name = function - .and_then(|f| f.name.as_deref()) - .filter(|n| !n.trim().is_empty()); - let arguments = function.and_then(|f| f.arguments.as_deref()); - let buffered; - let delta = match dtc.index { - Some(index) if !forwarded.contains(&index) => { - let entry = pending.entry(index).or_default(); - if entry.id.is_none() { - entry.id.clone_from(&dtc.id); - } - if entry.extra.is_null() - && let Some(extra) = &dtc.extra_content - { - entry.extra = extra.clone(); - } - if let Some(arguments) = arguments { - entry.arguments.push_str(arguments); - } - let Some(name) = name else { continue }; - buffered = pending.remove(&index).unwrap_or_default(); - forwarded.insert(index); - StreamingToolCallDelta { - index: Some(index), - id: buffered.id.as_deref(), - name: Some(name), - arguments: Some(&buffered.arguments), - provider_metadata: signature_metadata(Some(&buffered.extra)), - ..Default::default() - } - } - _ => StreamingToolCallDelta { - index: dtc.index, - id: dtc.id.as_deref(), - name: function.and_then(|f| f.name.as_deref()), - arguments, - provider_metadata: signature_metadata(dtc.extra_content.as_ref()), - ..Default::default() - }, - }; - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; - } - for part in tool_parts.drain(..) { - yield Ok(part); - } - } + if entry.extra.is_null() + && let Some(extra) = &dtc.extra_content + { + entry.extra = extra.clone(); } - - if let Some(annotations) = choice.delta.annotations { - for (i, annotation) in annotations.iter().enumerate() { - if annotation.get("type").and_then(Value::as_str) - == Some("url_citation") - && let Some(citation) = annotation.get("url_citation") - { - yield Ok(StreamPart::Source(Source::Url { - id: format!("annotation-{i}"), - url: citation - .get("url") - .and_then(Value::as_str) - .unwrap_or_default() - .to_string(), - title: citation - .get("title") - .and_then(Value::as_str) - .map(str::to_string), - provider_metadata: None, - })); - } - } + if let Some(arguments) = arguments { + entry.arguments.push_str(arguments); } - - if let Some(reason) = choice.finish_reason { - final_finish_reason = Some(parse_finish_reason(&reason)); + let Some(name) = name else { continue }; + buffered = self.pending.remove(&index).unwrap_or_default(); + self.forwarded.insert(index); + StreamingToolCallDelta { + index: Some(index), + id: buffered.id.as_deref(), + name: Some(name), + arguments: Some(&buffered.arguments), + provider_metadata: self.signature_metadata(Some(&buffered.extra)), + ..Default::default() } } + _ => StreamingToolCallDelta { + index: dtc.index, + id: dtc.id.as_deref(), + name: function.and_then(|f| f.name.as_deref()), + arguments, + provider_metadata: self.signature_metadata(dtc.extra_content.as_ref()), + ..Default::default() + }, + }; + if let Err(error) = self.tool_calls.process(delta, &mut self.tool_parts) { + controller.error(error.into()); + return; } - Err(error) => { - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - yield Ok(StreamPart::Error { error }); - } + self.enqueue_tool_parts(controller); } } - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, - }); - } - if text_started { - yield Ok(StreamPart::TextEnd { - id: text_id.to_string(), - provider_metadata: None, - }); - } - for (index, entry) in pending { - let delta = StreamingToolCallDelta { - index: Some(index), - id: entry.id.as_deref(), - arguments: Some(&entry.arguments), - provider_metadata: signature_metadata(Some(&entry.extra)), - ..Default::default() - }; - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; - } - for part in tool_parts.drain(..) { - yield Ok(part); + if let Some(annotations) = choice.delta.annotations { + for (i, annotation) in annotations.iter().enumerate() { + if annotation.get("type").and_then(Value::as_str) == Some("url_citation") + && let Some(citation) = annotation.get("url_citation") + { + controller.enqueue(StreamPart::Source(Source::Url { + id: format!("annotation-{i}"), + url: citation + .get("url") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(), + title: citation + .get("title") + .and_then(Value::as_str) + .map(str::to_string), + provider_metadata: None, + })); + } } } - // A parsable argument buffer can still be a prefix of a longer - // input: finalize only when the stream flushes. - tool_calls.finish(&mut tool_parts); - for part in tool_parts.drain(..) { - yield Ok(part); + + if let Some(reason) = choice.finish_reason { + self.final_finish_reason = Some(parse_finish_reason(&reason)); } + } + } - let mut metadata = HashMap::new(); - metadata.insert(metadata_key.clone(), Map::new()); + fn flush(mut self, controller: &mut TransformStreamController) { + if self.reasoning_started { + controller.enqueue(StreamPart::ReasoningEnd { + id: REASONING_ID.to_string(), + provider_metadata: None, + }); + } + if self.text_started { + controller.enqueue(StreamPart::TextEnd { + id: TEXT_ID.to_string(), + provider_metadata: None, + }); + } + for (index, entry) in std::mem::take(&mut self.pending) { + let delta = StreamingToolCallDelta { + index: Some(index), + id: entry.id.as_deref(), + arguments: Some(&entry.arguments), + provider_metadata: self.signature_metadata(Some(&entry.extra)), + ..Default::default() + }; + if let Err(error) = self.tool_calls.process(delta, &mut self.tool_parts) { + controller.error(error.into()); + return; + } + self.enqueue_tool_parts(controller); + } + // A parsable argument buffer can still be a prefix of a longer + // input: finalize only when the stream flushes. + self.tool_calls.finish(&mut self.tool_parts); + for part in self.tool_parts.drain(..) { + controller.enqueue(part); + } - prediction_tokens(final_usage_parsed.as_ref(), metadata.entry(metadata_key).or_default()); + let mut metadata = HashMap::new(); + metadata.insert(self.metadata_key.clone(), Map::new()); - if final_finish_reason.is_none() { - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData("Response stream ended without a finish reason.".to_string()), - }); - } - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - }), - usage: convert_usage.convert(final_usage_raw.as_ref()), - provider_metadata: Some(metadata), - }); - }; + prediction_tokens( + self.final_usage_parsed.as_ref(), + metadata.entry(self.metadata_key).or_default(), + ); - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(RequestInfo { body: Some(body) }), - response: Some(StreamResponseInfo { - headers: Some(response_headers), - }), - }) + if self.final_finish_reason.is_none() { + controller.enqueue(StreamPart::Error { + error: AiMuxError::InvalidResponseData( + "Response stream ended without a finish reason.".to_string(), + ), + }); + } + controller.enqueue(StreamPart::Finish { + finish_reason: self + .final_finish_reason + .unwrap_or_else(Self::error_finish_reason), + usage: self.convert_usage.convert(self.final_usage_raw.as_ref()), + provider_metadata: Some(metadata), + }); } } From 70d7e841fd63553fd87b1b82a90648fc98a9e90e Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 12:43:29 +0800 Subject: [PATCH 10/16] feat(provider-utils): TransformStreamController::enqueue_error for skippable errors Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-provider-utils/src/transform_stream.rs | 26 ++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/aimux-provider-utils/src/transform_stream.rs b/aimux-provider-utils/src/transform_stream.rs index eca71fa2f..424928b84 100644 --- a/aimux-provider-utils/src/transform_stream.rs +++ b/aimux-provider-utils/src/transform_stream.rs @@ -31,6 +31,15 @@ impl TransformStreamController { } } + /// Emit `error` and keep the stream open: the next chunk is transformed + /// as usual. For errors a consumer may skip (a malformed event inside an + /// otherwise valid stream). + pub fn enqueue_error(&mut self, error: AiMuxError) { + if !self.closed { + self.queue.push(Err(error)); + } + } + /// Emit `error` and end the stream: no further chunk is transformed and /// `flush` is not called. pub fn error(&mut self, error: AiMuxError) { @@ -109,8 +118,8 @@ where mod tests { use super::*; - /// Doubles each number, stops with an error on a negative one, and - /// emits the running total on flush. + /// Doubles each number, reports a skippable error for zero, stops with an + /// error on a negative one, and emits the running total on flush. #[derive(Default)] struct Doubler { total: i32, @@ -125,6 +134,10 @@ mod tests { } fn transform(&mut self, chunk: i32, controller: &mut TransformStreamController) { + if chunk == 0 { + controller.enqueue_error(AiMuxError::InvalidResponseData("zero".into())); + return; + } if chunk < 0 { controller.error(AiMuxError::InvalidResponseData("negative".into())); return; @@ -150,6 +163,15 @@ mod tests { assert_eq!(run(vec![1, 2]).await, vec![Ok(0), Ok(2), Ok(4), Ok(3)]); } + #[tokio::test] + async fn enqueue_error_keeps_transforming_and_flushes() { + let out = run(vec![1, 0, 2]).await; + assert_eq!(out.len(), 5); + assert_eq!(out[..2], [Ok(0), Ok(2)]); + assert!(out[2].is_err()); + assert_eq!(out[3..], [Ok(4), Ok(3)]); + } + #[tokio::test] async fn error_ends_the_stream_without_flush() { let out = run(vec![1, -1, 5]).await; From 1824376c2e77412b756871299feafb239e346613 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 12:52:03 +0800 Subject: [PATCH 11/16] refactor(providers): DeepSeek chat stream as a Transformer struct Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-providers/src/deepseek/model.rs | 542 ++++++++++++++------------ 1 file changed, 300 insertions(+), 242 deletions(-) diff --git a/aimux-providers/src/deepseek/model.rs b/aimux-providers/src/deepseek/model.rs index 8d9bd469d..272858c2d 100644 --- a/aimux-providers/src/deepseek/model.rs +++ b/aimux-providers/src/deepseek/model.rs @@ -4,7 +4,6 @@ use aimux_core::tool::RawToolCall; use std::collections::{BTreeMap, HashMap}; use async_trait::async_trait; -use futures::StreamExt; use serde_json::{Map, Value, json}; use aimux_core::error::AiMuxError; @@ -16,7 +15,10 @@ use aimux_core::stream_part::StreamPart; use aimux_core::types::{ FinishReason, FinishReasonUnified, ReasoningEffort, ResponseMetadata, Warning, }; -use aimux_provider_utils::{StreamingToolCallDelta, StreamingToolCallTracker, generate_id}; +use aimux_provider_utils::{ + StreamingToolCallDelta, StreamingToolCallTracker, TransformStreamController, Transformer, + generate_id, pipe_through, +}; use super::convert::convert_to_deepseek_chat_messages; use super::finish_reason::map_deepseek_finish_reason; @@ -496,279 +498,335 @@ impl LanguageModel for DeepSeekChatLanguageModel { ) .await?; let response_headers = resp.response_headers; - let sse_stream = resp.value; - - let provider_options_name = self.provider_options_name().to_string(); - let emit_raw_chunks = options.include_raw_chunks == Some(true); - let stream_error_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); - - let mut tool_calls = StreamingToolCallTracker::new().with_generate_id(generate_id); - let mut tool_parts = Vec::new(); - let mut finish_reason = FinishReason { unified: FinishReasonUnified::Other, raw: None }; - let mut usage: Option = None; - let mut system_fingerprint: Option = None; - let mut is_first_chunk = true; - let mut is_active_reasoning = false; - let mut is_active_text = false; - let mut response_object: Option = None; - let mut choice_index: Option = None; - let mut message_role: Option = None; - let mut tool_call_types: BTreeMap = BTreeMap::new(); - let mut content_logprobs: Vec = Vec::new(); - let mut reasoning_logprobs: Vec = Vec::new(); - - let mut events = sse_stream; - while let Some(event) = events.next().await { - let parsed = match event { - Ok(parsed) => parsed, - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; - } - continue; - } - }; - // Emit the raw chunk if requested (before anything else). - if emit_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); - } + let stream = pipe_through( + resp.value, + DeepSeekChatStream { + warnings, + emit_raw_chunks: options.include_raw_chunks == Some(true), + endpoint, + stream_error_body: body.clone(), + stream_response_headers: response_headers.clone(), + provider_options_name: self.provider_options_name().to_string(), + tool_calls: StreamingToolCallTracker::new().with_generate_id(generate_id), + tool_parts: Vec::new(), + finish_reason: FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }, + usage: None, + system_fingerprint: None, + is_first_chunk: true, + is_active_reasoning: false, + is_active_text: false, + response_object: None, + choice_index: None, + message_role: None, + tool_call_types: BTreeMap::new(), + content_logprobs: Vec::new(), + reasoning_logprobs: Vec::new(), + }, + ); - // handle error chunks: - if parsed.get("error").is_some() { - finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; - yield Ok(StreamPart::Error { - error: deepseek_stream_error( - &parsed, - &endpoint, - stream_error_body.clone(), - stream_response_headers.clone(), - ), - }); - continue; - } + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(RequestInfo { body: Some(body) }), + response: Some(StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} - // handle failed chunk validation: - let chunk: ChatChunk = match serde_json::from_value(parsed) { - Ok(chunk) => chunk, - Err(error) => { - finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; - yield Ok(StreamPart::Error { error: error.into() }); - continue; - } - }; +/// The `TransformStream` of `DeepSeekChatLanguageModel.doStream`. +struct DeepSeekChatStream { + warnings: Vec, + emit_raw_chunks: bool, + endpoint: String, + stream_error_body: Value, + stream_response_headers: HashMap, + provider_options_name: String, + tool_calls: StreamingToolCallTracker, + tool_parts: Vec, + finish_reason: FinishReason, + usage: Option, + system_fingerprint: Option, + is_first_chunk: bool, + is_active_reasoning: bool, + is_active_text: bool, + response_object: Option, + choice_index: Option, + message_role: Option, + tool_call_types: BTreeMap, + content_logprobs: Vec, + reasoning_logprobs: Vec, +} - if is_first_chunk { - is_first_chunk = false; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: chunk.id.clone(), - timestamp: timestamp(chunk.created), - model_id: chunk.model.clone(), - })); - } +impl DeepSeekChatStream { + fn end_active_reasoning(&mut self, controller: &mut TransformStreamController) { + if self.is_active_reasoning { + controller.enqueue(StreamPart::ReasoningEnd { + id: "reasoning-0".to_string(), + provider_metadata: None, + }); + self.is_active_reasoning = false; + } + } +} - if let Some(chunk_usage) = chunk.usage.filter(|usage| !usage.is_null()) { - usage = Some(chunk_usage); - } +impl Transformer for DeepSeekChatStream { + type Input = Result; + type Output = StreamPart; - if chunk.object.is_some() { - response_object = chunk.object; - } + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } - // The fingerprint is repeated on stream chunks; keep the - // latest non-null value in case it changes during the response. - if chunk.system_fingerprint.is_some() { - system_fingerprint = chunk.system_fingerprint; + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + let parsed = match event { + Ok(parsed) => parsed, + Err(error) => { + if error.is_recoverable_stream_error() { + controller.enqueue_error(error); + } else { + controller.error(error); } + return; + } + }; + + // Emit the raw chunk if requested (before anything else). + if self.emit_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: parsed.clone(), + }); + } + + // handle error chunks: + if parsed.get("error").is_some() { + self.finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }; + controller.enqueue(StreamPart::Error { + error: deepseek_stream_error( + &parsed, + &self.endpoint, + self.stream_error_body.clone(), + self.stream_response_headers.clone(), + ), + }); + return; + } - let Some(choice) = chunk.choices.into_iter().next() else { - continue; + // handle failed chunk validation: + let chunk: ChatChunk = match serde_json::from_value(parsed) { + Ok(chunk) => chunk, + Err(error) => { + self.finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, }; + controller.enqueue(StreamPart::Error { + error: error.into(), + }); + return; + } + }; - if choice.index.is_some() { - choice_index = choice.index; - } + if self.is_first_chunk { + self.is_first_chunk = false; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: chunk.id.clone(), + timestamp: timestamp(chunk.created), + model_id: chunk.model.clone(), + })); + } - if let Some(reason) = choice.finish_reason { - finish_reason = FinishReason { - unified: map_deepseek_finish_reason(Some(&reason)), - raw: Some(reason), - }; - } + if let Some(chunk_usage) = chunk.usage.filter(|usage| !usage.is_null()) { + self.usage = Some(chunk_usage); + } - if let Some(logprobs) = &choice.logprobs { - for (key, sink) in [ - ("content", &mut content_logprobs), - ("reasoning_content", &mut reasoning_logprobs), - ] { - if let Some(entries) = logprobs.get(key).and_then(Value::as_array) { - sink.extend(entries.iter().cloned()); - } - } - } + if chunk.object.is_some() { + self.response_object = chunk.object; + } - let Some(delta) = choice.delta else { - continue; - }; + // The fingerprint is repeated on stream chunks; keep the + // latest non-null value in case it changes during the response. + if chunk.system_fingerprint.is_some() { + self.system_fingerprint = chunk.system_fingerprint; + } - if delta.role.is_some() { - message_role = delta.role; - } + let Some(choice) = chunk.choices.into_iter().next() else { + return; + }; - // enqueue reasoning before text deltas: - if let Some(reasoning) = delta - .reasoning_content - .filter(|reasoning| !reasoning.is_empty()) - { - if !is_active_reasoning { - yield Ok(StreamPart::ReasoningStart { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = true; - } - yield Ok(StreamPart::ReasoningDelta { - id: "reasoning-0".to_string(), - delta: reasoning, - provider_metadata: None, - }); - } + if choice.index.is_some() { + self.choice_index = choice.index; + } - if let Some(text) = delta.content.filter(|text| !text.is_empty()) { - if !is_active_text { - yield Ok(StreamPart::TextStart { - id: "txt-0".to_string(), - provider_metadata: None, - }); - is_active_text = true; - } - - // end reasoning when text starts: - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = false; - } - - yield Ok(StreamPart::TextDelta { - id: "txt-0".to_string(), - delta: text, - provider_metadata: None, - }); - } + if let Some(reason) = choice.finish_reason { + self.finish_reason = FinishReason { + unified: map_deepseek_finish_reason(Some(&reason)), + raw: Some(reason), + }; + } - if let Some(deltas) = delta.tool_calls.filter(|deltas| !deltas.is_empty()) { - // end reasoning when tool calls start: - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = false; - } - - for tool_call in &deltas { - if let Some(kind) = &tool_call.r#type { - tool_call_types.insert(tool_call.index, kind.clone()); - } - let delta = StreamingToolCallDelta { - index: Some(tool_call.index), - id: tool_call.id.as_deref(), - r#type: tool_call.r#type.as_deref(), - name: tool_call.function.name.as_deref(), - arguments: tool_call.function.arguments.as_deref(), - provider_metadata: None, - }; - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; - } - for part in tool_parts.drain(..) { - yield Ok(part); - } - } + if let Some(logprobs) = &choice.logprobs { + for (key, sink) in [ + ("content", &mut self.content_logprobs), + ("reasoning_content", &mut self.reasoning_logprobs), + ] { + if let Some(entries) = logprobs.get(key).and_then(Value::as_array) { + sink.extend(entries.iter().cloned()); } } + } + + let Some(delta) = choice.delta else { + return; + }; + + if delta.role.is_some() { + self.message_role = delta.role; + } - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { + // enqueue reasoning before text deltas: + if let Some(reasoning) = delta + .reasoning_content + .filter(|reasoning| !reasoning.is_empty()) + { + if !self.is_active_reasoning { + controller.enqueue(StreamPart::ReasoningStart { id: "reasoning-0".to_string(), provider_metadata: None, }); + self.is_active_reasoning = true; } - if is_active_text { - yield Ok(StreamPart::TextEnd { + controller.enqueue(StreamPart::ReasoningDelta { + id: "reasoning-0".to_string(), + delta: reasoning, + provider_metadata: None, + }); + } + + if let Some(text) = delta.content.filter(|text| !text.is_empty()) { + if !self.is_active_text { + controller.enqueue(StreamPart::TextStart { id: "txt-0".to_string(), provider_metadata: None, }); - } - tool_calls.finish(&mut tool_parts); - for part in tool_parts.drain(..) { - yield Ok(part); + self.is_active_text = true; } - let mut metadata = Map::new(); - for (key, name) in [ - ("prompt_cache_hit_tokens", "promptCacheHitTokens"), - ("prompt_cache_miss_tokens", "promptCacheMissTokens"), - ] { - if let Some(value) = usage - .as_ref() - .and_then(|usage| usage.get(key)) - .filter(|value| !value.is_null()) - { - metadata.insert(name.into(), value.clone()); + // end reasoning when text starts: + self.end_active_reasoning(controller); + + controller.enqueue(StreamPart::TextDelta { + id: "txt-0".to_string(), + delta: text, + provider_metadata: None, + }); + } + + if let Some(deltas) = delta.tool_calls.filter(|deltas| !deltas.is_empty()) { + // end reasoning when tool calls start: + self.end_active_reasoning(controller); + + for tool_call in &deltas { + if let Some(kind) = &tool_call.r#type { + self.tool_call_types.insert(tool_call.index, kind.clone()); + } + let delta = StreamingToolCallDelta { + index: Some(tool_call.index), + id: tool_call.id.as_deref(), + r#type: tool_call.r#type.as_deref(), + name: tool_call.function.name.as_deref(), + arguments: tool_call.function.arguments.as_deref(), + provider_metadata: None, + }; + if let Err(error) = self.tool_calls.process(delta, &mut self.tool_parts) { + controller.error(error.into()); + return; + } + for part in self.tool_parts.drain(..) { + controller.enqueue(part); } } - if let Some(object) = response_object { - metadata.insert("responseObject".into(), json!(object)); - } - if let Some(index) = choice_index { - metadata.insert("choiceIndex".into(), json!(index)); - } - if let Some(role) = message_role { - metadata.insert("messageRole".into(), json!(role)); - } - if !tool_call_types.is_empty() { - metadata.insert( - "toolCallTypes".into(), - json!(tool_call_types.into_values().collect::>()), - ); - } - let mut logprobs = Map::new(); - if !content_logprobs.is_empty() { - logprobs.insert("content".into(), Value::Array(content_logprobs)); - } - if !reasoning_logprobs.is_empty() { - logprobs.insert("reasoning_content".into(), Value::Array(reasoning_logprobs)); - } - if !logprobs.is_empty() { - metadata.insert("logprobs".into(), Value::Object(logprobs)); - } - if let Some(system_fingerprint) = system_fingerprint { - metadata.insert("systemFingerprint".into(), json!(system_fingerprint)); - } - yield Ok(StreamPart::Finish { - finish_reason, - usage: convert_deepseek_usage(usage.as_ref()), - provider_metadata: Some(provider_namespace(&provider_options_name, Value::Object(metadata)).expect("metadata payload is an object")), + } + } + + fn flush(mut self, controller: &mut TransformStreamController) { + self.end_active_reasoning(controller); + if self.is_active_text { + controller.enqueue(StreamPart::TextEnd { + id: "txt-0".to_string(), + provider_metadata: None, }); - }; + } + self.tool_calls.finish(&mut self.tool_parts); + for part in self.tool_parts.drain(..) { + controller.enqueue(part); + } - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(RequestInfo { body: Some(body) }), - response: Some(StreamResponseInfo { - headers: Some(response_headers), - }), - }) + let mut metadata = Map::new(); + for (key, name) in [ + ("prompt_cache_hit_tokens", "promptCacheHitTokens"), + ("prompt_cache_miss_tokens", "promptCacheMissTokens"), + ] { + if let Some(value) = self + .usage + .as_ref() + .and_then(|usage| usage.get(key)) + .filter(|value| !value.is_null()) + { + metadata.insert(name.into(), value.clone()); + } + } + if let Some(object) = self.response_object { + metadata.insert("responseObject".into(), json!(object)); + } + if let Some(index) = self.choice_index { + metadata.insert("choiceIndex".into(), json!(index)); + } + if let Some(role) = self.message_role { + metadata.insert("messageRole".into(), json!(role)); + } + if !self.tool_call_types.is_empty() { + metadata.insert( + "toolCallTypes".into(), + json!(self.tool_call_types.into_values().collect::>()), + ); + } + let mut logprobs = Map::new(); + if !self.content_logprobs.is_empty() { + logprobs.insert("content".into(), Value::Array(self.content_logprobs)); + } + if !self.reasoning_logprobs.is_empty() { + logprobs.insert( + "reasoning_content".into(), + Value::Array(self.reasoning_logprobs), + ); + } + if !logprobs.is_empty() { + metadata.insert("logprobs".into(), Value::Object(logprobs)); + } + if let Some(system_fingerprint) = self.system_fingerprint { + metadata.insert("systemFingerprint".into(), json!(system_fingerprint)); + } + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason, + usage: convert_deepseek_usage(self.usage.as_ref()), + provider_metadata: Some( + provider_namespace(&self.provider_options_name, Value::Object(metadata)) + .expect("metadata payload is an object"), + ), + }); } } From ed72d74028d7c624591e29e84397fd737d4d2ae5 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 13:07:42 +0800 Subject: [PATCH 12/16] refactor(providers): Google, Vertex and OpenAI-compatible completion streams as Transformer structs Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-providers/src/google/model.rs | 950 ++++++++------- .../src/openai_compatible/completion.rs | 180 ++- aimux-providers/src/vertex/model.rs | 1029 ++++++++++------- 3 files changed, 1277 insertions(+), 882 deletions(-) diff --git a/aimux-providers/src/google/model.rs b/aimux-providers/src/google/model.rs index daadf682f..c119c77b8 100644 --- a/aimux-providers/src/google/model.rs +++ b/aimux-providers/src/google/model.rs @@ -1,5 +1,7 @@ //! Google Gemini language model — implements `LanguageModel`. +use std::collections::{HashMap, HashSet}; + use aimux_core::tool::RawToolCall; use aimux_core::tool::ToolResult; use async_trait::async_trait; @@ -15,8 +17,9 @@ use aimux_core::result::{ use aimux_core::shared::{FileBytes, GeneratedFileData}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{ - FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, + FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, Warning, }; +use aimux_provider_utils::{TransformStreamController, Transformer, pipe_through}; use aimux_core::language_model::SupportedUrls; @@ -25,7 +28,7 @@ use super::convert::{ parse_finish_reason, validate_call_options, }; use super::options::google_metadata; -use super::types::{Candidate, GenerateContentResponse, GoogleStreamEvent}; +use super::types::{Candidate, GenerateContentResponse, GoogleStreamEvent, StreamChunk}; use crate::shared::EndpointConfig; /// A Google Gemini language model. @@ -198,444 +201,553 @@ impl LanguageModel for GoogleModel { let generate_id = self.generate_id; let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings: tool_warnings }); - - let mut sse_stream = futures::stream::iter(first_event.into_iter()).chain(sse_stream); - let mut text_id: Option = None; - let mut reasoning_id: Option = None; - let mut block_counter = 0usize; - let mut final_usage: Usage = Usage::default(); - let mut final_finish_reason: Option = None; - let mut has_tool_calls = false; - let mut response_metadata_emitted = false; - let mut stream_errored = false; - let mut prompt_blocked = false; - - // Provider-metadata accumulators (mirrors TS `lastGroundingMetadata` / - // `lastUrlContextMetadata` + the finishReason-chunk snapshot). - let mut last_grounding_metadata: Option = None; - let mut last_url_context_metadata: Option = None; - let mut last_prompt_feedback: Option = None; - let mut last_safety_ratings: Option = None; - let mut last_finish_message: Option = None; - let mut last_usage_metadata_value: Option = None; - - // Source dedup across chunks (url sources only). - let mut emitted_source_urls: std::collections::HashSet = - std::collections::HashSet::new(); - let mut source_id = 0usize; - - // Associates code-execution results / server tool responses with - // their preceding call. - let mut last_code_execution_tool_call_id: Option = None; - let mut last_server_tool_call_id: Option = None; - - while let Some(event) = sse_stream.next().await { - if stream_errored { - break; - } - let event = match event { - Ok(raw_value) => { - if include_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: raw_value.clone() }); - } - serde_json::from_value::(raw_value).map_err(AiMuxError::from) - } - Err(error) => Err(error), - }; - match event { - Ok(GoogleStreamEvent::Chunk(chunk)) => { - - // Emit ResponseMetadata from the first chunk that has - // a responseId (matches TS behaviour). - if !response_metadata_emitted - && let Some(id) = &chunk.response_id { - response_metadata_emitted = true; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: Some(id.clone()), - timestamp: None, - model_id: None, - })); - } + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + GoogleLanguageModelStream { + warnings: tool_warnings, + include_raw_chunks, + code_execution_tool_name, + stream_error_url, + stream_request_body, + stream_error_headers, + generate_id, + text_id: None, + reasoning_id: None, + block_counter: 0, + final_usage: Usage::default(), + final_finish_reason: None, + has_tool_calls: false, + response_metadata_emitted: false, + stream_errored: false, + prompt_blocked: false, + last_grounding_metadata: None, + last_url_context_metadata: None, + last_prompt_feedback: None, + last_safety_ratings: None, + last_finish_message: None, + last_usage_metadata_value: None, + emitted_source_urls: HashSet::new(), + source_id: 0, + last_code_execution_tool_call_id: None, + last_server_tool_call_id: None, + }, + ); - if let Some(usage) = &chunk.usage_metadata { - final_usage = convert_usage(usage); - last_usage_metadata_value = - serde_json::to_value(usage).ok(); - } - if let Some(pf) = &chunk.prompt_feedback - && !prompt_blocked { - last_prompt_feedback = Some(pf.clone()); - if let Some(reason) = confirmed_prompt_block_reason(Some(pf)) { - prompt_blocked = true; - final_finish_reason = Some(FinishReason { - unified: FinishReasonUnified::ContentFilter, raw: Some(reason), - }); - } - } + 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), + }), + }) + } +} - let Some(candidates) = chunk.candidates else { - continue; - }; - let Some(candidate) = candidates.into_iter().next() else { - continue; - }; +/// The `TransformStream` of `GoogleLanguageModel.doStream`. +struct GoogleLanguageModelStream { + warnings: Vec, + include_raw_chunks: bool, + code_execution_tool_name: String, + stream_error_url: String, + stream_request_body: Value, + stream_error_headers: HashMap, + generate_id: aimux_provider_utils::IdGenerator, + text_id: Option, + reasoning_id: Option, + block_counter: usize, + final_usage: Usage, + final_finish_reason: Option, + has_tool_calls: bool, + response_metadata_emitted: bool, + stream_errored: bool, + prompt_blocked: bool, + // Provider-metadata accumulators (mirrors TS `lastGroundingMetadata` / + // `lastUrlContextMetadata` + the finishReason-chunk snapshot). + last_grounding_metadata: Option, + last_url_context_metadata: Option, + last_prompt_feedback: Option, + last_safety_ratings: Option, + last_finish_message: Option, + last_usage_metadata_value: Option, + // Source dedup across chunks (url sources only). + emitted_source_urls: HashSet, + source_id: usize, + // Associates code-execution results / server tool responses with their + // preceding call. + last_code_execution_tool_call_id: Option, + last_server_tool_call_id: Option, +} - if let Some(gm) = &candidate.grounding_metadata { - last_grounding_metadata = Some(gm.clone()); - } - if let Some(ucm) = &candidate.url_context_metadata { - last_url_context_metadata = Some(ucm.clone()); - } +impl GoogleLanguageModelStream { + fn transform_chunk( + &mut self, + chunk: Box, + controller: &mut TransformStreamController, + ) { + // Emit ResponseMetadata from the first chunk that has + // a responseId (matches TS behaviour). + if !self.response_metadata_emitted + && let Some(id) = &chunk.response_id + { + self.response_metadata_emitted = true; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: Some(id.clone()), + timestamp: None, + model_id: None, + })); + } - if let Some(sr) = &candidate.safety_ratings { - last_safety_ratings = serde_json::to_value(sr).ok(); - } - if let Some(fm) = &candidate.finish_message { - last_finish_message = Some(json!(fm)); - } - if prompt_blocked { continue; } - - // Extract url sources from this chunk's grounding metadata - // (deduplicated across chunks; document sources are not - // emitted in the stream, matching TS). - let chunk_sources = - extract_sources(candidate.grounding_metadata.as_ref(), &mut source_id); - for src in chunk_sources { - if let GenerateContent::Source(Source::Url { - url, - id: _, - title, - provider_metadata: None, - }) = src - && emitted_source_urls.insert(url.clone()) { - yield Ok(StreamPart::Source(Source::Url { - id: generate_id(), - url, - title, - provider_metadata: None, - })); - } - } + if let Some(usage) = &chunk.usage_metadata { + self.final_usage = convert_usage(usage); + self.last_usage_metadata_value = serde_json::to_value(usage).ok(); + } + if let Some(pf) = &chunk.prompt_feedback + && !self.prompt_blocked + { + self.last_prompt_feedback = Some(pf.clone()); + if let Some(reason) = confirmed_prompt_block_reason(Some(pf)) { + self.prompt_blocked = true; + self.final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::ContentFilter, + raw: Some(reason), + }); + } + } - if let Some(parts) = - candidate.content.as_ref().and_then(|c| c.parts.as_ref()) + let Some(candidates) = chunk.candidates else { + return; + }; + let Some(candidate) = candidates.into_iter().next() else { + return; + }; + + if let Some(gm) = &candidate.grounding_metadata { + self.last_grounding_metadata = Some(gm.clone()); + } + if let Some(ucm) = &candidate.url_context_metadata { + self.last_url_context_metadata = Some(ucm.clone()); + } + + if let Some(sr) = &candidate.safety_ratings { + self.last_safety_ratings = serde_json::to_value(sr).ok(); + } + if let Some(fm) = &candidate.finish_message { + self.last_finish_message = Some(json!(fm)); + } + if self.prompt_blocked { + return; + } + + // Extract url sources from this chunk's grounding metadata + // (deduplicated across chunks; document sources are not + // emitted in the stream, matching TS). + let chunk_sources = + extract_sources(candidate.grounding_metadata.as_ref(), &mut self.source_id); + for src in chunk_sources { + if let GenerateContent::Source(Source::Url { + url, + id: _, + title, + provider_metadata: None, + }) = src + && self.emitted_source_urls.insert(url.clone()) + { + controller.enqueue(StreamPart::Source(Source::Url { + id: (self.generate_id)(), + url, + title, + provider_metadata: None, + })); + } + } + + if let Some(parts) = candidate.content.as_ref().and_then(|c| c.parts.as_ref()) { + for part in parts { + // thoughtSignature → provider_metadata (upstream :778-782) + let thought_sig_meta: Option = part + .get("thoughtSignature") + .and_then(|v| v.as_str()) + .filter(|signature| !signature.is_empty()) + .map(|s| google_metadata(json!({ "thoughtSignature": s }))); + + // text part + if let Some(text) = part.get("text").and_then(|v| v.as_str()) { + if text.is_empty() { + // Empty-text part may carry a thoughtSignature + // (upstream :784-794): emit as a zero-length + // delta with the signature metadata on + // whichever block is open — text or + // reasoning (mirrors the non-streaming + // path, which attaches to the last + // content item regardless of type). + if let Some(meta) = &thought_sig_meta + && let Some(id) = &self.text_id { - for part in parts { - // thoughtSignature → provider_metadata (upstream :778-782) - let thought_sig_meta: Option = part - .get("thoughtSignature") - .and_then(|v| v.as_str()) - .filter(|signature| !signature.is_empty()) - .map(|s| google_metadata(json!({ "thoughtSignature": s }))); - - // text part - if let Some(text) = part.get("text").and_then(|v| v.as_str()) { - if text.is_empty() { - // Empty-text part may carry a thoughtSignature - // (upstream :784-794): emit as a zero-length - // delta with the signature metadata on - // whichever block is open — text or - // reasoning (mirrors the non-streaming - // path, which attaches to the last - // content item regardless of type). - if let Some(meta) = &thought_sig_meta - && let Some(id) = &text_id { - yield Ok(StreamPart::TextDelta { - id: id.clone(), - delta: String::new(), - provider_metadata: Some(meta.clone()), - }); - - } - } else { - let is_thought = part.get("thought").and_then(serde_json::Value::as_bool).unwrap_or(false); - if is_thought { - // Close any open text block before starting reasoning - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); - } - // Start a reasoning block if not already active - if reasoning_id.is_none() { - let id = format!("{block_counter}"); - block_counter += 1; - reasoning_id = Some(id.clone()); - yield Ok(StreamPart::ReasoningStart { - id, - provider_metadata: thought_sig_meta.clone(), - }); - } - if let Some(id) = &reasoning_id { - yield Ok(StreamPart::ReasoningDelta { - id: id.clone(), - delta: text.to_string(), - provider_metadata: thought_sig_meta.clone(), - }); - } - } else { - // Close any open reasoning block before starting text - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); - } - if text_id.is_none() { - let id = format!("{block_counter}"); - block_counter += 1; - text_id = Some(id.clone()); - yield Ok(StreamPart::TextStart { id, provider_metadata: thought_sig_meta.clone() }); - } - if let Some(id) = &text_id { - yield Ok(StreamPart::TextDelta { - id: id.clone(), - delta: text.to_string(), - provider_metadata: thought_sig_meta.clone(), - }); - } - } - } - } else if let Some(ec) = part.get("executableCode") { - // Provider-executed code execution. - let has_code = ec - .get("code") - .and_then(|v| v.as_str()) - .map(|s| !s.is_empty()) - .unwrap_or(false); - if has_code { - let id = generate_id(); - last_code_execution_tool_call_id = Some(id.clone()); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id.clone(), - tool_name: code_execution_tool_name.clone(), - input: ec.to_string(), - provider_executed: Some(true), - dynamic: None, - provider_metadata: None, - })); - // provider-executed → does NOT set has_tool_calls - } - } else if let Some(cer) = part.get("codeExecutionResult") { - // Result corresponds to the most recent - // executableCode part. Gemini may emit - // several results for that one call, so - // retain the association until a new call. - if let Some(call_id) = - last_code_execution_tool_call_id.as_ref() - { - let outcome = - cer.get("outcome").cloned().unwrap_or(json!(null)); - let output = cer - .get("output") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - .unwrap_or_default(); - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: call_id.clone(), - tool_name: code_execution_tool_name.clone(), - result: json!({ "outcome": outcome, "output": output }), - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: None, - })); - } - } else if let Some(tc) = part.get("toolCall") { - // Server-side tool call (provider-executed). - let tool_type = tc - .get("toolType") - .and_then(|v| v.as_str()) - .unwrap_or(""); - let id = tc - .get("id") - .and_then(|v| v.as_str()) - .filter(|id| !id.is_empty()) - .map(std::string::ToString::to_string) - .unwrap_or_else(generate_id); - last_server_tool_call_id = Some(id.clone()); - let args = tc.get("args").cloned().unwrap_or(json!({})); - let server_meta = server_tool_metadata( - &id, - tool_type, - part.get("thoughtSignature").and_then(|v| v.as_str()), - ); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: format!("server:{tool_type}"), - input: args.to_string(), - provider_executed: Some(true), - dynamic: Some(true), - provider_metadata: Some(server_meta), - })); - // provider-executed → does NOT set has_tool_calls - } else if let Some(tr) = part.get("toolResponse") { - // Server-side tool response. - let tool_type = tr.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); - let id = last_server_tool_call_id - .take() - .or_else(|| { - tr.get("id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .filter(|id| !id.is_empty()) - .unwrap_or_else(generate_id); - let response = - tr.get("response").cloned().unwrap_or(json!({})); - let server_meta = server_tool_metadata( - &id, - tool_type, - part.get("thoughtSignature").and_then(|v| v.as_str()), - ); - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: id, - tool_name: format!("server:{tool_type}"), - result: response, - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: Some(server_meta), - })); - } else if let Some(inline) = part.get("inlineData") { - // File output — upstream :847-877. - // Close any open text/reasoning block before file. - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); - } - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); - } - if let (Some(data), Some(mime)) = ( - inline.get("data").and_then(|v| v.as_str()), - inline.get("mimeType").and_then(|v| v.as_str()), - ) { - let file = GeneratedFile { - data: GeneratedFileData::Data { data: FileBytes::Base64(data.to_string()) }, - media_type: mime.to_string(), - provider_metadata: thought_sig_meta.clone(), - }; - yield Ok(if part.get("thought").and_then(serde_json::Value::as_bool).unwrap_or(false) { StreamPart::ReasoningFile(file) } else { StreamPart::File(file) }); - } - } + controller.enqueue(StreamPart::TextDelta { + id: id.clone(), + delta: String::new(), + provider_metadata: Some(meta.clone()), + }); + } + } else { + let is_thought = part + .get("thought") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + if is_thought { + // Close any open text block before starting reasoning + if let Some(id) = self.text_id.take() { + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: None, + }); } - for part in parts { - let thought_sig_meta = part.get("thoughtSignature").and_then(Value::as_str) - .filter(|signature| !signature.is_empty()) - .map(|signature| google_metadata(json!({"thoughtSignature": signature}))); - if let Some(fc) = part.get("functionCall") { - // Single-chunk complete function call (the - // common case for the public Gemini API). - if fc.get("name").and_then(Value::as_str).is_none() || fc.get("partialArgs").is_some_and(|value| !value.is_null()) || fc.get("willContinue").and_then(Value::as_bool) == Some(true) { continue; } - let name = fc - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or(""); - let id = fc - .get("id") - .and_then(|v| v.as_str()) - .filter(|id| !id.is_empty()) - .filter(|id| !id.is_empty()) - .map(std::string::ToString::to_string) - .unwrap_or_else(generate_id); - let args = fc.get("args").filter(|value| !value.is_null()).cloned().unwrap_or(json!({})); - - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: name.to_string(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: thought_sig_meta.clone(), - }); - let args_str = args.as_str().map(str::to_owned).unwrap_or_else(|| args.to_string()); - if fc.get("args").is_some_and(|value| !value.is_null()) { - yield Ok(StreamPart::ToolInputDelta { - id: id.clone(), - delta: args_str.clone(), - provider_metadata: thought_sig_meta.clone(), - }); - } - yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: thought_sig_meta.clone()}); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: name.to_string(), - input: args_str, - provider_executed: None, - dynamic: None, - provider_metadata: thought_sig_meta.clone(), - })); - has_tool_calls = true; - } + // Start a reasoning block if not already active + if self.reasoning_id.is_none() { + let id = format!("{}", self.block_counter); + self.block_counter += 1; + self.reasoning_id = Some(id.clone()); + controller.enqueue(StreamPart::ReasoningStart { + id, + provider_metadata: thought_sig_meta.clone(), + }); + } + if let Some(id) = &self.reasoning_id { + controller.enqueue(StreamPart::ReasoningDelta { + id: id.clone(), + delta: text.to_string(), + provider_metadata: thought_sig_meta.clone(), + }); + } + } else { + // Close any open reasoning block before starting text + if let Some(id) = self.reasoning_id.take() { + controller.enqueue(StreamPart::ReasoningEnd { + id, + provider_metadata: None, + }); + } + if self.text_id.is_none() { + let id = format!("{}", self.block_counter); + self.block_counter += 1; + self.text_id = Some(id.clone()); + controller.enqueue(StreamPart::TextStart { + id, + provider_metadata: thought_sig_meta.clone(), + }); + } + if let Some(id) = &self.text_id { + controller.enqueue(StreamPart::TextDelta { + id: id.clone(), + delta: text.to_string(), + provider_metadata: thought_sig_meta.clone(), + }); } } - - if let Some(reason) = candidate.finish_reason.as_deref() { - final_finish_reason = - Some(parse_finish_reason(reason, has_tool_calls)); - } } - Ok(GoogleStreamEvent::Error(error)) => { - yield Ok(StreamPart::Error { - error: super::google_stream_error( - &error.error, - &stream_error_url, - stream_request_body.clone(), - stream_error_headers.clone(), - ), + } else if let Some(ec) = part.get("executableCode") { + // Provider-executed code execution. + let has_code = ec + .get("code") + .and_then(|v| v.as_str()) + .map(|s| !s.is_empty()) + .unwrap_or(false); + if has_code { + let id = (self.generate_id)(); + self.last_code_execution_tool_call_id = Some(id.clone()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id.clone(), + tool_name: self.code_execution_tool_name.clone(), + input: ec.to_string(), + provider_executed: Some(true), + dynamic: None, + provider_metadata: None, + })); + // provider-executed → does NOT set has_tool_calls + } + } else if let Some(cer) = part.get("codeExecutionResult") { + // Result corresponds to the most recent + // executableCode part. Gemini may emit + // several results for that one call, so + // retain the association until a new call. + if let Some(call_id) = self.last_code_execution_tool_call_id.as_ref() { + let outcome = cer.get("outcome").cloned().unwrap_or(json!(null)); + let output = cer + .get("output") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string) + .unwrap_or_default(); + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: call_id.clone(), + tool_name: self.code_execution_tool_name.clone(), + result: json!({ "outcome": outcome, "output": output }), + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: None, + })); + } + } else if let Some(tc) = part.get("toolCall") { + // Server-side tool call (provider-executed). + let tool_type = tc.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); + let id = tc + .get("id") + .and_then(|v| v.as_str()) + .filter(|id| !id.is_empty()) + .map(std::string::ToString::to_string) + .unwrap_or_else(self.generate_id); + self.last_server_tool_call_id = Some(id.clone()); + let args = tc.get("args").cloned().unwrap_or(json!({})); + let server_meta = server_tool_metadata( + &id, + tool_type, + part.get("thoughtSignature").and_then(|v| v.as_str()), + ); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id, + tool_name: format!("server:{tool_type}"), + input: args.to_string(), + provider_executed: Some(true), + dynamic: Some(true), + provider_metadata: Some(server_meta), + })); + // provider-executed → does NOT set has_tool_calls + } else if let Some(tr) = part.get("toolResponse") { + // Server-side tool response. + let tool_type = tr.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); + let id = self + .last_server_tool_call_id + .take() + .or_else(|| { + tr.get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string) + }) + .filter(|id| !id.is_empty()) + .unwrap_or_else(self.generate_id); + let response = tr.get("response").cloned().unwrap_or(json!({})); + let server_meta = server_tool_metadata( + &id, + tool_type, + part.get("thoughtSignature").and_then(|v| v.as_str()), + ); + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: id, + tool_name: format!("server:{tool_type}"), + result: response, + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: Some(server_meta), + })); + } else if let Some(inline) = part.get("inlineData") { + // File output — upstream :847-877. + // Close any open text/reasoning block before file. + if let Some(id) = self.text_id.take() { + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: None, }); - stream_errored = true; - break; } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Ok(StreamPart::Error { error }); - if !recoverable { - return; - } + if let Some(id) = self.reasoning_id.take() { + controller.enqueue(StreamPart::ReasoningEnd { + id, + provider_metadata: None, + }); + } + if let (Some(data), Some(mime)) = ( + inline.get("data").and_then(|v| v.as_str()), + inline.get("mimeType").and_then(|v| v.as_str()), + ) { + let file = GeneratedFile { + data: GeneratedFileData::Data { + data: FileBytes::Base64(data.to_string()), + }, + media_type: mime.to_string(), + provider_metadata: thought_sig_meta.clone(), + }; + controller.enqueue( + if part + .get("thought") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false) + { + StreamPart::ReasoningFile(file) + } else { + StreamPart::File(file) + }, + ); } } } + for part in parts { + let thought_sig_meta = part + .get("thoughtSignature") + .and_then(Value::as_str) + .filter(|signature| !signature.is_empty()) + .map(|signature| google_metadata(json!({"thoughtSignature": signature}))); + if let Some(fc) = part.get("functionCall") { + // Single-chunk complete function call (the + // common case for the public Gemini API). + if fc.get("name").and_then(Value::as_str).is_none() + || fc.get("partialArgs").is_some_and(|value| !value.is_null()) + || fc.get("willContinue").and_then(Value::as_bool) == Some(true) + { + continue; + } + let name = fc.get("name").and_then(|v| v.as_str()).unwrap_or(""); + let id = fc + .get("id") + .and_then(|v| v.as_str()) + .filter(|id| !id.is_empty()) + .filter(|id| !id.is_empty()) + .map(std::string::ToString::to_string) + .unwrap_or_else(self.generate_id); + let args = fc + .get("args") + .filter(|value| !value.is_null()) + .cloned() + .unwrap_or(json!({})); - // Close any remaining open text/reasoning segment. - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); + controller.enqueue(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: name.to_string(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: thought_sig_meta.clone(), + }); + let args_str = args + .as_str() + .map(str::to_owned) + .unwrap_or_else(|| args.to_string()); + if fc.get("args").is_some_and(|value| !value.is_null()) { + controller.enqueue(StreamPart::ToolInputDelta { + id: id.clone(), + delta: args_str.clone(), + provider_metadata: thought_sig_meta.clone(), + }); + } + controller.enqueue(StreamPart::ToolInputEnd { + id: id.clone(), + provider_metadata: thought_sig_meta.clone(), + }); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id, + tool_name: name.to_string(), + input: args_str, + provider_executed: None, + dynamic: None, + provider_metadata: thought_sig_meta.clone(), + })); + self.has_tool_calls = true; + } } - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None}); + } + + if let Some(reason) = candidate.finish_reason.as_deref() { + self.final_finish_reason = Some(parse_finish_reason(reason, self.has_tool_calls)); + } + } +} + +impl Transformer for GoogleLanguageModelStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + if self.stream_errored { + return; + } + let event = match event { + Ok(raw_value) => { + if self.include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: raw_value.clone(), + }); + } + serde_json::from_value::(raw_value).map_err(AiMuxError::from) } + Err(error) => Err(error), + }; - let provider_metadata = Some(google_metadata(serde_json::json!({ - "promptFeedback": last_prompt_feedback, - "groundingMetadata": last_grounding_metadata, - "urlContextMetadata": last_url_context_metadata, - "safetyRatings": last_safety_ratings, - "usageMetadata": last_usage_metadata_value, - "finishMessage": last_finish_message, - "serviceTier": last_usage_metadata_value.as_ref().and_then(|usage| usage.get("serviceTier")), - }))); - - yield Ok(StreamPart::Finish { - finish_reason: if stream_errored { - FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - } - } else { - final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }) - }, - usage: if stream_errored { Usage::default() } else { final_usage }, - provider_metadata, + match event { + Ok(GoogleStreamEvent::Chunk(chunk)) => self.transform_chunk(chunk, controller), + Ok(GoogleStreamEvent::Error(error)) => { + controller.enqueue(StreamPart::Error { + error: super::google_stream_error( + &error.error, + &self.stream_error_url, + self.stream_request_body.clone(), + self.stream_error_headers.clone(), + ), + }); + self.stream_errored = true; + } + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + controller.enqueue(StreamPart::Error { error }); + if !recoverable { + controller.terminate(); + } + } + } + } + + fn flush(mut self, controller: &mut TransformStreamController) { + // Close any remaining open text/reasoning segment. + if let Some(id) = self.text_id.take() { + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: None, }); - }; + } + if let Some(id) = self.reasoning_id.take() { + controller.enqueue(StreamPart::ReasoningEnd { + id, + provider_metadata: None, + }); + } - 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), - }), - }) + let provider_metadata = Some(google_metadata(serde_json::json!({ + "promptFeedback": self.last_prompt_feedback, + "groundingMetadata": self.last_grounding_metadata, + "urlContextMetadata": self.last_url_context_metadata, + "safetyRatings": self.last_safety_ratings, + "usageMetadata": self.last_usage_metadata_value, + "finishMessage": self.last_finish_message, + "serviceTier": self.last_usage_metadata_value.as_ref().and_then(|usage| usage.get("serviceTier")), + }))); + + controller.enqueue(StreamPart::Finish { + finish_reason: if self.stream_errored { + FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + } + } else { + self.final_finish_reason.unwrap_or(FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }) + }, + usage: if self.stream_errored { + Usage::default() + } else { + self.final_usage + }, + provider_metadata, + }); } } diff --git a/aimux-providers/src/openai_compatible/completion.rs b/aimux-providers/src/openai_compatible/completion.rs index ea5508854..f31c67a6a 100644 --- a/aimux-providers/src/openai_compatible/completion.rs +++ b/aimux-providers/src/openai_compatible/completion.rs @@ -1,5 +1,7 @@ //! Text-completion endpoint shared by the native and compatible packages. +use std::collections::HashMap; + use async_trait::async_trait; use futures::StreamExt; use serde::Deserialize; @@ -17,7 +19,7 @@ use aimux_core::types::{ FinishReason, FinishReasonUnified, InputTokenUsage, OutputTokenUsage, ResponseMetadata, Usage, Warning, }; -use aimux_provider_utils::HttpRequest; +use aimux_provider_utils::{HttpRequest, TransformStreamController, Transformer, pipe_through}; use super::config::CompatModelConfig; use super::convert::{parse_finish_reason, to_camel_case}; @@ -497,44 +499,24 @@ impl LanguageModel for OpenAICompatibleCompletionModel { } } } - let stream_body = body.clone(); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); - let mut events = futures::stream::iter(buffered).chain(events); - let mut started = false; - let mut finish_reason = FinishReason { unified: FinishReasonUnified::Other, raw: None }; - let mut final_usage = None; - let mut logprobs = None; - while let Some(event) = events.next().await { - let event = match event { - Ok(event) => event, - Err(error) => { finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; yield Ok(StreamPart::Error { error }); continue; } - }; - if raw_chunks { yield Ok(StreamPart::Raw { raw_value: event.clone() }); } - if let Some(error) = event.get("error") { - finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; - yield Ok(StreamPart::Error { error: crate::openai::openai_stream_error(error, &error_url, stream_body.clone(), response_headers.clone()) }); - continue; - } - let parsed: CompletionResponse = match parse_response(event, native, true) { - Ok(parsed) => parsed, - Err(error) => { finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; yield Ok(StreamPart::Error { error }); continue; } - }; - if !started { - started = true; - yield Ok(StreamPart::ResponseMetadata(metadata(&parsed))); - yield Ok(StreamPart::TextStart { id: "0".into(), provider_metadata: None }); - } - if parsed.usage.is_some() { final_usage = parsed.usage; } - if let Some(choice) = parsed.choices.into_iter().next() { - if let Some(reason) = choice.finish_reason { finish_reason = parse_finish_reason(&reason); } - if native && choice.logprobs.is_some() { logprobs = choice.logprobs; } - if !native || !choice.text.is_empty() { yield Ok(StreamPart::TextDelta { id: "0".into(), delta: choice.text, provider_metadata: None }); } - } - } - if started { yield Ok(StreamPart::TextEnd { id: "0".into(), provider_metadata: None }); } - yield Ok(StreamPart::Finish { finish_reason, usage: usage(final_usage, native), provider_metadata: native.then(|| provider_namespace("openai", match logprobs { Some(logprobs) => json!({ "logprobs": logprobs }), None => json!({}) }).expect("provider metadata must be an object")) }); - }; + let stream = pipe_through( + futures::stream::iter(buffered).chain(events), + OpenAICompatibleCompletionStream { + warnings, + emit_raw_chunks: raw_chunks, + native, + error_url, + stream_error_body: body.clone(), + stream_response_headers: response_headers, + started: false, + finish_reason: FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }, + usage: None, + logprobs: None, + }, + ); Ok(StreamResult { stream: Box::pin(stream), request: Some(RequestInfo { body: Some(body) }), @@ -544,3 +526,123 @@ impl LanguageModel for OpenAICompatibleCompletionModel { }) } } + +/// The `TransformStream` of `OpenAICompatibleCompletionLanguageModel.doStream` +/// (and of `OpenAICompletionLanguageModel.doStream` when native). +struct OpenAICompatibleCompletionStream { + warnings: Vec, + emit_raw_chunks: bool, + native: bool, + error_url: String, + stream_error_body: Value, + stream_response_headers: HashMap, + started: bool, + finish_reason: FinishReason, + usage: Option>, + logprobs: Option, +} + +impl OpenAICompatibleCompletionStream { + fn report_error( + &mut self, + error: AiMuxError, + controller: &mut TransformStreamController, + ) { + self.finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }; + controller.enqueue(StreamPart::Error { error }); + } +} + +impl Transformer for OpenAICompatibleCompletionStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + let event = match event { + Ok(event) => event, + Err(error) => return self.report_error(error, controller), + }; + if self.emit_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: event.clone(), + }); + } + if let Some(error) = event.get("error") { + let error = crate::openai::openai_stream_error( + error, + &self.error_url, + self.stream_error_body.clone(), + self.stream_response_headers.clone(), + ); + return self.report_error(error, controller); + } + let parsed: CompletionResponse = match parse_response(event, self.native, true) { + Ok(parsed) => parsed, + Err(error) => return self.report_error(error, controller), + }; + if !self.started { + self.started = true; + controller.enqueue(StreamPart::ResponseMetadata(metadata(&parsed))); + controller.enqueue(StreamPart::TextStart { + id: "0".into(), + provider_metadata: None, + }); + } + if parsed.usage.is_some() { + self.usage = parsed.usage; + } + if let Some(choice) = parsed.choices.into_iter().next() { + if let Some(reason) = choice.finish_reason { + self.finish_reason = parse_finish_reason(&reason); + } + if self.native && choice.logprobs.is_some() { + self.logprobs = choice.logprobs; + } + if !self.native || !choice.text.is_empty() { + controller.enqueue(StreamPart::TextDelta { + id: "0".into(), + delta: choice.text, + provider_metadata: None, + }); + } + } + } + + fn flush(self, controller: &mut TransformStreamController) { + if self.started { + controller.enqueue(StreamPart::TextEnd { + id: "0".into(), + provider_metadata: None, + }); + } + let native = self.native; + let logprobs = self.logprobs; + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason, + usage: usage(self.usage, native), + provider_metadata: native.then(|| { + provider_namespace( + "openai", + match logprobs { + Some(logprobs) => json!({ "logprobs": logprobs }), + None => json!({}), + }, + ) + .expect("provider metadata must be an object") + }), + }); + } +} diff --git a/aimux-providers/src/vertex/model.rs b/aimux-providers/src/vertex/model.rs index 50d9da4b8..d4b39f22e 100644 --- a/aimux-providers/src/vertex/model.rs +++ b/aimux-providers/src/vertex/model.rs @@ -4,10 +4,11 @@ //! [`crate::google::types`] response types. Only the endpoint construction and //! authentication differ from the public Gemini API provider. +use std::collections::{HashMap, HashSet}; + use aimux_core::tool::RawToolCall; use aimux_core::tool::ToolResult; use async_trait::async_trait; -use futures::StreamExt; use serde_json::{Value, json}; use aimux_core::error::AiMuxError; @@ -20,12 +21,13 @@ use aimux_core::stream_part::StreamPart; use aimux_core::types::{ FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, Warning, }; +use aimux_provider_utils::{TransformStreamController, Transformer, pipe_through}; use crate::google::convert::{ build_vertex_request_body, code_execution_tool_name, extract_sources, parse_finish_reason, prepare_all_tools, validate_call_options_for_namespace, }; -use crate::google::types::{Candidate, GenerateContentResponse, GoogleStreamEvent}; +use crate::google::types::{Candidate, GenerateContentResponse, GoogleStreamEvent, StreamChunk}; use crate::google::utils::{GoogleJsonAccumulator, PartialArg}; use crate::google::options::{GOOGLE, Namespace}; @@ -187,463 +189,642 @@ impl LanguageModel for VertexModel { let stream_request_body = body.clone(); let stream_error_headers = response_headers.clone(); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); + let stream = pipe_through( + sse_stream, + VertexModelStream { + warnings, + include_raw_chunks, + code_execution_tool_name, + stream_error_url, + stream_request_body, + stream_error_headers, + text_id: None, + reasoning_id: None, + block_counter: 0, + final_usage: Usage::default(), + final_finish_reason: None, + has_tool_calls: false, + active_tool_calls: Vec::new(), + response_metadata_emitted: false, + stream_errored: false, + prompt_blocked: false, + last_grounding_metadata: None, + last_url_context_metadata: None, + last_prompt_feedback: None, + last_safety_ratings: None, + last_finish_message: None, + last_usage_metadata_value: None, + emitted_source_urls: HashSet::new(), + source_id: 0, + last_code_execution_tool_call_id: None, + last_server_tool_call_id: None, + }, + ); - let mut sse_stream = sse_stream; - let mut text_id: Option = None; - let mut reasoning_id: Option = None; - let mut prompt_blocked = false; - let mut block_counter = 0usize; - let mut final_usage: Usage = Usage::default(); - let mut final_finish_reason: Option = None; - let mut has_tool_calls = false; - let mut active_tool_calls: Vec<(RawToolCall, GoogleJsonAccumulator)> = Vec::new(); - let mut response_metadata_emitted = false; - let mut stream_errored = false; + 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), + }), + }) + } +} - // Provider-metadata accumulators (mirrors TS `lastGroundingMetadata` / - // `lastUrlContextMetadata` + the finishReason-chunk snapshot). Same - // behaviour as the public Gemini API provider. - let mut last_grounding_metadata: Option = None; - let mut last_url_context_metadata: Option = None; - let mut last_prompt_feedback: Option = None; - let mut last_safety_ratings: Option = None; - let mut last_finish_message: Option = None; - let mut last_usage_metadata_value: Option = None; +/// The `TransformStream` of `GoogleLanguageModel.doStream` as configured by +/// `@ai-sdk/google-vertex`. +struct VertexModelStream { + warnings: Vec, + include_raw_chunks: bool, + code_execution_tool_name: String, + stream_error_url: String, + stream_request_body: Value, + stream_error_headers: HashMap, + text_id: Option, + reasoning_id: Option, + block_counter: usize, + final_usage: Usage, + final_finish_reason: Option, + has_tool_calls: bool, + active_tool_calls: Vec<(RawToolCall, GoogleJsonAccumulator)>, + response_metadata_emitted: bool, + stream_errored: bool, + prompt_blocked: bool, + // Provider-metadata accumulators (mirrors TS `lastGroundingMetadata` / + // `lastUrlContextMetadata` + the finishReason-chunk snapshot). + last_grounding_metadata: Option, + last_url_context_metadata: Option, + last_prompt_feedback: Option, + last_safety_ratings: Option, + last_finish_message: Option, + last_usage_metadata_value: Option, + // Source dedup across chunks (url sources only). + emitted_source_urls: HashSet, + source_id: usize, + // Associates code-execution results / server tool responses with their + // preceding call. + last_code_execution_tool_call_id: Option, + last_server_tool_call_id: Option, +} - // Source dedup across chunks (url sources only). - let mut emitted_source_urls: std::collections::HashSet = - std::collections::HashSet::new(); - let mut source_id = 0usize; +impl VertexModelStream { + fn transform_chunk( + &mut self, + chunk: Box, + controller: &mut TransformStreamController, + ) { + if !self.response_metadata_emitted + && let Some(id) = &chunk.response_id + { + self.response_metadata_emitted = true; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: Some(id.clone()), + timestamp: None, + model_id: None, + })); + } - // Associates code-execution results / server tool responses with - // their preceding call. - let mut last_code_execution_tool_call_id: Option = None; - let mut last_server_tool_call_id: Option = None; + if !self.prompt_blocked + && let Some(pf) = &chunk.prompt_feedback + { + self.last_prompt_feedback = Some(pf.clone()); + if let Some(reason) = confirmed_block_reason(Some(pf)) { + self.prompt_blocked = true; + self.final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::ContentFilter, + raw: Some(reason), + }); + } + } - while let Some(event) = sse_stream.next().await { - if stream_errored { - break; - } - let event = match event { - Ok(raw) => { - if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: raw.clone() }); } - if let Some(usage) = raw.get("usageMetadata") { - final_usage = vertex_usage(Some(usage)); - last_usage_metadata_value = Some(usage.clone()); - } - serde_json::from_value::(raw) - .map_err(|error| AiMuxError::InvalidResponseData(error.to_string())) - } - Err(error) => Err(error), - }; - match event { - Ok(GoogleStreamEvent::Chunk(chunk)) => { + let Some(candidates) = chunk.candidates else { + return; + }; + let Some(candidate) = candidates.into_iter().next() else { + return; + }; - if !response_metadata_emitted - && let Some(id) = &chunk.response_id { - response_metadata_emitted = true; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: Some(id.clone()), - timestamp: None, - model_id: None, - })); - } + if let Some(gm) = &candidate.grounding_metadata { + self.last_grounding_metadata = Some(gm.clone()); + } + if let Some(ucm) = &candidate.url_context_metadata { + self.last_url_context_metadata = Some(ucm.clone()); + } - if !prompt_blocked && let Some(pf) = &chunk.prompt_feedback { - last_prompt_feedback = Some(pf.clone()); - if let Some(reason) = confirmed_block_reason(Some(pf)) { - prompt_blocked = true; - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::ContentFilter, raw: Some(reason) }); - } - } + if let Some(sr) = &candidate.safety_ratings { + self.last_safety_ratings = serde_json::to_value(sr).ok(); + } + if let Some(fm) = &candidate.finish_message { + self.last_finish_message = Some(json!(fm)); + } - let Some(candidates) = chunk.candidates else { - continue; - }; - let Some(candidate) = candidates.into_iter().next() else { - continue; - }; + if self.prompt_blocked { + return; + } - if let Some(gm) = &candidate.grounding_metadata { - last_grounding_metadata = Some(gm.clone()); - } - if let Some(ucm) = &candidate.url_context_metadata { - last_url_context_metadata = Some(ucm.clone()); - } + // Extract url sources from this chunk's grounding metadata + // (deduplicated across chunks; document sources are not + // emitted in the stream, matching TS). + let chunk_sources = + extract_sources(candidate.grounding_metadata.as_ref(), &mut self.source_id); + for src in chunk_sources { + if let GenerateContent::Source(Source::Url { + url, + id, + title, + provider_metadata: None, + }) = src + && self.emitted_source_urls.insert(url.clone()) + { + controller.enqueue(StreamPart::Source(Source::Url { + id, + url, + title, + provider_metadata: None, + })); + } + } - if let Some(sr) = &candidate.safety_ratings { - last_safety_ratings = serde_json::to_value(sr).ok(); + if let Some(parts) = candidate.content.as_ref().and_then(|c| c.parts.as_ref()) { + for part in parts + .iter() + .filter(|part| part.get("functionCall").is_none()) + .chain( + parts + .iter() + .filter(|part| part.get("functionCall").is_some()), + ) + { + if let Some(text) = part.get("text").and_then(Value::as_str) { + let metadata = thought_metadata(part); + if text.is_empty() { + if metadata.is_some() + && let Some(id) = &self.text_id + { + controller.enqueue(StreamPart::TextDelta { + id: id.clone(), + delta: String::new(), + provider_metadata: metadata, + }); } - if let Some(fm) = &candidate.finish_message { - last_finish_message = Some(json!(fm)); + } else if part.get("thought") == Some(&Value::Bool(true)) { + if let Some(id) = self.text_id.take() { + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: None, + }); } - - if prompt_blocked { continue; } - - // Extract url sources from this chunk's grounding metadata - // (deduplicated across chunks; document sources are not - // emitted in the stream, matching TS). - let chunk_sources = - extract_sources(candidate.grounding_metadata.as_ref(), &mut source_id); - for src in chunk_sources { - if let GenerateContent::Source(Source::Url { - url, + if self.reasoning_id.is_none() { + let id = format!("{}", self.block_counter); + self.block_counter += 1; + self.reasoning_id = Some(id.clone()); + controller.enqueue(StreamPart::ReasoningStart { + id, + provider_metadata: metadata.clone(), + }); + } + controller.enqueue(StreamPart::ReasoningDelta { + id: self.reasoning_id.clone().unwrap(), + delta: text.to_string(), + provider_metadata: metadata, + }); + } else { + if let Some(id) = self.reasoning_id.take() { + controller.enqueue(StreamPart::ReasoningEnd { id, - title, provider_metadata: None, - }) = src - && emitted_source_urls.insert(url.clone()) { - yield Ok(StreamPart::Source(Source::Url { - id, - url, - title, - provider_metadata: None, - })); - } + }); } - - if let Some(parts) = - candidate.content.as_ref().and_then(|c| c.parts.as_ref()) + if self.text_id.is_none() { + let id = format!("{}", self.block_counter); + self.block_counter += 1; + self.text_id = Some(id.clone()); + controller.enqueue(StreamPart::TextStart { + id, + provider_metadata: metadata.clone(), + }); + } + controller.enqueue(StreamPart::TextDelta { + id: self.text_id.clone().unwrap(), + delta: text.to_string(), + provider_metadata: metadata, + }); + } + } else if let Some(file) = inline_file(part) { + if let Some(id) = self.text_id.take() { + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: None, + }); + } + if let Some(id) = self.reasoning_id.take() { + controller.enqueue(StreamPart::ReasoningEnd { + id, + provider_metadata: None, + }); + } + controller.enqueue(if part.get("thought") == Some(&Value::Bool(true)) { + StreamPart::ReasoningFile(file) + } else { + StreamPart::File(file) + }); + } else if let Some(fc) = part.get("functionCall") { + let name = fc.get("name").and_then(Value::as_str); + let args = fc.get("args").filter(|value| !value.is_null()); + let partial_args = fc.get("partialArgs").filter(|value| !value.is_null()); + let will_continue = + fc.get("willContinue").and_then(Value::as_bool) == Some(true); + let thought_signature = part + .get("thoughtSignature") + .and_then(Value::as_str) + .filter(|signature| !signature.is_empty()) + .map(str::to_owned); + let tool_metadata = thought_signature.as_deref().map(|signature| { + vertex_provider_metadata(json!({ "thoughtSignature": signature })) + }); + if partial_args.is_some() || (name.is_some() && will_continue) { + if let Some(name) = name { + let id = fc + .get("id") + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .map(str::to_owned) + .unwrap_or_else(aimux_provider_utils::generate_id); + controller.enqueue(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: name.to_owned(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: tool_metadata.clone(), + }); + self.active_tool_calls.push(( + RawToolCall { + tool_call_id: id, + tool_name: name.to_owned(), + input: String::new(), + provider_executed: None, + dynamic: None, + provider_metadata: tool_metadata.clone(), + }, + GoogleJsonAccumulator::new(), + )); + } + if let Some(partial_args) = partial_args.and_then(Value::as_array) + && let Some((call, accumulator)) = self.active_tool_calls.last_mut() { - for part in parts.iter().filter(|part| part.get("functionCall").is_none()) - .chain(parts.iter().filter(|part| part.get("functionCall").is_some())) { - if let Some(text) = part.get("text").and_then(Value::as_str) { - let metadata = thought_metadata(part); - if text.is_empty() { - if metadata.is_some() && let Some(id) = &text_id { - yield Ok(StreamPart::TextDelta { id: id.clone(), delta: String::new(), provider_metadata: metadata }); - } - } else if part.get("thought") == Some(&Value::Bool(true)) { - if let Some(id) = text_id.take() { yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); } - if reasoning_id.is_none() { - let id = format!("{block_counter}"); block_counter += 1; - reasoning_id = Some(id.clone()); - yield Ok(StreamPart::ReasoningStart { id, provider_metadata: metadata.clone() }); - } - yield Ok(StreamPart::ReasoningDelta { id: reasoning_id.clone().unwrap(), delta: text.to_string(), provider_metadata: metadata }); - } else { - if let Some(id) = reasoning_id.take() { yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); } - if text_id.is_none() { - let id = format!("{block_counter}"); block_counter += 1; - text_id = Some(id.clone()); - yield Ok(StreamPart::TextStart { id, provider_metadata: metadata.clone() }); - } - yield Ok(StreamPart::TextDelta { id: text_id.clone().unwrap(), delta: text.to_string(), provider_metadata: metadata }); - } - } else if let Some(file) = inline_file(part) { - if let Some(id) = text_id.take() { yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); } - if let Some(id) = reasoning_id.take() { yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); } - yield Ok(if part.get("thought") == Some(&Value::Bool(true)) { StreamPart::ReasoningFile(file) } else { StreamPart::File(file) }); - } else if let Some(fc) = part.get("functionCall") { - let name = fc.get("name").and_then(Value::as_str); - let args = fc.get("args").filter(|value| !value.is_null()); - let partial_args = fc.get("partialArgs").filter(|value| !value.is_null()); - let will_continue = fc.get("willContinue").and_then(Value::as_bool) == Some(true); - let thought_signature = part.get("thoughtSignature") - .and_then(Value::as_str).filter(|signature| !signature.is_empty()) - .map(str::to_owned); - let tool_metadata = thought_signature.as_deref().map(|signature| { - vertex_provider_metadata(json!({ "thoughtSignature": signature })) - }); - if partial_args.is_some() || (name.is_some() && will_continue) { - if let Some(name) = name { - let id = fc.get("id").and_then(Value::as_str) - .filter(|id| !id.is_empty()).map(str::to_owned) - .unwrap_or_else(aimux_provider_utils::generate_id); - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), tool_name: name.to_owned(), - provider_executed: None, dynamic: None, title: None, - provider_metadata: tool_metadata.clone(), - }); - active_tool_calls.push((RawToolCall { - tool_call_id: id, tool_name: name.to_owned(), input: String::new(), - provider_executed: None, dynamic: None, - provider_metadata: tool_metadata.clone(), - }, GoogleJsonAccumulator::new())); - } - if let Some(partial_args) = partial_args.and_then(Value::as_array) - && let Some((call, accumulator)) = active_tool_calls.last_mut() { - let partial_args: Vec = partial_args.iter().map(|arg| PartialArg { - json_path: arg.get("jsonPath").and_then(Value::as_str).unwrap_or_default().to_owned(), - string_value: arg.get("stringValue").and_then(Value::as_str).map(str::to_owned), - number_value: arg.get("numberValue").and_then(Value::as_f64), - bool_value: arg.get("boolValue").and_then(Value::as_bool), - null_value: arg.get("nullValue").map(|_| ()), - will_continue: arg.get("willContinue").and_then(Value::as_bool), - }).collect(); - match accumulator.process_partial_args(&partial_args) { - Ok(result) => { - if !result.text_delta.is_empty() { - yield Ok(StreamPart::ToolInputDelta { - id: call.tool_call_id.clone(), delta: result.text_delta, - provider_metadata: tool_metadata, - }); - } - } - Err(error) => { yield Ok(StreamPart::Error { error }); return; } - } - if !will_continue && partial_args.iter().all(|arg| arg.will_continue != Some(true)) { - for event in finish_streaming_tool_call(&mut active_tool_calls) { yield Ok(event); } - has_tool_calls = true; - } - } - continue; - } - if name.is_none() && args.is_none() && partial_args.is_none() - && fc.get("willContinue").is_none_or(Value::is_null) { - if !active_tool_calls.is_empty() { - for event in finish_streaming_tool_call(&mut active_tool_calls) { yield Ok(event); } - has_tool_calls = true; - } - continue; - } - let Some(name) = name else { continue; }; - let id = fc.get("id").and_then(Value::as_str) - .filter(|id| !id.is_empty()).map(str::to_owned) - .unwrap_or_else(aimux_provider_utils::generate_id); - let args = args.cloned().unwrap_or(json!({})); - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: name.to_string(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: tool_metadata.clone(), - }); - let args_str = args.as_str().map(str::to_string).unwrap_or_else(|| args.to_string()); - if fc.get("args").is_some_and(|value| !value.is_null()) { yield Ok(StreamPart::ToolInputDelta { - id: id.clone(), - delta: args_str.clone(), - provider_metadata: tool_metadata.clone(), - }); - } - yield Ok(StreamPart::ToolInputEnd { - id: id.clone(), - provider_metadata: tool_metadata.clone(), - }); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: name.to_string(), - input: args_str, - provider_executed: None, - dynamic: None, - provider_metadata: tool_metadata, - })); - has_tool_calls = true; - } else if let Some(inline) = part.get("inlineData") { - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); - } - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); - } - if let (Some(data), Some(media_type)) = ( - inline.get("data").and_then(Value::as_str), - inline.get("mimeType").and_then(Value::as_str), - ) { - let file = GeneratedFile { - data: GeneratedFileData::Data { data: FileBytes::Base64(data.to_string()) }, - media_type: media_type.to_string(), - provider_metadata: thought_metadata(part), - }; - yield Ok(if part.get("thought").and_then(Value::as_bool) == Some(true) { - StreamPart::ReasoningFile(file) - } else { - StreamPart::File(file) + let partial_args: Vec = partial_args + .iter() + .map(|arg| PartialArg { + json_path: arg + .get("jsonPath") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + string_value: arg + .get("stringValue") + .and_then(Value::as_str) + .map(str::to_owned), + number_value: arg.get("numberValue").and_then(Value::as_f64), + bool_value: arg.get("boolValue").and_then(Value::as_bool), + null_value: arg.get("nullValue").map(|_| ()), + will_continue: arg.get("willContinue").and_then(Value::as_bool), + }) + .collect(); + match accumulator.process_partial_args(&partial_args) { + Ok(result) => { + if !result.text_delta.is_empty() { + controller.enqueue(StreamPart::ToolInputDelta { + id: call.tool_call_id.clone(), + delta: result.text_delta, + provider_metadata: tool_metadata, }); } - } else if let Some(ec) = part.get("executableCode") { - // Provider-executed code execution. - let has_code = ec - .get("code") - .and_then(|v| v.as_str()) - .map(|s| !s.is_empty()) - .unwrap_or(false); - if has_code { - let id = format!("call-{block_counter}"); - block_counter += 1; - last_code_execution_tool_call_id = Some(id.clone()); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id.clone(), - tool_name: code_execution_tool_name.clone(), - input: ec.to_string(), - provider_executed: Some(true), - dynamic: None, - provider_metadata: Some(vertex_server_tool_metadata( - &id, - "code_execution", - None, - )), - })); - // provider-executed → does NOT set has_tool_calls - } - } else if let Some(cer) = part.get("codeExecutionResult") { - // Result corresponds to the most recent - // executableCode part. Gemini may emit - // several results for that one call, so - // retain the association until a new call. - if let Some(call_id) = - last_code_execution_tool_call_id.as_ref() - { - let outcome = - cer.get("outcome").cloned().unwrap_or(json!(null)); - let output = cer - .get("output") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - .unwrap_or_default(); - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: call_id.clone(), - tool_name: code_execution_tool_name.clone(), - result: json!({ "outcome": outcome, "output": output }), - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: Some(vertex_server_tool_metadata( - call_id, - "code_execution", - None, - )), - })); - } - } else if let Some(tc) = part.get("toolCall") { - // Server-side tool call (provider-executed). - let tool_type = tc - .get("toolType") - .and_then(|v| v.as_str()) - .unwrap_or(""); - let id = tc - .get("id") - .and_then(|v| v.as_str()) - .filter(|id| !id.is_empty()) - .map(std::string::ToString::to_string) - .unwrap_or_else(aimux_provider_utils::generate_id); - block_counter += 1; - last_server_tool_call_id = Some(id.clone()); - let args = tc.get("args").cloned().unwrap_or(json!({})); - let thought_signature = part - .get("thoughtSignature") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let server_meta = vertex_server_tool_metadata( - &id, - tool_type, - thought_signature.as_deref(), - ); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: format!("server:{tool_type}"), - input: args.to_string(), - provider_executed: Some(true), - dynamic: Some(true), - provider_metadata: Some(server_meta), - })); - // provider-executed → does NOT set has_tool_calls - } else if let Some(tr) = part.get("toolResponse") { - // Server-side tool response. - let tool_type = tr - .get("toolType") - .and_then(|v| v.as_str()) - .unwrap_or(""); - let id = last_server_tool_call_id - .take() - .or_else(|| { - tr.get("id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .unwrap_or_else(|| format!("call-{block_counter}")); - block_counter += 1; - let response = - tr.get("response").cloned().unwrap_or(json!({})); - let server_meta = vertex_server_tool_metadata( - &id, - tool_type, - part.get("thoughtSignature").and_then(|v| v.as_str()), - ); - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: id, - tool_name: format!("server:{tool_type}"), - result: response, - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: Some(server_meta), - })); + } + Err(error) => { + controller.enqueue(StreamPart::Error { error }); + controller.terminate(); + return; } } + if !will_continue + && partial_args + .iter() + .all(|arg| arg.will_continue != Some(true)) + { + for event in finish_streaming_tool_call(&mut self.active_tool_calls) + { + controller.enqueue(event); + } + self.has_tool_calls = true; + } } - - if let Some(reason) = candidate.finish_reason.as_deref() { - final_finish_reason = - Some(parse_finish_reason(reason, has_tool_calls)); + continue; + } + if name.is_none() + && args.is_none() + && partial_args.is_none() + && fc.get("willContinue").is_none_or(Value::is_null) + { + if !self.active_tool_calls.is_empty() { + for event in finish_streaming_tool_call(&mut self.active_tool_calls) { + controller.enqueue(event); + } + self.has_tool_calls = true; } + continue; } - Ok(GoogleStreamEvent::Error(error)) => { - yield Ok(StreamPart::Error { - error: crate::google::google_stream_error( - &error.error, - &stream_error_url, - stream_request_body.clone(), - stream_error_headers.clone(), - ), + let Some(name) = name else { + continue; + }; + let id = fc + .get("id") + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .map(str::to_owned) + .unwrap_or_else(aimux_provider_utils::generate_id); + let args = args.cloned().unwrap_or(json!({})); + controller.enqueue(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: name.to_string(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: tool_metadata.clone(), + }); + let args_str = args + .as_str() + .map(str::to_string) + .unwrap_or_else(|| args.to_string()); + if fc.get("args").is_some_and(|value| !value.is_null()) { + controller.enqueue(StreamPart::ToolInputDelta { + id: id.clone(), + delta: args_str.clone(), + provider_metadata: tool_metadata.clone(), }); - stream_errored = true; - break; } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Ok(StreamPart::Error { error }); - if !recoverable { - return; - } + controller.enqueue(StreamPart::ToolInputEnd { + id: id.clone(), + provider_metadata: tool_metadata.clone(), + }); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id, + tool_name: name.to_string(), + input: args_str, + provider_executed: None, + dynamic: None, + provider_metadata: tool_metadata, + })); + self.has_tool_calls = true; + } else if let Some(inline) = part.get("inlineData") { + if let Some(id) = self.text_id.take() { + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: None, + }); } + if let Some(id) = self.reasoning_id.take() { + controller.enqueue(StreamPart::ReasoningEnd { + id, + provider_metadata: None, + }); + } + if let (Some(data), Some(media_type)) = ( + inline.get("data").and_then(Value::as_str), + inline.get("mimeType").and_then(Value::as_str), + ) { + let file = GeneratedFile { + data: GeneratedFileData::Data { + data: FileBytes::Base64(data.to_string()), + }, + media_type: media_type.to_string(), + provider_metadata: thought_metadata(part), + }; + controller.enqueue( + if part.get("thought").and_then(Value::as_bool) == Some(true) { + StreamPart::ReasoningFile(file) + } else { + StreamPart::File(file) + }, + ); + } + } else if let Some(ec) = part.get("executableCode") { + // Provider-executed code execution. + let has_code = ec + .get("code") + .and_then(|v| v.as_str()) + .map(|s| !s.is_empty()) + .unwrap_or(false); + if has_code { + let id = format!("call-{}", self.block_counter); + self.block_counter += 1; + self.last_code_execution_tool_call_id = Some(id.clone()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id.clone(), + tool_name: self.code_execution_tool_name.clone(), + input: ec.to_string(), + provider_executed: Some(true), + dynamic: None, + provider_metadata: Some(vertex_server_tool_metadata( + &id, + "code_execution", + None, + )), + })); + // provider-executed → does NOT set has_tool_calls + } + } else if let Some(cer) = part.get("codeExecutionResult") { + // Result corresponds to the most recent + // executableCode part. Gemini may emit + // several results for that one call, so + // retain the association until a new call. + if let Some(call_id) = self.last_code_execution_tool_call_id.as_ref() { + let outcome = cer.get("outcome").cloned().unwrap_or(json!(null)); + let output = cer + .get("output") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string) + .unwrap_or_default(); + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: call_id.clone(), + tool_name: self.code_execution_tool_name.clone(), + result: json!({ "outcome": outcome, "output": output }), + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: Some(vertex_server_tool_metadata( + call_id, + "code_execution", + None, + )), + })); + } + } else if let Some(tc) = part.get("toolCall") { + // Server-side tool call (provider-executed). + let tool_type = tc.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); + let id = tc + .get("id") + .and_then(|v| v.as_str()) + .filter(|id| !id.is_empty()) + .map(std::string::ToString::to_string) + .unwrap_or_else(aimux_provider_utils::generate_id); + self.block_counter += 1; + self.last_server_tool_call_id = Some(id.clone()); + let args = tc.get("args").cloned().unwrap_or(json!({})); + let thought_signature = part + .get("thoughtSignature") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let server_meta = + vertex_server_tool_metadata(&id, tool_type, thought_signature.as_deref()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id, + tool_name: format!("server:{tool_type}"), + input: args.to_string(), + provider_executed: Some(true), + dynamic: Some(true), + provider_metadata: Some(server_meta), + })); + // provider-executed → does NOT set has_tool_calls + } else if let Some(tr) = part.get("toolResponse") { + // Server-side tool response. + let tool_type = tr.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); + let id = self + .last_server_tool_call_id + .take() + .or_else(|| { + tr.get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string) + }) + .unwrap_or_else(|| format!("call-{}", self.block_counter)); + self.block_counter += 1; + let response = tr.get("response").cloned().unwrap_or(json!({})); + let server_meta = vertex_server_tool_metadata( + &id, + tool_type, + part.get("thoughtSignature").and_then(|v| v.as_str()), + ); + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: id, + tool_name: format!("server:{tool_type}"), + result: response, + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: Some(server_meta), + })); } } + } + + if let Some(reason) = candidate.finish_reason.as_deref() { + self.final_finish_reason = Some(parse_finish_reason(reason, self.has_tool_calls)); + } + } +} + +impl Transformer for VertexModelStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + if self.stream_errored { + return; + } + let event = match event { + Ok(raw) => { + if self.include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: raw.clone(), + }); + } + if let Some(usage) = raw.get("usageMetadata") { + self.final_usage = vertex_usage(Some(usage)); + self.last_usage_metadata_value = Some(usage.clone()); + } + serde_json::from_value::(raw) + .map_err(|error| AiMuxError::InvalidResponseData(error.to_string())) } + Err(error) => Err(error), + }; - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); + match event { + Ok(GoogleStreamEvent::Chunk(chunk)) => self.transform_chunk(chunk, controller), + Ok(GoogleStreamEvent::Error(error)) => { + controller.enqueue(StreamPart::Error { + error: crate::google::google_stream_error( + &error.error, + &self.stream_error_url, + self.stream_request_body.clone(), + self.stream_error_headers.clone(), + ), + }); + self.stream_errored = true; } - let provider_metadata = Some(vertex_provider_metadata(json!({ - "promptFeedback": last_prompt_feedback, - "groundingMetadata": last_grounding_metadata, - "urlContextMetadata": last_url_context_metadata, - "safetyRatings": last_safety_ratings, - "serviceTier": last_usage_metadata_value.as_ref().and_then(|usage| usage.get("serviceTier")), - "usageMetadata": last_usage_metadata_value, - "finishMessage": last_finish_message, - }))); + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + controller.enqueue(StreamPart::Error { error }); + if !recoverable { + controller.terminate(); + } + } + } + } - yield Ok(StreamPart::Finish { - finish_reason: if stream_errored { - FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - } - } else { - final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }) - }, - usage: if stream_errored { Usage::default() } else { final_usage }, - provider_metadata, + fn flush(mut self, controller: &mut TransformStreamController) { + if let Some(id) = self.text_id.take() { + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: None, }); - }; + } - 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), - }), - }) + if let Some(id) = self.reasoning_id.take() { + controller.enqueue(StreamPart::ReasoningEnd { + id, + provider_metadata: None, + }); + } + let provider_metadata = Some(vertex_provider_metadata(json!({ + "promptFeedback": self.last_prompt_feedback, + "groundingMetadata": self.last_grounding_metadata, + "urlContextMetadata": self.last_url_context_metadata, + "safetyRatings": self.last_safety_ratings, + "serviceTier": self.last_usage_metadata_value.as_ref().and_then(|usage| usage.get("serviceTier")), + "usageMetadata": self.last_usage_metadata_value, + "finishMessage": self.last_finish_message, + }))); + + controller.enqueue(StreamPart::Finish { + finish_reason: if self.stream_errored { + FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + } + } else { + self.final_finish_reason.unwrap_or(FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }) + }, + usage: if self.stream_errored { + Usage::default() + } else { + self.final_usage + }, + provider_metadata, + }); } } From bcde50dba0e19f266accbe22ae988b4e9d675ce5 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 13:11:13 +0800 Subject: [PATCH 13/16] refactor(providers): Anthropic, Bedrock and Cohere streams as Transformer structs Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-providers/src/anthropic/stream.rs | 1090 +++++++++++++---------- aimux-providers/src/bedrock/model.rs | 757 ++++++++++------ aimux-providers/src/cohere/model.rs | 340 +++++-- 3 files changed, 1360 insertions(+), 827 deletions(-) diff --git a/aimux-providers/src/anthropic/stream.rs b/aimux-providers/src/anthropic/stream.rs index 1a5b0e063..025c6e192 100644 --- a/aimux-providers/src/anthropic/stream.rs +++ b/aimux-providers/src/anthropic/stream.rs @@ -24,14 +24,17 @@ use aimux_core::shared::provider_namespace; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; use aimux_core::types::{ProviderMetadata, Warning}; -use aimux_provider_utils::HttpRequest; +use aimux_provider_utils::{HttpRequest, TransformStreamController, Transformer, pipe_through}; use serde_json::{Value, json}; use super::config::AnthropicModelConfig; use super::convert::parse_stop_reason; use super::options::CANONICAL; use super::tool_name_mapping::ToolNameMapping; -use super::types::{AnthropicResponse, ContentBlock, StreamErrorData, StreamEvent, ToolCallCaller}; +use super::types::{ + AnthropicResponse, AnthropicUsage, ContentBlock, DeltaBlock, MessageDeltaBody, + MessageStartData, StreamErrorData, StreamEvent, ToolCallCaller, +}; pub(crate) fn anthropic_stream_error( error: &StreamErrorData, @@ -1092,7 +1095,7 @@ pub(crate) async fn anthropic_stream_core( tool_names: ToolNameMapping, used_custom_options_key: bool, uses_json_response_tool: bool, - mut citation_documents: Vec, + citation_documents: Vec, generate_id: aimux_provider_utils::IdGenerator, ) -> Result { let endpoint = request.url.clone(); @@ -1119,471 +1122,664 @@ pub(crate) async fn anthropic_stream_core( response_headers, )); } - let stream_error_url = endpoint; - let stream_request_body = body.clone(); - let stream_response_headers = response_headers.clone(); - let stream = async_stream::stream! { - // First part: StreamStart. - yield Ok(StreamPart::StreamStart { warnings }); + let raw_usage = Value::Object(serde_json::Map::new()); + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + AnthropicMessagesStream { + warnings, + stream_error_url: endpoint, + stream_request_body: body.clone(), + stream_response_headers: response_headers.clone(), + options_name, + tool_names, + used_custom_options_key, + uses_json_response_tool, + citation_documents, + generate_id, + blocks: HashMap::new(), + is_json_response_from_tool: false, + final_usage: Usage::default(), + final_finish_reason: None, + effective_usage: raw_usage.clone(), + raw_usage, + stop_sequence: None, + container: None, + context_management: None, + stop_details: None, + input_transformations: None, + safeguard_results: None, + response_meta_emitted: false, + active_message_id: None, + message_stopped: false, + mcp_tool_calls: HashMap::new(), + server_tool_calls: HashMap::new(), + }, + ); - let mut sse = futures::stream::iter(first_event.into_iter()).chain(sse_stream); - let mut blocks: HashMap = HashMap::new(); - let mut is_json_response_from_tool = false; - let mut final_usage = Usage::default(); - let mut final_finish_reason: Option = None; - // Result-level providerMetadata: the raw usage (`message_start`'s, - // updated by every `message_delta`'s), the stop sequence, the - // container and the context-management edits. - let mut raw_usage = Value::Object(serde_json::Map::new()); - let mut effective_usage = raw_usage.clone(); - let mut stop_sequence: Option = None; - let mut container: Option = None; - let mut context_management: Option = None; - let mut stop_details: Option = None; - let mut input_transformations: Option = None; - let mut safeguard_results: Option = None; - let mut response_meta_emitted = false; - let mut active_message_id: Option = None; - let mut message_stopped = false; + 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), + }), + }) +} - // id → (tool name, server name), so `mcp_tool_result` can inherit them - // from the `mcp_tool_use` it answers. - let mut mcp_tool_calls: HashMap = HashMap::new(); - // tool_use_id → provider tool name. Both tool-search variants share - // one result block type, so the id is required to disambiguate aliases. - let mut server_tool_calls: HashMap = HashMap::new(); +/// The `TransformStream` of `AnthropicMessagesLanguageModel.doStream`. +struct AnthropicMessagesStream { + warnings: Vec, + stream_error_url: String, + stream_request_body: Value, + stream_response_headers: HashMap, + options_name: String, + tool_names: ToolNameMapping, + used_custom_options_key: bool, + uses_json_response_tool: bool, + citation_documents: Vec, + generate_id: aimux_provider_utils::IdGenerator, + blocks: HashMap, + is_json_response_from_tool: bool, + final_usage: Usage, + final_finish_reason: Option, + /// Result-level providerMetadata: the raw usage (`message_start`'s, + /// updated by every `message_delta`'s), the stop sequence, the + /// container and the context-management edits. + raw_usage: Value, + effective_usage: Value, + stop_sequence: Option, + container: Option, + context_management: Option, + stop_details: Option, + input_transformations: Option, + safeguard_results: Option, + response_meta_emitted: bool, + active_message_id: Option, + /// Set by `message_stop`; later events are ignored and only then does + /// `flush` emit the finish part. + message_stopped: bool, + /// id → (tool name, server name), so `mcp_tool_result` can inherit them + /// from the `mcp_tool_use` it answers. + mcp_tool_calls: HashMap, + /// tool_use_id → provider tool name. Both tool-search variants share + /// one result block type, so the id is required to disambiguate aliases. + server_tool_calls: HashMap, +} - while let Some(event) = sse.next().await { - match event { - Ok(stream_event) => { - match stream_event { - StreamEvent::MessageStart { message } => { - if let Some(active_id) = &active_message_id { - if active_id == &message.id { - continue; - } - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(format!( - "Received message_start for message {:?} while message {:?} is still open.", - message.id, active_id, - )), - }); - return; - } - active_message_id = Some(message.id.clone()); - input_transformations = message.input_transformations.or(input_transformations); - container = message.container.or(container); - if let Some(reason) = message.stop_reason.as_deref() { - final_finish_reason = Some(parse_stop_reason(reason)); - } - if let Some(usage) = &message.usage { - // RFC-0015 P0-2: full input side incl. cache - // fields + raw (Anthropic reports cache only - // in message_start). - final_usage = super::usage::usage_from_anthropic(usage); - raw_usage = serde_json::to_value(usage).unwrap_or(raw_usage); - effective_usage = raw_usage.clone(); - } - if !response_meta_emitted { - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: Some(message.id.clone()), - timestamp: None, - model_id: Some(message.model.clone()), - })); - response_meta_emitted = true; - } - } - StreamEvent::ContentBlockStart { index, content_block } => { - match content_block { - ContentBlock::Text { .. } if uses_json_response_tool => {} - ContentBlock::Text { .. } => { - blocks.insert(index, BlockState::Text { citations: Vec::new() }); - yield Ok(StreamPart::TextStart { id: index.to_string(), provider_metadata: None }); - } - ContentBlock::Compaction { content, signature } => { - blocks.insert(index, BlockState::Text { citations: Vec::new() }); - yield Ok(StreamPart::TextStart { id: index.to_string(), provider_metadata: Some(compaction_metadata(signature.as_deref())) }); - if signature.is_some() - && let Some(text) = content.filter(|text| !text.is_empty()) { - yield Ok(StreamPart::TextDelta { id: index.to_string(), delta: text, provider_metadata: None }); - } - } - ContentBlock::Thinking { .. } => { - yield Ok(StreamPart::ReasoningStart { id: index.to_string(), provider_metadata: None }); - blocks.insert( - index, - BlockState::Thinking, - ); - } - ContentBlock::ToolUse { name, .. } if uses_json_response_tool && name == "json" => { - is_json_response_from_tool = true; - blocks.insert(index, BlockState::Text { citations: Vec::new() }); - yield Ok(StreamPart::TextStart { id: index.to_string(), provider_metadata: None }); - } - ContentBlock::ToolUse { - id, - name, - input, - toolset_name, - caller, - } => { - let custom_name = toolset_name.as_deref().map_or_else( - || name.clone(), - |name| tool_names.to_custom_tool_name(name).to_string(), - ); - let initial_input = initial_tool_input(&input); - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: custom_name.clone(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - blocks.insert(index, BlockState::ToolUse { - id, - name: custom_name, - first_delta: initial_input.is_empty(), - accumulated_json: initial_input, - provider_executed: None, - dynamic: None, - provider_tool_name: None, - provider_tool_input_type: None, - provider_metadata: tool_call_metadata(caller.as_ref(), toolset_name.as_deref(), CANONICAL), - toolset_member_name: toolset_name.map(|_| name), - }); - } - // Server-side tool use follows the same input - // lifecycle as client tools because some code - // execution inputs arrive entirely via deltas. - ContentBlock::ServerToolUse { id, name, input, caller } => { - if is_tool_search_provider_name(&name) { - server_tool_calls.insert(id.clone(), name.clone()); - } - let provider_name = server_tool_provider_name(&name); - let custom_name = tool_names - .to_custom_tool_name(provider_name) - .to_string(); - let dynamic = (provider_name == "code_execution" - && tool_names.mark_code_execution_dynamic()) - .then_some(true); - let initial_input = initial_tool_input(&input); - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: custom_name.clone(), - provider_executed: Some(true), - dynamic, - title: None, - provider_metadata: None, - }); - blocks.insert(index, BlockState::ToolUse { - id, - name: custom_name, - first_delta: initial_input.is_empty(), - accumulated_json: initial_input, - provider_executed: Some(true), - dynamic, - provider_tool_name: Some(provider_name.to_string()), - provider_tool_input_type: match name.as_str() { - "text_editor_code_execution" | "bash_code_execution" => { - Some(name) - } - "code_execution" => { - Some("programmatic-tool-call".to_string()) - } - _ => None, - }, - provider_metadata: tool_call_caller_metadata(caller.as_ref(), CANONICAL), - toolset_member_name: None, - }); - } - // MCP tool use — provider-executed + dynamic. - ContentBlock::McpToolUse { id, name, input, server_name } => { - mcp_tool_calls - .insert(id.clone(), (name.clone(), server_name.clone())); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id.clone(), - tool_name: name.clone(), - input: input.to_string(), - provider_executed: Some(true), - dynamic: Some(true), - provider_metadata: Some(provider_namespace(CANONICAL, json!({ - "type": "mcp-tool-use", - "serverName": server_name, - })).expect("provider metadata must be an object")), - })); - } - // Redacted thinking — emit as ReasoningStart. - ContentBlock::RedactedThinking { data } => { - let id = index.to_string(); - yield Ok(StreamPart::ReasoningStart { - id: id.clone(), - provider_metadata: Some(provider_namespace(CANONICAL, json!({ "redactedData": data })).expect("provider metadata must be an object")), - }); - blocks.insert( - index, - BlockState::Thinking, - ); - } - // Server-tool result blocks arrive whole on - // `content_block_start`, so they reuse the - // non-streaming payload mapping. Upstream - // mirrors its `doGenerate` switch here too - // (anthropic-language-model.ts:1901-2178). - other => { - if let ContentBlock::WebFetchToolResult { content, .. } = &other - && let Some(document) = web_fetch_document(content) { - citation_documents.push(document); - } - for part in stream_parts_for_result_block( - CANONICAL, - &other, - &tool_names, - &mcp_tool_calls, - &server_tool_calls, - generate_id, - ) { - yield Ok(part); - } - } - } - } - StreamEvent::ContentBlockDelta { index, delta } => { - if !uses_json_response_tool && let Some(text) = delta.text { - yield Ok(StreamPart::TextDelta { - id: index.to_string(), - delta: text, - provider_metadata: None, - }); - } - if let Some(text) = delta.content { - yield Ok(StreamPart::TextDelta { id: index.to_string(), delta: text, provider_metadata: None }); - } - if let Some(citation) = delta.citation { - if let Some(BlockState::Text { citations }) = blocks.get_mut(&index) { - citations.push(citation.clone()); - } - if let Some(source) = citation_source(&citation, &citation_documents, generate_id) { - yield Ok(StreamPart::Source(source)); - } - } - if let Some(partial) = delta.partial_json { - // Accumulate the partial JSON fragment and emit - // a ToolInputDelta. Empty fragments (the - // leading `input_json_delta` with - // `partial_json: ""`) are skipped, matching the - // TS SDK. - if is_json_response_from_tool && matches!(blocks.get(&index), Some(BlockState::Text { .. })) && !partial.is_empty() { - yield Ok(StreamPart::TextDelta { id: index.to_string(), delta: partial, provider_metadata: None }); - continue; - } - if is_json_response_from_tool { continue; } - let delta_event: Option<(String, String)> = - match blocks.get_mut(&index) { - Some(BlockState::ToolUse { - id, - accumulated_json, - provider_tool_input_type, - first_delta, - toolset_member_name, - .. - }) if !partial.is_empty() => { - let emitted_delta = if *first_delta { - if let Some(input_type) = provider_tool_input_type { - format!( - "{{\"type\": \"{input_type}\",{}", - partial.strip_prefix('{').unwrap_or(&partial) - ) - } else { - partial - } - } else { - partial - }; - accumulated_json.push_str(&emitted_delta); - *first_delta = false; - toolset_member_name.is_none().then(|| (id.clone(), emitted_delta)) - } - _ => None, - }; - if let Some((id, delta)) = delta_event { - yield Ok(StreamPart::ToolInputDelta { - id, - delta, - provider_metadata: None, - }); - } - } - if let Some(thinking) = delta.thinking { - yield Ok(StreamPart::ReasoningDelta { - id: index.to_string(), - delta: thinking, - provider_metadata: None, - }); - } - if let Some(sig) = delta.signature - && let Some(BlockState::Thinking) = blocks.get(&index) { - yield Ok(StreamPart::ReasoningDelta { - id: index.to_string(), delta: String::new(), - provider_metadata: Some(provider_namespace(CANONICAL, json!({ "signature": sig })).expect("provider metadata must be an object")), - }); - } - } - StreamEvent::ContentBlockStop { index } => { - // Removing the block releases the borrow before any - // yield. - if let Some(state) = blocks.remove(&index) { - match state { - BlockState::Text { citations } => { - yield Ok(StreamPart::TextEnd { - id: index.to_string(), - provider_metadata: citation_metadata(&citations), - }); - } - BlockState::Thinking => { - yield Ok(StreamPart::ReasoningEnd { - id: index.to_string(), - provider_metadata: None, - }); - } - BlockState::ToolUse { - id, - name, - mut accumulated_json, - provider_executed, - dynamic, - provider_tool_name, - provider_tool_input_type, - provider_metadata, - toolset_member_name, - .. - } => { - if let Some(member_name) = toolset_member_name { - let parsed = if accumulated_json.is_empty() { - Ok(json!({})) - } else { - serde_json::from_str::(&accumulated_json) - }; - if let Ok(input) = parsed { - accumulated_json = toolset_member_input(&member_name, &input).to_string(); - } - yield Ok(StreamPart::ToolInputDelta { - id: id.clone(), - delta: accumulated_json.clone(), - provider_metadata: None, - }); - } - yield Ok(StreamPart::ToolInputEnd { - id: id.clone(), - provider_metadata: None, - }); - let input = - finalize_streamed_tool_input( - accumulated_json, - provider_tool_name.as_deref(), - provider_tool_input_type.as_deref(), - ); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: name, - input, - provider_executed, - dynamic, - provider_metadata, - })); - } - } - } - } - StreamEvent::MessageDelta { delta, usage, context_management: edits, input_transformations: transformations } => { - if let Some(reason) = delta.stop_reason { - final_finish_reason = Some(parse_stop_reason(&reason)); - } - stop_sequence = delta.stop_sequence; - container = delta.container; - context_management = edits.or(context_management); - stop_details = delta.stop_details; - input_transformations = transformations.or(input_transformations); - safeguard_results = delta.safeguard_results.or(safeguard_results); - if let Some(u) = usage { - if let (Value::Object(raw), Ok(Value::Object(update))) = - (&mut raw_usage, serde_json::to_value(&u)) - { - if let Value::Object(effective) = &mut effective_usage { - effective.extend(update.iter().filter(|(_, value)| !value.is_null()).map(|(key, value)| (key.clone(), value.clone()))); - } - raw.extend(update); - } - if let Ok(usage) = serde_json::from_value(effective_usage.clone()) { - final_usage = super::usage::usage_from_anthropic(&usage); - final_usage.raw = raw_usage.as_object().cloned(); - } - } - } - StreamEvent::MessageStop => { - message_stopped = true; - break; +impl AnthropicMessagesStream { + fn message_start( + &mut self, + message: MessageStartData, + controller: &mut TransformStreamController, + ) { + if let Some(active_id) = &self.active_message_id { + if active_id == &message.id { + return; + } + controller.enqueue(StreamPart::Error { + error: AiMuxError::InvalidResponseData(format!( + "Received message_start for message {:?} while message {:?} is still open.", + message.id, active_id, + )), + }); + controller.terminate(); + return; + } + self.active_message_id = Some(message.id.clone()); + self.input_transformations = message + .input_transformations + .or(self.input_transformations.take()); + self.container = message.container.or(self.container.take()); + if let Some(reason) = message.stop_reason.as_deref() { + self.final_finish_reason = Some(parse_stop_reason(reason)); + } + if let Some(usage) = &message.usage { + // RFC-0015 P0-2: full input side incl. cache + // fields + raw (Anthropic reports cache only + // in message_start). + self.final_usage = super::usage::usage_from_anthropic(usage); + if let Ok(raw_usage) = serde_json::to_value(usage) { + self.raw_usage = raw_usage; + } + self.effective_usage = self.raw_usage.clone(); + } + if !self.response_meta_emitted { + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: Some(message.id.clone()), + timestamp: None, + model_id: Some(message.model.clone()), + })); + self.response_meta_emitted = true; + } + } + + fn content_block_start( + &mut self, + index: usize, + content_block: ContentBlock, + controller: &mut TransformStreamController, + ) { + match content_block { + ContentBlock::Text { .. } if self.uses_json_response_tool => {} + ContentBlock::Text { .. } => { + self.blocks.insert( + index, + BlockState::Text { + citations: Vec::new(), + }, + ); + controller.enqueue(StreamPart::TextStart { + id: index.to_string(), + provider_metadata: None, + }); + } + ContentBlock::Compaction { content, signature } => { + self.blocks.insert( + index, + BlockState::Text { + citations: Vec::new(), + }, + ); + controller.enqueue(StreamPart::TextStart { + id: index.to_string(), + provider_metadata: Some(compaction_metadata(signature.as_deref())), + }); + if signature.is_some() + && let Some(text) = content.filter(|text| !text.is_empty()) + { + controller.enqueue(StreamPart::TextDelta { + id: index.to_string(), + delta: text, + provider_metadata: None, + }); + } + } + ContentBlock::Thinking { .. } => { + controller.enqueue(StreamPart::ReasoningStart { + id: index.to_string(), + provider_metadata: None, + }); + self.blocks.insert(index, BlockState::Thinking); + } + ContentBlock::ToolUse { name, .. } + if self.uses_json_response_tool && name == "json" => + { + self.is_json_response_from_tool = true; + self.blocks.insert( + index, + BlockState::Text { + citations: Vec::new(), + }, + ); + controller.enqueue(StreamPart::TextStart { + id: index.to_string(), + provider_metadata: None, + }); + } + ContentBlock::ToolUse { + id, + name, + input, + toolset_name, + caller, + } => { + let custom_name = toolset_name.as_deref().map_or_else( + || name.clone(), + |name| self.tool_names.to_custom_tool_name(name).to_string(), + ); + let initial_input = initial_tool_input(&input); + controller.enqueue(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: custom_name.clone(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + }); + self.blocks.insert( + index, + BlockState::ToolUse { + id, + name: custom_name, + first_delta: initial_input.is_empty(), + accumulated_json: initial_input, + provider_executed: None, + dynamic: None, + provider_tool_name: None, + provider_tool_input_type: None, + provider_metadata: tool_call_metadata( + caller.as_ref(), + toolset_name.as_deref(), + CANONICAL, + ), + toolset_member_name: toolset_name.map(|_| name), + }, + ); + } + // Server-side tool use follows the same input + // lifecycle as client tools because some code + // execution inputs arrive entirely via deltas. + ContentBlock::ServerToolUse { + id, + name, + input, + caller, + } => { + if is_tool_search_provider_name(&name) { + self.server_tool_calls.insert(id.clone(), name.clone()); + } + let provider_name = server_tool_provider_name(&name); + let custom_name = self + .tool_names + .to_custom_tool_name(provider_name) + .to_string(); + let dynamic = (provider_name == "code_execution" + && self.tool_names.mark_code_execution_dynamic()) + .then_some(true); + let initial_input = initial_tool_input(&input); + controller.enqueue(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: custom_name.clone(), + provider_executed: Some(true), + dynamic, + title: None, + provider_metadata: None, + }); + self.blocks.insert( + index, + BlockState::ToolUse { + id, + name: custom_name, + first_delta: initial_input.is_empty(), + accumulated_json: initial_input, + provider_executed: Some(true), + dynamic, + provider_tool_name: Some(provider_name.to_string()), + provider_tool_input_type: match name.as_str() { + "text_editor_code_execution" | "bash_code_execution" => Some(name), + "code_execution" => Some("programmatic-tool-call".to_string()), + _ => None, }, - StreamEvent::Error { error } => { - yield Ok(StreamPart::Error { - error: anthropic_stream_error( - &error, - &stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ), - }); + provider_metadata: tool_call_caller_metadata(caller.as_ref(), CANONICAL), + toolset_member_name: None, + }, + ); + } + // MCP tool use — provider-executed + dynamic. + ContentBlock::McpToolUse { + id, + name, + input, + server_name, + } => { + self.mcp_tool_calls + .insert(id.clone(), (name.clone(), server_name.clone())); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id.clone(), + tool_name: name.clone(), + input: input.to_string(), + provider_executed: Some(true), + dynamic: Some(true), + provider_metadata: Some( + provider_namespace( + CANONICAL, + json!({ + "type": "mcp-tool-use", + "serverName": server_name, + }), + ) + .expect("provider metadata must be an object"), + ), + })); + } + // Redacted thinking — emit as ReasoningStart. + ContentBlock::RedactedThinking { data } => { + let id = index.to_string(); + controller.enqueue(StreamPart::ReasoningStart { + id: id.clone(), + provider_metadata: Some( + provider_namespace(CANONICAL, json!({ "redactedData": data })) + .expect("provider metadata must be an object"), + ), + }); + self.blocks.insert(index, BlockState::Thinking); + } + // Server-tool result blocks arrive whole on + // `content_block_start`, so they reuse the + // non-streaming payload mapping. Upstream + // mirrors its `doGenerate` switch here too + // (anthropic-language-model.ts:1901-2178). + other => { + if let ContentBlock::WebFetchToolResult { content, .. } = &other + && let Some(document) = web_fetch_document(content) + { + self.citation_documents.push(document); + } + for part in stream_parts_for_result_block( + CANONICAL, + &other, + &self.tool_names, + &self.mcp_tool_calls, + &self.server_tool_calls, + self.generate_id, + ) { + controller.enqueue(part); + } + } + } + } + fn content_block_delta( + &mut self, + index: usize, + delta: DeltaBlock, + controller: &mut TransformStreamController, + ) { + if !self.uses_json_response_tool + && let Some(text) = delta.text + { + controller.enqueue(StreamPart::TextDelta { + id: index.to_string(), + delta: text, + provider_metadata: None, + }); + } + if let Some(text) = delta.content { + controller.enqueue(StreamPart::TextDelta { + id: index.to_string(), + delta: text, + provider_metadata: None, + }); + } + if let Some(citation) = delta.citation { + if let Some(BlockState::Text { citations }) = self.blocks.get_mut(&index) { + citations.push(citation.clone()); + } + if let Some(source) = + citation_source(&citation, &self.citation_documents, self.generate_id) + { + controller.enqueue(StreamPart::Source(source)); + } + } + if let Some(partial) = delta.partial_json { + // Accumulate the partial JSON fragment and emit + // a ToolInputDelta. Empty fragments (the + // leading `input_json_delta` with + // `partial_json: ""`) are skipped, matching the + // TS SDK. + if self.is_json_response_from_tool + && matches!(self.blocks.get(&index), Some(BlockState::Text { .. })) + && !partial.is_empty() + { + controller.enqueue(StreamPart::TextDelta { + id: index.to_string(), + delta: partial, + provider_metadata: None, + }); + return; + } + if self.is_json_response_from_tool { + return; + } + let delta_event: Option<(String, String)> = match self.blocks.get_mut(&index) { + Some(BlockState::ToolUse { + id, + accumulated_json, + provider_tool_input_type, + first_delta, + toolset_member_name, + .. + }) if !partial.is_empty() => { + let emitted_delta = if *first_delta { + if let Some(input_type) = provider_tool_input_type { + format!( + "{{\"type\": \"{input_type}\",{}", + partial.strip_prefix('{').unwrap_or(&partial) + ) + } else { + partial } - _ => {} - } + } else { + partial + }; + accumulated_json.push_str(&emitted_delta); + *first_delta = false; + toolset_member_name + .is_none() + .then(|| (id.clone(), emitted_delta)) } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Ok(StreamPart::Error { error }); - if !recoverable { - return; + _ => None, + }; + if let Some((id, delta)) = delta_event { + controller.enqueue(StreamPart::ToolInputDelta { + id, + delta, + provider_metadata: None, + }); + } + } + if let Some(thinking) = delta.thinking { + controller.enqueue(StreamPart::ReasoningDelta { + id: index.to_string(), + delta: thinking, + provider_metadata: None, + }); + } + if let Some(sig) = delta.signature + && let Some(BlockState::Thinking) = self.blocks.get(&index) + { + controller.enqueue(StreamPart::ReasoningDelta { + id: index.to_string(), + delta: String::new(), + provider_metadata: Some( + provider_namespace(CANONICAL, json!({ "signature": sig })) + .expect("provider metadata must be an object"), + ), + }); + } + } + + fn content_block_stop( + &mut self, + index: usize, + controller: &mut TransformStreamController, + ) { + let Some(state) = self.blocks.remove(&index) else { + return; + }; + match state { + BlockState::Text { citations } => { + controller.enqueue(StreamPart::TextEnd { + id: index.to_string(), + provider_metadata: citation_metadata(&citations), + }); + } + BlockState::Thinking => { + controller.enqueue(StreamPart::ReasoningEnd { + id: index.to_string(), + provider_metadata: None, + }); + } + BlockState::ToolUse { + id, + name, + mut accumulated_json, + provider_executed, + dynamic, + provider_tool_name, + provider_tool_input_type, + provider_metadata, + toolset_member_name, + .. + } => { + if let Some(member_name) = toolset_member_name { + let parsed = if accumulated_json.is_empty() { + Ok(json!({})) + } else { + serde_json::from_str::(&accumulated_json) + }; + if let Ok(input) = parsed { + accumulated_json = toolset_member_input(&member_name, &input).to_string(); } + controller.enqueue(StreamPart::ToolInputDelta { + id: id.clone(), + delta: accumulated_json.clone(), + provider_metadata: None, + }); } + controller.enqueue(StreamPart::ToolInputEnd { + id: id.clone(), + provider_metadata: None, + }); + let input = finalize_streamed_tool_input( + accumulated_json, + provider_tool_name.as_deref(), + provider_tool_input_type.as_deref(), + ); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id, + tool_name: name, + input, + provider_executed, + dynamic, + provider_metadata, + })); } } + } + fn message_delta( + &mut self, + delta: MessageDeltaBody, + usage: Option, + edits: Option, + transformations: Option, + ) { + if let Some(reason) = delta.stop_reason { + self.final_finish_reason = Some(parse_stop_reason(&reason)); + } + self.stop_sequence = delta.stop_sequence; + self.container = delta.container; + self.context_management = edits.or(self.context_management.take()); + self.stop_details = delta.stop_details; + self.input_transformations = transformations.or(self.input_transformations.take()); + self.safeguard_results = delta.safeguard_results.or(self.safeguard_results.take()); + if let Some(u) = usage { + if let (Value::Object(raw), Ok(Value::Object(update))) = + (&mut self.raw_usage, serde_json::to_value(&u)) + { + if let Value::Object(effective) = &mut self.effective_usage { + effective.extend( + update + .iter() + .filter(|(_, value)| !value.is_null()) + .map(|(key, value)| (key.clone(), value.clone())), + ); + } + raw.extend(update); + } + if let Ok(usage) = serde_json::from_value(self.effective_usage.clone()) { + self.final_usage = super::usage::usage_from_anthropic(&usage); + self.final_usage.raw = self.raw_usage.as_object().cloned(); + } + } + } +} + +impl Transformer for AnthropicMessagesStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut TransformStreamController) { + // First part: StreamStart. + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + if self.message_stopped { + return; + } + let stream_event = match event { + Ok(stream_event) => stream_event, + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + controller.enqueue(StreamPart::Error { error }); + if !recoverable { + controller.terminate(); + } + return; + } + }; + match stream_event { + StreamEvent::MessageStart { message } => self.message_start(message, controller), + StreamEvent::ContentBlockStart { + index, + content_block, + } => self.content_block_start(index, content_block, controller), + StreamEvent::ContentBlockDelta { index, delta } => { + self.content_block_delta(index, delta, controller); + } + StreamEvent::ContentBlockStop { index } => self.content_block_stop(index, controller), + StreamEvent::MessageDelta { + delta, + usage, + context_management: edits, + input_transformations: transformations, + } => self.message_delta(delta, usage, edits, transformations), + StreamEvent::MessageStop => self.message_stopped = true, + StreamEvent::Error { error } => { + controller.enqueue(StreamPart::Error { + error: anthropic_stream_error( + &error, + &self.stream_error_url, + self.stream_request_body.clone(), + self.stream_response_headers.clone(), + ), + }); + } + StreamEvent::Other => {} + } + } + + fn flush(mut self, controller: &mut TransformStreamController) { // The provider completes a message explicitly; EOF after an error // does not synthesize a finish event. - if !message_stopped { + if !self.message_stopped { return; } - if is_json_response_from_tool && let Some(reason) = &mut final_finish_reason - && reason.unified == FinishReasonUnified::ToolCalls { + if self.is_json_response_from_tool + && let Some(reason) = &mut self.final_finish_reason + && reason.unified == FinishReasonUnified::ToolCalls + { reason.unified = FinishReasonUnified::Stop; } - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { + let mut metadata = super::usage::result_provider_metadata( + self.options_name.as_str(), + &self.raw_usage, + self.stop_sequence.as_deref(), + self.container.as_ref(), + self.context_management.as_ref(), + self.used_custom_options_key, + ); + super::usage::extend_result_metadata( + &mut metadata, + self.stop_details.as_ref(), + self.input_transformations.as_ref(), + self.safeguard_results.as_ref(), + ); + controller.enqueue(StreamPart::Finish { + finish_reason: self.final_finish_reason.unwrap_or(FinishReason { unified: FinishReasonUnified::Other, raw: None, }), - usage: final_usage, - provider_metadata: Some({ - let mut metadata = super::usage::result_provider_metadata( - options_name.as_str(), - &raw_usage, - stop_sequence.as_deref(), - container.as_ref(), - context_management.as_ref(), - used_custom_options_key, - ); - super::usage::extend_result_metadata(&mut metadata, stop_details.as_ref(), input_transformations.as_ref(), safeguard_results.as_ref()); - metadata - }), + usage: self.final_usage, + provider_metadata: Some(metadata), }); - }; - - 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), - }), - }) + } } diff --git a/aimux-providers/src/bedrock/model.rs b/aimux-providers/src/bedrock/model.rs index ea9534adf..08ec57395 100644 --- a/aimux-providers/src/bedrock/model.rs +++ b/aimux-providers/src/bedrock/model.rs @@ -18,16 +18,17 @@ use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{ - FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, + FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, Warning, }; use serde_json::json; -use aimux_provider_utils::HttpBody; +use aimux_provider_utils::{HttpBody, TransformStreamController, Transformer, pipe_through}; use super::convert::{ build_request_body_checked, convert_usage, map_finish_reason, normalize_tool_call_id, }; +use super::event_stream::EventStreamMessage; use super::options; use super::types::{BedrockContentBlock, BedrockConverseResponse}; use crate::shared::EndpointConfig; @@ -276,305 +277,491 @@ impl LanguageModel for BedrockModel { // (RFC1123). Keeps stream/non-stream timestamps consistent. let response_timestamp = response_headers.get("date").cloned(); - let model_id = self.model_id.clone(); + let stream = pipe_through( + super::event_stream::decode_stream(body_stream), + AmazonBedrockChatStream { + warnings, + request_id, + response_timestamp, + model_id: self.model_id.clone(), + uses_json_tool, + include_raw_chunks, + error_url, + error_body, + extractor: uses_json_instruction.then(JsonObjectTextExtractor::default), + is_json_response_from_tool: false, + text_blocks: HashSet::new(), + reasoning_blocks: HashMap::new(), + block_counter: 0, + tool_blocks: HashMap::new(), + final_usage: Usage::default(), + final_finish_reason: None, + finish_meta: serde_json::Map::new(), + stop_sequence: None, + }, + ); + + 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), + }), + }) + } +} - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); +/// The `TransformStream` of `BedrockChatLanguageModel.doStream`. +struct AmazonBedrockChatStream { + warnings: Vec, + request_id: Option, + response_timestamp: Option, + model_id: String, + uses_json_tool: bool, + include_raw_chunks: bool, + error_url: String, + error_body: serde_json::Value, + extractor: Option, + is_json_response_from_tool: bool, + text_blocks: HashSet, + reasoning_blocks: HashMap>, + block_counter: usize, + /// Tool call state: block_index → (id, name, accumulated_json) + tool_blocks: HashMap, + final_usage: Usage, + final_finish_reason: Option, + /// Provider metadata accumulated for the Finish chunk from `metadata` + /// and `messageStop` events (findings #26, #27). Mirrors the TS + /// `providerMetadata` / `stopSequence` accumulation in + /// `amazon-bedrock-chat-language-model.ts` (`doStream`). + finish_meta: serde_json::Map, + stop_sequence: Option, +} - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: request_id, - timestamp: response_timestamp, - model_id: Some(model_id.clone()), - })); +impl AmazonBedrockChatStream { + fn content_block_start( + &mut self, + payload: &serde_json::Value, + controller: &mut TransformStreamController, + ) { + let idx = payload + .get("contentBlockIndex") + .and_then(serde_json::Value::as_u64) + .unwrap_or(self.block_counter as u64) as usize; + + // Check if this is a tool use block. + if let Some(start) = payload.get("start") { + if let Some(tool_use) = start.get("toolUse") { + let name = tool_use + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let id = tool_use + .get("toolUseId") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let id = normalize_tool_call_id( + &if id.is_empty() { + aimux_provider_utils::generate_id() + } else { + id + }, + self.model_id.contains("mistral."), + ); + if !(self.uses_json_tool && name == "json") { + controller.enqueue(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: name.clone(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + }); + } + self.tool_blocks.insert(idx, (id, name, String::new())); + } else { + // Text block. + self.block_counter = idx + 1; + let id = idx.to_string(); + self.text_blocks.insert(idx); + controller.enqueue(StreamPart::TextStart { + id, + provider_metadata: None, + }); + } + } else { + // Default: text block. + let id = idx.to_string(); + self.text_blocks.insert(idx); + controller.enqueue(StreamPart::TextStart { + id, + provider_metadata: None, + }); + } + } + + fn content_block_delta( + &mut self, + payload: &serde_json::Value, + controller: &mut TransformStreamController, + ) { + let idx = payload + .get("contentBlockIndex") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + + let Some(delta) = payload.get("delta") else { + return; + }; + // Text delta + if let Some(text) = delta.get("text").and_then(|v| v.as_str()) + && !text.is_empty() + { + if self.text_blocks.insert(idx) { + let id = idx.to_string(); + self.text_blocks.insert(idx); + controller.enqueue(StreamPart::TextStart { + id, + provider_metadata: None, + }); + } + let delta = self + .extractor + .as_mut() + .map(|extractor| extractor.process(text)) + .unwrap_or_else(|| text.to_string()); + if !delta.is_empty() { + controller.enqueue(StreamPart::TextDelta { + id: idx.to_string(), + delta, + provider_metadata: None, + }); + } + } + // Tool use input delta + if let Some(partial) = delta.get("toolUse").and_then(|t| t.get("input")) + && let Some(partial_str) = partial.as_str() + && let Some((id, name, acc)) = self.tool_blocks.get_mut(&idx) + && !partial_str.is_empty() + { + acc.push_str(partial_str); + if !(self.uses_json_tool && name == "json") { + let id = id.clone(); + controller.enqueue(StreamPart::ToolInputDelta { + id, + delta: partial_str.to_string(), + provider_metadata: None, + }); + } + } + if let Some(rc) = delta.get("reasoningContent") { + if !["text", "signature", "data", "redactedContent"] + .iter() + .any(|key| { + rc.get(key) + .and_then(|v| v.as_str()) + .is_some_and(|s| !s.is_empty()) + }) + { + return; + } + let id = idx.to_string(); + if let std::collections::hash_map::Entry::Vacant(entry) = + self.reasoning_blocks.entry(idx) + { + entry.insert(None); + controller.enqueue(StreamPart::ReasoningStart { + id: id.clone(), + provider_metadata: None, + }); + } + if let Some(text) = rc.get("text").and_then(|v| v.as_str()) { + controller.enqueue(StreamPart::ReasoningDelta { + id: id.clone(), + delta: text.to_string(), + provider_metadata: None, + }); + } else if let Some(sig) = rc + .get("signature") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + { + controller.enqueue(StreamPart::ReasoningDelta { + id: id.clone(), + delta: String::new(), + provider_metadata: Some(options::metadata(json!({"signature": sig}))), + }); + } else if let Some(data) = rc + .get("data") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + { + controller.enqueue(StreamPart::ReasoningDelta { + id: id.clone(), + delta: String::new(), + provider_metadata: Some(options::metadata(json!({"redactedData": data}))), + }); + } else if let Some(data) = rc.get("redactedContent").and_then(|v| v.as_str()) { + self.reasoning_blocks + .entry(idx) + .or_default() + .get_or_insert_with(String::new) + .push_str(data); + } + } + } - let mut messages = super::event_stream::decode_stream(body_stream); - - let mut extractor = uses_json_instruction.then(JsonObjectTextExtractor::default); - let mut is_json_response_from_tool = false; - let mut text_blocks = HashSet::new(); - let mut reasoning_blocks: HashMap> = HashMap::new(); - let mut block_counter = 0usize; - // Tool call state: block_index → (id, name, accumulated_json) - let mut tool_blocks: HashMap = HashMap::new(); - let mut final_usage: Usage = Usage::default(); - let mut final_finish_reason: Option = None; - // Provider metadata accumulated for the Finish chunk from `metadata` - // and `messageStop` events (findings #26, #27). Mirrors the TS - // `providerMetadata` / `stopSequence` accumulation in - // `amazon-bedrock-chat-language-model.ts` (`doStream`). - let mut finish_meta: serde_json::Map = - serde_json::Map::new(); - let mut stop_sequence: Option = None; - - while let Some(message) = messages.next().await { - let msg = match message { - Ok(msg) => msg, - Err(error) => { yield Err(error); return; } + fn content_block_stop( + &mut self, + payload: &serde_json::Value, + controller: &mut TransformStreamController, + ) { + let idx = payload + .get("contentBlockIndex") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + + if let Some((id, name, acc)) = self.tool_blocks.remove(&idx) { + if self.uses_json_tool && name == "json" { + self.is_json_response_from_tool = true; + let id = idx.to_string(); + controller.enqueue(StreamPart::TextStart { + id: id.clone(), + provider_metadata: None, + }); + controller.enqueue(StreamPart::TextDelta { + id: id.clone(), + delta: acc, + provider_metadata: None, + }); + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: None, + }); + } else { + controller.enqueue(StreamPart::ToolInputEnd { + id: id.clone(), + provider_metadata: None, + }); + // Empty input normalizes to "{}" per the upstream + // provider. + let input = if acc.is_empty() { + "{}".to_string() + } else { + acc }; - let payload: serde_json::Value = match serde_json::from_str(&msg.data) { - Ok(v) => v, + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: normalize_tool_call_id(&id, self.model_id.contains("mistral.")), + tool_name: name, + input, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); + } + } else if let Some(redacted) = self.reasoning_blocks.remove(&idx) { + controller.enqueue(StreamPart::ReasoningEnd { + id: idx.to_string(), + provider_metadata: reasoning_redacted_meta(redacted), + }); + } else if self.text_blocks.remove(&idx) { + controller.enqueue(StreamPart::TextEnd { + id: idx.to_string(), + provider_metadata: None, + }); + } + } + + fn message_stop(&mut self, payload: &serde_json::Value) { + if let Some(reason) = payload.get("stopReason").and_then(|v| v.as_str()) { + let mut finish_reason = map_finish_reason(reason); + if self.is_json_response_from_tool + && finish_reason.unified == FinishReasonUnified::ToolCalls + { + finish_reason.unified = FinishReasonUnified::Stop; + } + self.final_finish_reason = Some(finish_reason); + } + // #26: surface which stop sequence sentinel fired + // (additionalModelResponseFields.delta.stop_sequence), if any. + if let Some(seq) = payload + .get("additionalModelResponseFields") + .and_then(|f| f.get("delta")) + .and_then(|d| d.get("stop_sequence")) + .and_then(|v| v.as_str()) + { + self.stop_sequence = Some(seq.to_string()); + } + } + + fn metadata( + &mut self, + payload: &serde_json::Value, + controller: &mut TransformStreamController, + ) { + if let Some(usage) = payload.get("usage") { + let bedrock_usage: super::types::BedrockUsage = + match serde_json::from_value(usage.clone()) { + Ok(usage) => usage, Err(error) => { - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - yield Ok(StreamPart::Error { error: AiMuxError::from(error) }); - continue; + self.final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }); + controller.enqueue(StreamPart::Error { + error: AiMuxError::InvalidResponseData(format!( + "{error}; usage: {usage}" + )), + }); + return; } }; - if include_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: json!({ msg.event_type.clone(): payload.clone() }) }); - } - if msg.message_type != "event" || msg.event_type.ends_with("Exception") { - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - let (status_code, is_retryable) = match msg.event_type.as_str() { - "internalServerException" | "InternalServerException" => (Some(500), true), - "modelStreamErrorException" | "ModelStreamErrorException" => (Some(424), true), - "serviceUnavailableException" | "ServiceUnavailableException" => (Some(503), true), - "throttlingException" | "ThrottlingException" => (Some(429), true), - "validationException" | "ValidationException" => (Some(400), false), - _ => (None, false), - }; - let message = payload.get("message").and_then(|v| v.as_str()).map(str::to_string) - .unwrap_or_else(|| format!("Amazon Bedrock stream failed with {}", msg.event_type)); - yield Ok(StreamPart::Error { error: AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { - status_code, is_retryable, provider_code: Some(msg.event_type.clone()), - data: Some(json!({msg.event_type.clone(): payload.clone()})), - ..aimux_core::ApiCallError::new(message, error_url.clone(), error_body.clone()) - })) }); - continue; - } - match msg.event_type.as_str() { - "messageStart" => { - // Nothing to emit; response metadata already sent. - } - "contentBlockStart" => { - let idx = payload - .get("contentBlockIndex") - .and_then(serde_json::Value::as_u64) - .unwrap_or(block_counter as u64) as usize; - - // Check if this is a tool use block. - if let Some(start) = payload.get("start") { - if let Some(tool_use) = start.get("toolUse") { - let name = tool_use - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let id = tool_use - .get("toolUseId") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let id = normalize_tool_call_id(&if id.is_empty() { - aimux_provider_utils::generate_id() - } else { - id - }, model_id.contains("mistral.")); - if !(uses_json_tool && name == "json") { - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: name.clone(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } - tool_blocks.insert(idx, (id, name, String::new())); - } else { - // Text block. - block_counter = idx + 1; - let id = idx.to_string(); - text_blocks.insert(idx); - yield Ok(StreamPart::TextStart { id, provider_metadata: None}); - } - } else { - // Default: text block. - let id = idx.to_string(); - text_blocks.insert(idx); - yield Ok(StreamPart::TextStart { id, provider_metadata: None}); - } - } - "contentBlockDelta" => { - let idx = payload - .get("contentBlockIndex") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - - if let Some(delta) = payload.get("delta") { - // Text delta - if let Some(text) = delta.get("text").and_then(|v| v.as_str()) - && !text.is_empty() { - if text_blocks.insert(idx) { - let id = idx.to_string(); - text_blocks.insert(idx); - yield Ok(StreamPart::TextStart { id, provider_metadata: None}); - } - let delta = extractor.as_mut().map(|extractor| extractor.process(text)).unwrap_or_else(|| text.to_string()); - if !delta.is_empty() { - yield Ok(StreamPart::TextDelta { id: idx.to_string(), delta, provider_metadata: None }); - } - } - // Tool use input delta - if let Some(partial) = - delta.get("toolUse").and_then(|t| t.get("input")) - && let Some(partial_str) = partial.as_str() - && let Some((id, name, acc)) = tool_blocks.get_mut(&idx) - && !partial_str.is_empty() { - acc.push_str(partial_str); - if !(uses_json_tool && name == "json") { - let id = id.clone(); - yield Ok(StreamPart::ToolInputDelta { - id, - delta: partial_str.to_string(), - provider_metadata: None, - }); - } - } - if let Some(rc) = delta.get("reasoningContent") { - if !["text", "signature", "data", "redactedContent"].iter().any(|key| rc.get(key).and_then(|v| v.as_str()).is_some_and(|s| !s.is_empty())) { continue; } - let id = idx.to_string(); - if let std::collections::hash_map::Entry::Vacant(entry) = reasoning_blocks.entry(idx) { - entry.insert(None); - yield Ok(StreamPart::ReasoningStart { id: id.clone(), provider_metadata: None }); - } - if let Some(text) = rc.get("text").and_then(|v| v.as_str()) { - yield Ok(StreamPart::ReasoningDelta { id: id.clone(), delta: text.to_string(), provider_metadata: None }); - } else if let Some(sig) = rc.get("signature").and_then(|v| v.as_str()).filter(|s| !s.is_empty()) { - yield Ok(StreamPart::ReasoningDelta { id: id.clone(), delta: String::new(), provider_metadata: Some(options::metadata(json!({"signature": sig}))) }); - } else if let Some(data) = rc.get("data").and_then(|v| v.as_str()).filter(|s| !s.is_empty()) { - yield Ok(StreamPart::ReasoningDelta { id: id.clone(), delta: String::new(), provider_metadata: Some(options::metadata(json!({"redactedData": data}))) }); - } else if let Some(data) = rc.get("redactedContent").and_then(|v| v.as_str()) { - reasoning_blocks.entry(idx).or_default().get_or_insert_with(String::new).push_str(data); - } - } - } - } - "contentBlockStop" => { - let idx = payload - .get("contentBlockIndex") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - - if let Some((id, name, acc)) = tool_blocks.remove(&idx) { - if uses_json_tool && name == "json" { - is_json_response_from_tool = true; - let id = idx.to_string(); - yield Ok(StreamPart::TextStart { id: id.clone(), provider_metadata: None }); - yield Ok(StreamPart::TextDelta { id: id.clone(), delta: acc, provider_metadata: None }); - yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); - } else { - yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: None}); - // Empty input normalizes to "{}" per the upstream - // provider. - let input = if acc.is_empty() { - "{}".to_string() - } else { - acc - }; - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: normalize_tool_call_id(&id, model_id.contains("mistral.")), - tool_name: name, - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - } - } else if let Some(redacted) = reasoning_blocks.remove(&idx) { - yield Ok(StreamPart::ReasoningEnd { id: idx.to_string(), provider_metadata: reasoning_redacted_meta(redacted) }); - } else if text_blocks.remove(&idx) { - yield Ok(StreamPart::TextEnd { id: idx.to_string(), provider_metadata: None }); - } - } - "messageStop" => { - if let Some(reason) = - payload.get("stopReason").and_then(|v| v.as_str()) - { - let mut finish_reason = map_finish_reason(reason); - if is_json_response_from_tool && finish_reason.unified == FinishReasonUnified::ToolCalls { finish_reason.unified = FinishReasonUnified::Stop; } - final_finish_reason = Some(finish_reason); - } - // #26: surface which stop sequence sentinel fired - // (additionalModelResponseFields.delta.stop_sequence), if any. - if let Some(seq) = payload - .get("additionalModelResponseFields") - .and_then(|f| f.get("delta")) - .and_then(|d| d.get("stop_sequence")) - .and_then(|v| v.as_str()) - { - stop_sequence = Some(seq.to_string()); - } - } - "metadata" => { - if let Some(usage) = payload.get("usage") { - let bedrock_usage: super::types::BedrockUsage = - match serde_json::from_value(usage.clone()) { - Ok(usage) => usage, - Err(error) => { - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - yield Ok(StreamPart::Error { error: AiMuxError::InvalidResponseData(format!("{error}; usage: {usage}")) }); - continue; - } - }; - final_usage = convert_usage(Some(&bedrock_usage)); - final_usage.raw = usage.as_object().cloned(); - if let Some(cache_usage) = cache_usage_metadata(Some(&bedrock_usage)) { finish_meta.insert("usage".into(), cache_usage); } - } - // #27: surface guardrails trace, performanceConfig, and - // serviceTier into the Finish chunk's provider_metadata - // (mirrors TS `doStream` metadata handling). - if let Some(trace) = payload.get("trace") { - finish_meta.insert("trace".to_string(), trace.clone()); - } - if let Some(pc) = payload.get("performanceConfig") { - finish_meta - .insert("performanceConfig".to_string(), pc.clone()); - } - if let Some(st) = payload.get("serviceTier") { - finish_meta.insert("serviceTier".to_string(), st.clone()); - } - } - _ => {} - } + self.final_usage = convert_usage(Some(&bedrock_usage)); + self.final_usage.raw = usage.as_object().cloned(); + if let Some(cache_usage) = cache_usage_metadata(Some(&bedrock_usage)) { + self.finish_meta.insert("usage".into(), cache_usage); } + } + // #27: surface guardrails trace, performanceConfig, and + // serviceTier into the Finish chunk's provider_metadata + // (mirrors TS `doStream` metadata handling). + if let Some(trace) = payload.get("trace") { + self.finish_meta.insert("trace".to_string(), trace.clone()); + } + if let Some(pc) = payload.get("performanceConfig") { + self.finish_meta + .insert("performanceConfig".to_string(), pc.clone()); + } + if let Some(st) = payload.get("serviceTier") { + self.finish_meta + .insert("serviceTier".to_string(), st.clone()); + } + } +} + +impl Transformer for AmazonBedrockChatStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: self.request_id.take(), + timestamp: self.response_timestamp.take(), + model_id: Some(self.model_id.clone()), + })); + } - // #26: merge the stop sentinel into the metadata payload, then build - // the Finish provider_metadata under `amazonBedrock` (mirrors the - // TS `doStream` flush handler). - if is_json_response_from_tool { - finish_meta.insert("isJsonResponseFromTool".into(), json!(true)); - finish_meta.insert("stopSequence".into(), serde_json::Value::Null); + fn transform( + &mut self, + message: Result, + controller: &mut TransformStreamController, + ) { + let msg = match message { + Ok(msg) => msg, + Err(error) => { + controller.error(error); + return; } - if let Some(seq) = stop_sequence { - finish_meta.insert( - "stopSequence".to_string(), - serde_json::Value::String(seq), - ); + }; + let payload: serde_json::Value = match serde_json::from_str(&msg.data) { + Ok(v) => v, + Err(error) => { + self.final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }); + controller.enqueue(StreamPart::Error { + error: AiMuxError::from(error), + }); + return; } - let provider_metadata = if finish_meta.is_empty() { - None - } else { - let payload = serde_json::Value::Object(finish_meta); - Some(options::metadata(payload)) + }; + if self.include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: json!({ msg.event_type.clone(): payload.clone() }), + }); + } + if msg.message_type != "event" || msg.event_type.ends_with("Exception") { + self.final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }); + let (status_code, is_retryable) = match msg.event_type.as_str() { + "internalServerException" | "InternalServerException" => (Some(500), true), + "modelStreamErrorException" | "ModelStreamErrorException" => (Some(424), true), + "serviceUnavailableException" | "ServiceUnavailableException" => (Some(503), true), + "throttlingException" | "ThrottlingException" => (Some(429), true), + "validationException" | "ValidationException" => (Some(400), false), + _ => (None, false), }; - - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }), - usage: final_usage, - provider_metadata, + let message = payload + .get("message") + .and_then(|v| v.as_str()) + .map(str::to_string) + .unwrap_or_else(|| format!("Amazon Bedrock stream failed with {}", msg.event_type)); + controller.enqueue(StreamPart::Error { + error: AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + status_code, + is_retryable, + provider_code: Some(msg.event_type.clone()), + data: Some(json!({msg.event_type.clone(): payload.clone()})), + ..aimux_core::ApiCallError::new( + message, + self.error_url.clone(), + self.error_body.clone(), + ) + })), }); + return; + } + match msg.event_type.as_str() { + "messageStart" => { + // Nothing to emit; response metadata already sent. + } + "contentBlockStart" => self.content_block_start(&payload, controller), + "contentBlockDelta" => self.content_block_delta(&payload, controller), + "contentBlockStop" => self.content_block_stop(&payload, controller), + "messageStop" => self.message_stop(&payload), + "metadata" => self.metadata(&payload, controller), + _ => {} + } + } + + fn flush(mut self, controller: &mut TransformStreamController) { + // #26: merge the stop sentinel into the metadata payload, then build + // the Finish provider_metadata under `amazonBedrock` (mirrors the + // TS `doStream` flush handler). + if self.is_json_response_from_tool { + self.finish_meta + .insert("isJsonResponseFromTool".into(), json!(true)); + self.finish_meta + .insert("stopSequence".into(), serde_json::Value::Null); + } + if let Some(seq) = self.stop_sequence { + self.finish_meta + .insert("stopSequence".to_string(), serde_json::Value::String(seq)); + } + let provider_metadata = if self.finish_meta.is_empty() { + None + } else { + let payload = serde_json::Value::Object(self.finish_meta); + Some(options::metadata(payload)) }; - 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), + controller.enqueue(StreamPart::Finish { + finish_reason: self.final_finish_reason.unwrap_or(FinishReason { + unified: FinishReasonUnified::Other, + raw: None, }), - }) + usage: self.final_usage, + provider_metadata, + }); } } diff --git a/aimux-providers/src/cohere/model.rs b/aimux-providers/src/cohere/model.rs index a129e4e59..861881943 100644 --- a/aimux-providers/src/cohere/model.rs +++ b/aimux-providers/src/cohere/model.rs @@ -6,7 +6,6 @@ use aimux_core::tool::RawToolCall; use async_trait::async_trait; -use futures::StreamExt; use serde_json::Value; use aimux_core::error::AiMuxError; @@ -14,7 +13,9 @@ use aimux_core::language_model::LanguageModel; 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 aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage, Warning}; + +use aimux_provider_utils::{TransformStreamController, Transformer, pipe_through}; use crate::shared::EndpointConfig; @@ -246,99 +247,20 @@ impl LanguageModel for CohereModel { let sse_stream = resp.value; let stream_warnings = body_result.warnings; let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings: stream_warnings }); - let mut final_usage = Usage::default(); - let mut finish_reason = FinishReason { unified: FinishReasonUnified::Other, raw: None }; - let mut pending_tool_call: Option = None; - let mut is_reasoning = false; - let mut sse_iter = sse_stream; - while let Some(event) = sse_iter.next().await { - let raw = match event { - Ok(raw) => raw, - Err(error) => { - if !error.is_recoverable_stream_error() { yield Err(error); return; } - if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: Value::Null }); } - finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; - yield Ok(StreamPart::Error { error }); - continue; - } - }; - if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: raw.clone() }); } - let parsed = match StreamEvent::parse(raw) { - Ok(parsed) => parsed, - Err(error) => { - finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; - yield Ok(StreamPart::Error { error }); - continue; - } - }; - let id = parsed.index.map(|index| index.to_string()).unwrap_or_default(); - match parsed.event_type.as_str() { - "message-start" => yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { id: parsed.id, timestamp: None, model_id: None })), - "content-start" => { - let content = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.content.as_ref()).unwrap(); - if content["type"] == "thinking" { - is_reasoning = true; - yield Ok(StreamPart::ReasoningStart { id, provider_metadata: None }); - } else { - yield Ok(StreamPart::TextStart { id, provider_metadata: None }); - } - } - "content-delta" => { - let content = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.content.as_ref()).unwrap(); - if let Some(text) = content.get("text").and_then(Value::as_str) { - yield Ok(StreamPart::TextDelta { id, delta: text.into(), provider_metadata: None }); - } else { - yield Ok(StreamPart::ReasoningDelta { id, delta: content["thinking"].as_str().unwrap().into(), provider_metadata: None }); - } - } - "content-end" => { - if is_reasoning { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); - is_reasoning = false; - } else { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); - } - } - "tool-call-start" => { - let tool = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.tool_calls.as_ref()).unwrap(); - let id = tool["id"].as_str().unwrap().to_string(); - let name = tool["function"]["name"].as_str().unwrap().to_string(); - let arguments = tool["function"]["arguments"].as_str().unwrap().to_string(); - pending_tool_call = Some(PendingToolCall { id: id.clone(), name: name.clone(), arguments: arguments.clone() }); - yield Ok(StreamPart::ToolInputStart { id: id.clone(), tool_name: name, provider_executed: None, dynamic: None, title: None, provider_metadata: None }); - if !arguments.is_empty() { yield Ok(StreamPart::ToolInputDelta { id, delta: arguments, provider_metadata: None }); } - } - "tool-call-delta" => { - if let Some(tool) = &mut pending_tool_call { - let delta = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.tool_calls.as_ref()).unwrap()["function"]["arguments"].as_str().unwrap().to_string(); - tool.arguments.push_str(&delta); - yield Ok(StreamPart::ToolInputDelta { id: tool.id.clone(), delta, provider_metadata: None }); - } - } - "tool-call-end" => { - if let Some(tool) = pending_tool_call.take() { - yield Ok(StreamPart::ToolInputEnd { id: tool.id.clone(), provider_metadata: None }); - let text = tool.arguments.trim(); - let input = match serde_json::from_str::(if text.is_empty() { "{}" } else { text }) { - Ok(value) if !contains_prototype_key(&value) => value.to_string(), - Ok(_) => { yield Err(AiMuxError::InvalidResponseData("Object contains forbidden prototype property".into())); return; } - Err(error) => { yield Err(AiMuxError::JsonParse(error.to_string())); return; } - }; - yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: tool.id, tool_name: tool.name, input, provider_executed: None, dynamic: None, provider_metadata: None })); - } - } - "message-end" => { - let delta = parsed.delta.unwrap(); - finish_reason = parse_finish_reason(delta.finish_reason.as_ref().unwrap()); - final_usage = convert_usage(delta.usage.as_ref().unwrap()); - } - _ => {} - } - } - yield Ok(StreamPart::Finish { finish_reason, usage: final_usage, provider_metadata: None }); - }; + let stream = pipe_through( + sse_stream, + CohereChatStream { + warnings: stream_warnings, + include_raw_chunks, + final_usage: Usage::default(), + finish_reason: FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }, + pending_tool_call: None, + is_reasoning: false, + }, + ); Ok(StreamResult { stream: Box::pin(stream), @@ -350,6 +272,234 @@ impl LanguageModel for CohereModel { } } +/// The `TransformStream` of `CohereChatLanguageModel.doStream`. +struct CohereChatStream { + warnings: Vec, + include_raw_chunks: bool, + final_usage: Usage, + finish_reason: FinishReason, + pending_tool_call: Option, + is_reasoning: bool, +} + +impl Transformer for CohereChatStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + let raw = match event { + Ok(raw) => raw, + Err(error) => { + if !error.is_recoverable_stream_error() { + controller.error(error); + return; + } + if self.include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: Value::Null, + }); + } + self.finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }; + controller.enqueue(StreamPart::Error { error }); + return; + } + }; + if self.include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: raw.clone(), + }); + } + let parsed = match StreamEvent::parse(raw) { + Ok(parsed) => parsed, + Err(error) => { + self.finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }; + controller.enqueue(StreamPart::Error { error }); + return; + } + }; + let id = parsed + .index + .map(|index| index.to_string()) + .unwrap_or_default(); + match parsed.event_type.as_str() { + "message-start" => controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: parsed.id, + timestamp: None, + model_id: None, + })), + "content-start" => { + let content = parsed + .delta + .as_ref() + .and_then(|delta| delta.message.as_ref()) + .and_then(|message| message.content.as_ref()) + .unwrap(); + if content["type"] == "thinking" { + self.is_reasoning = true; + controller.enqueue(StreamPart::ReasoningStart { + id, + provider_metadata: None, + }); + } else { + controller.enqueue(StreamPart::TextStart { + id, + provider_metadata: None, + }); + } + } + "content-delta" => { + let content = parsed + .delta + .as_ref() + .and_then(|delta| delta.message.as_ref()) + .and_then(|message| message.content.as_ref()) + .unwrap(); + if let Some(text) = content.get("text").and_then(Value::as_str) { + controller.enqueue(StreamPart::TextDelta { + id, + delta: text.into(), + provider_metadata: None, + }); + } else { + controller.enqueue(StreamPart::ReasoningDelta { + id, + delta: content["thinking"].as_str().unwrap().into(), + provider_metadata: None, + }); + } + } + "content-end" => { + if self.is_reasoning { + controller.enqueue(StreamPart::ReasoningEnd { + id, + provider_metadata: None, + }); + self.is_reasoning = false; + } else { + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: None, + }); + } + } + "tool-call-start" => { + let tool = parsed + .delta + .as_ref() + .and_then(|delta| delta.message.as_ref()) + .and_then(|message| message.tool_calls.as_ref()) + .unwrap(); + let id = tool["id"].as_str().unwrap().to_string(); + let name = tool["function"]["name"].as_str().unwrap().to_string(); + let arguments = tool["function"]["arguments"].as_str().unwrap().to_string(); + self.pending_tool_call = Some(PendingToolCall { + id: id.clone(), + name: name.clone(), + arguments: arguments.clone(), + }); + controller.enqueue(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: name, + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + }); + if !arguments.is_empty() { + controller.enqueue(StreamPart::ToolInputDelta { + id, + delta: arguments, + provider_metadata: None, + }); + } + } + "tool-call-delta" => { + if let Some(tool) = &mut self.pending_tool_call { + let delta = parsed + .delta + .as_ref() + .and_then(|delta| delta.message.as_ref()) + .and_then(|message| message.tool_calls.as_ref()) + .unwrap()["function"]["arguments"] + .as_str() + .unwrap() + .to_string(); + tool.arguments.push_str(&delta); + controller.enqueue(StreamPart::ToolInputDelta { + id: tool.id.clone(), + delta, + provider_metadata: None, + }); + } + } + "tool-call-end" => { + if let Some(tool) = self.pending_tool_call.take() { + controller.enqueue(StreamPart::ToolInputEnd { + id: tool.id.clone(), + provider_metadata: None, + }); + let text = tool.arguments.trim(); + let input = match serde_json::from_str::(if text.is_empty() { + "{}" + } else { + text + }) { + Ok(value) if !contains_prototype_key(&value) => value.to_string(), + Ok(_) => { + controller.error(AiMuxError::InvalidResponseData( + "Object contains forbidden prototype property".into(), + )); + return; + } + Err(error) => { + controller.error(AiMuxError::JsonParse(error.to_string())); + return; + } + }; + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: tool.id, + tool_name: tool.name, + input, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); + } + } + "message-end" => { + let delta = parsed.delta.unwrap(); + self.finish_reason = parse_finish_reason(delta.finish_reason.as_ref().unwrap()); + self.final_usage = convert_usage(delta.usage.as_ref().unwrap()); + } + _ => {} + } + } + + fn flush(self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason, + usage: self.final_usage, + provider_metadata: None, + }); + } +} + fn contains_prototype_key(value: &Value) -> bool { match value { Value::Object(object) => object.iter().any(|(key, value)| { From 824aa9beea5e34a42da96dba1e48453f954db201 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 13:16:11 +0800 Subject: [PATCH 14/16] refactor(providers): Responses streams as Transformer structs Convert the async_stream generators of the OpenAI, xAI, Hugging Face and Open Responses doStream paths into OpenAIResponsesStream, XaiResponsesStream, HuggingFaceResponsesStream and OpenResponsesStream driven by pipe_through. Emitted part sequences are unchanged. Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- aimux-providers/src/huggingface/responses.rs | 567 ++++--- aimux-providers/src/open_responses.rs | 603 +++---- .../src/openai/responses/responses_convert.rs | 1507 ++++++++++------- aimux-providers/src/xai/responses/mod.rs | 1201 +++++++------ 4 files changed, 2223 insertions(+), 1655 deletions(-) diff --git a/aimux-providers/src/huggingface/responses.rs b/aimux-providers/src/huggingface/responses.rs index 3b9792e04..3fa39e504 100644 --- a/aimux-providers/src/huggingface/responses.rs +++ b/aimux-providers/src/huggingface/responses.rs @@ -37,7 +37,8 @@ use aimux_core::tool::Tool; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage, Warning}; use aimux_provider_utils::{ - MediaTypeData, detect_media_type, get_top_level_media_type, is_full_media_type, + MediaTypeData, TransformStreamController, Transformer, detect_media_type, + get_top_level_media_type, is_full_media_type, pipe_through, }; use crate::shared::EndpointConfig; @@ -233,290 +234,338 @@ impl LanguageModel for HuggingFaceResponsesModel { let stream_response_headers = response_headers.clone(); let warnings = request.warnings.clone(); - let stream = async_stream::stream! { - // First part: StreamStart. - yield Ok(StreamPart::StreamStart { warnings }); - - let mut finish_reason = FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }; - let mut response_id: Option = None; - let mut usage_raw: Option = None; - let mut event_iter = futures::stream::iter(first_event.into_iter()).chain(sse_stream); - - while let Some(event) = event_iter.next().await { - match event { - Ok(parsed) => { - if let Some(error) = huggingface_stream_error( - &parsed, - &stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ) { - yield Ok(StreamPart::Error { error }); - break; + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + HuggingFaceResponsesStream { + warnings, + stream_error_url, + stream_request_body, + stream_response_headers, + finish_reason: FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }, + response_id: None, + usage_raw: None, + finished: false, + }, + ); + + 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 Hugging Face Responses `doStream`. +struct HuggingFaceResponsesStream { + warnings: Vec, + stream_error_url: String, + stream_request_body: Value, + stream_response_headers: HashMap, + finish_reason: FinishReason, + response_id: Option, + usage_raw: Option, + /// Set by an error chunk: later chunks are ignored, the stream still + /// finishes. + finished: bool, +} + +impl Transformer for HuggingFaceResponsesStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + if self.finished { + return; + } + let Self { + ref mut finish_reason, + ref mut response_id, + ref mut usage_raw, + .. + } = *self; + match event { + Ok(parsed) => { + if let Some(error) = huggingface_stream_error( + &parsed, + &self.stream_error_url, + self.stream_request_body.clone(), + self.stream_response_headers.clone(), + ) { + controller.enqueue(StreamPart::Error { error }); + self.finished = true; + return; + } + let chunk_type = parsed.get("type").and_then(|v| v.as_str()).unwrap_or(""); + + match chunk_type { + "response.created" => { + if let Some(resp) = parsed.get("response") { + *response_id = resp + .get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let created_at = resp + .get("created_at") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + let model = resp + .get("model") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: response_id.clone(), + timestamp: format_timestamp(created_at), + model_id: model, + })); } - let chunk_type = parsed - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or(""); + } - match chunk_type { - "response.created" => { - if let Some(resp) = parsed.get("response") { - response_id = resp - .get("id") + "response.output_item.added" => { + if let Some(item) = parsed.get("item") { + let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); + match item_type { + "message" => { + let role = + item.get("role").and_then(|v| v.as_str()).unwrap_or(""); + if role == "assistant" { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::TextStart { + id: id.clone(), + provider_metadata: Some( + super::options::huggingface_metadata( + json!({ "itemId": id }), + ), + ), + }); + } + } + "function_call" => { + let call_id = item + .get("call_id") .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let created_at = resp - .get("created_at") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0); - let model = resp - .get("model") + .unwrap_or("") + .to_string(); + let name = item + .get("name") .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: response_id.clone(), - timestamp: format_timestamp(created_at), - model_id: model, - })); + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::ToolInputStart { + id: call_id, + tool_name: name, + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + }); } - } - - "response.output_item.added" => { - if let Some(item) = parsed.get("item") { - let item_type = item - .get("type") + "reasoning" => { + let id = item + .get("id") .and_then(|v| v.as_str()) - .unwrap_or(""); - match item_type { - "message" => { - let role = item - .get("role") - .and_then(|v| v.as_str()) - .unwrap_or(""); - if role == "assistant" { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextStart { - id: id.clone(), - provider_metadata: Some(super::options::huggingface_metadata(json!({ "itemId": id }))), - }); - } - } - "function_call" => { - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ToolInputStart { - id: call_id, - tool_name: name, - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } - "reasoning" => { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningStart { - id: id.clone(), - provider_metadata: Some(super::options::huggingface_metadata(json!({ "itemId": id }))), - }); - } - _ => {} - } + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::ReasoningStart { + id: id.clone(), + provider_metadata: Some( + super::options::huggingface_metadata( + json!({ "itemId": id }), + ), + ), + }); } + _ => {} } + } + } - "response.output_item.done" => { - if let Some(item) = parsed.get("item") { - let item_type = item - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or(""); - match item_type { - "message" => { - let role = item - .get("role") - .and_then(|v| v.as_str()) - .unwrap_or(""); - if role == "assistant" { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); - } - } - "function_call" => { - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let arguments = item - .get("arguments") - .and_then(|v| v.as_str()) - .unwrap_or("{}"); - let input = arguments.to_string(); - - yield Ok(StreamPart::ToolInputEnd { - id: call_id.clone(), - provider_metadata: None, - }); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: call_id.clone(), - tool_name: name.clone(), - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - - if let Some(output) = - item.get("output").and_then(|v| v.as_str()) - { - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: call_id, - tool_name: name.clone(), - result: Value::String(output.to_string()), - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: None, - })); - } - } - _ => {} + "response.output_item.done" => { + if let Some(item) = parsed.get("item") { + let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); + match item_type { + "message" => { + let role = + item.get("role").and_then(|v| v.as_str()).unwrap_or(""); + if role == "assistant" { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: None, + }); } } - } - - "response.output_text.delta" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let delta = parsed - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextDelta { id: item_id, delta, provider_metadata: None}); - } - - "response.reasoning_text.delta" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let delta = parsed - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningDelta { - id: item_id, - delta, - provider_metadata: None, - }); - } - - "response.reasoning_text.done" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningEnd { id: item_id, - provider_metadata: None, - }); - } - - "response.completed" => { - if let Some(resp) = parsed.get("response") { - response_id = resp - .get("id") + "function_call" => { + let call_id = item + .get("call_id") .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - - let incomplete_reason = resp - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|r| r.as_str()); - - finish_reason = FinishReason { - unified: map_finish_reason( - incomplete_reason.unwrap_or("stop"), - ), - raw: incomplete_reason.map(std::string::ToString::to_string), - }; + .unwrap_or("") + .to_string(); + let name = item + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let arguments = item + .get("arguments") + .and_then(|v| v.as_str()) + .unwrap_or("{}"); + let input = arguments.to_string(); + + controller.enqueue(StreamPart::ToolInputEnd { + id: call_id.clone(), + provider_metadata: None, + }); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: call_id.clone(), + tool_name: name.clone(), + input, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); - if let Some(usage) = resp.get("usage") - && !usage.is_null() + if let Some(output) = + item.get("output").and_then(|v| v.as_str()) { - usage_raw = Some(usage.clone()); + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: call_id, + tool_name: name.clone(), + result: Value::String(output.to_string()), + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: None, + })); } } + _ => {} } + } + } - _ => { - // Unknown chunk type — ignore (matches the TS - // fallback schema which silently drops unknowns). + "response.output_text.delta" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let delta = parsed + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::TextDelta { + id: item_id, + delta, + provider_metadata: None, + }); + } + + "response.reasoning_text.delta" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let delta = parsed + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::ReasoningDelta { + id: item_id, + delta, + provider_metadata: None, + }); + } + + "response.reasoning_text.done" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::ReasoningEnd { + id: item_id, + provider_metadata: None, + }); + } + + "response.completed" => { + if let Some(resp) = parsed.get("response") { + *response_id = resp + .get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + + let incomplete_reason = resp + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|r| r.as_str()); + + *finish_reason = FinishReason { + unified: map_finish_reason(incomplete_reason.unwrap_or("stop")), + raw: incomplete_reason.map(std::string::ToString::to_string), + }; + + if let Some(usage) = resp.get("usage") + && !usage.is_null() + { + *usage_raw = Some(usage.clone()); } } } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; - } + + _ => { + // Unknown chunk type — ignore (matches the TS + // fallback schema which silently drops unknowns). } } } + Err(error) => { + if error.is_recoverable_stream_error() { + controller.enqueue_error(error); + } else { + controller.error(error); + } + } + } + } - let usage = convert_usage(usage_raw.as_ref()); - - yield Ok(StreamPart::Finish { - finish_reason, - usage, - provider_metadata: Some(super::options::huggingface_metadata(json!({ "responseId": response_id }))), - }); - }; + fn flush(self, controller: &mut TransformStreamController) { + let usage = convert_usage(self.usage_raw.as_ref()); - 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), - }), - }) + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason, + usage, + provider_metadata: Some(super::options::huggingface_metadata( + json!({ "responseId": self.response_id }), + )), + }); } } diff --git a/aimux-providers/src/open_responses.rs b/aimux-providers/src/open_responses.rs index d89dd6b04..4d3223a5c 100644 --- a/aimux-providers/src/open_responses.rs +++ b/aimux-providers/src/open_responses.rs @@ -33,7 +33,10 @@ use aimux_core::types::{ use aimux_core::embedding_model::EmbeddingModel; use aimux_core::image_model::ImageModel; use aimux_core::provider::Provider; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; +use aimux_provider_utils::{ + FetchFunction, HeaderMapOpt, Resolvable, TransformStreamController, Transformer, pipe_through, + validate_base_url, +}; use crate::shared::{Credential, EndpointConfig, EndpointHeaders, ProviderHeaders}; @@ -453,319 +456,347 @@ impl LanguageModel for OpenResponsesModel { )); } - let stream = async_stream::stream! { - // First part: StreamStart. - yield Ok(StreamPart::StreamStart { warnings }); + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + OpenResponsesStream { + warnings, + final_usage: Usage::default(), + has_tool_calls: false, + finish_reason: FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }, + is_active_reasoning: false, + tool_calls: HashMap::new(), + }, + ); - let mut final_usage = Usage::default(); - let mut has_tool_calls = false; - let mut finish_reason = FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }; - let mut is_active_reasoning = false; - - // Tool-call accumulators keyed by item_id. - let mut tool_calls: HashMap = HashMap::new(); - - let mut event_iter = - futures::stream::iter(first_event.into_iter()).chain(sse_stream); - - while let Some(event) = event_iter.next().await { - match event { - Ok(chunk) => { - let chunk_type = chunk - .get("type") - .and_then(|t| t.as_str()) - .unwrap_or(""); - - match chunk_type { - // -- Tool call / reasoning / message item added -- - "response.output_item.added" => { - if let Some(item) = chunk.get("item") { - let item_type = item - .get("type") - .and_then(|t| t.as_str()) - .unwrap_or(""); - match item_type { - "function_call" => { - let id = item - .get("id") + 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 `OpenResponsesLanguageModel.doStream`. +struct OpenResponsesStream { + warnings: Vec, + final_usage: Usage, + has_tool_calls: bool, + finish_reason: FinishReason, + is_active_reasoning: bool, + /// Tool-call accumulators keyed by item_id. + tool_calls: HashMap, +} + +impl Transformer for OpenResponsesStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + let Self { + ref mut final_usage, + ref mut has_tool_calls, + ref mut finish_reason, + ref mut is_active_reasoning, + ref mut tool_calls, + .. + } = *self; + match event { + Ok(chunk) => { + let chunk_type = chunk.get("type").and_then(|t| t.as_str()).unwrap_or(""); + + match chunk_type { + // -- Tool call / reasoning / message item added -- + "response.output_item.added" => { + if let Some(item) = chunk.get("item") { + let item_type = item.get("type").and_then(|t| t.as_str()).unwrap_or(""); + match item_type { + "function_call" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + tool_calls.insert( + id, + ToolCallAccum { + tool_name: item + .get("name") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - tool_calls.insert( - id, - ToolCallAccum { - tool_name: item - .get("name") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - tool_call_id: item - .get("call_id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - arguments: item - .get("arguments") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - }, - ); - } - "reasoning" => { - let id = item - .get("id") + .map(std::string::ToString::to_string), + tool_call_id: item + .get("call_id") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningStart { - id, - provider_metadata: None, - }); - is_active_reasoning = true; - } - "message" => { - let id = item - .get("id") + .map(std::string::ToString::to_string), + arguments: item + .get("arguments") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextStart { id, provider_metadata: None}); - } - _ => {} - } + .map(std::string::ToString::to_string), + }, + ); } - } - "response.function_call_arguments.delta" => { - let item_id = chunk - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let delta = chunk - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - tool_calls - .entry(item_id) - .and_modify(|tc| { - let existing = tc.arguments.take().unwrap_or_default(); - tc.arguments = Some(existing + &delta); - }) - .or_insert(ToolCallAccum { - tool_name: None, - tool_call_id: None, - arguments: Some(delta), + "reasoning" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::ReasoningStart { + id, + provider_metadata: None, }); - } - "response.function_call_arguments.done" => { - let item_id = chunk - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let arguments = chunk - .get("arguments") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - tool_calls - .entry(item_id) - .and_modify(|tc| { - tc.arguments = Some(arguments.clone()); - }) - .or_insert(ToolCallAccum { - tool_name: None, - tool_call_id: None, - arguments: Some(arguments), + *is_active_reasoning = true; + } + "message" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::TextStart { + id, + provider_metadata: None, }); + } + _ => {} } - "response.output_item.done" => { - if let Some(item) = chunk.get("item") { - let item_type = item - .get("type") - .and_then(|t| t.as_str()) - .unwrap_or(""); - match item_type { - "function_call" => { - let id = item - .get("id") + } + } + "response.function_call_arguments.delta" => { + let item_id = chunk + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let delta = chunk + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + tool_calls + .entry(item_id) + .and_modify(|tc| { + let existing = tc.arguments.take().unwrap_or_default(); + tc.arguments = Some(existing + &delta); + }) + .or_insert(ToolCallAccum { + tool_name: None, + tool_call_id: None, + arguments: Some(delta), + }); + } + "response.function_call_arguments.done" => { + let item_id = chunk + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let arguments = chunk + .get("arguments") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + tool_calls + .entry(item_id) + .and_modify(|tc| { + tc.arguments = Some(arguments.clone()); + }) + .or_insert(ToolCallAccum { + tool_name: None, + tool_call_id: None, + arguments: Some(arguments), + }); + } + "response.output_item.done" => { + if let Some(item) = chunk.get("item") { + let item_type = item.get("type").and_then(|t| t.as_str()).unwrap_or(""); + match item_type { + "function_call" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let accum = tool_calls.remove(&id); + let tool_name = accum + .as_ref() + .and_then(|a| a.tool_name.clone()) + .or_else(|| { + item.get("name") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let accum = tool_calls.remove(&id); - let tool_name = accum - .as_ref() - .and_then(|a| a.tool_name.clone()) - .or_else(|| { - item.get("name") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .unwrap_or_default(); - let tool_call_id = accum - .as_ref() - .and_then(|a| a.tool_call_id.clone()) - .or_else(|| { - item.get("call_id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .unwrap_or_default(); - let arguments = accum - .as_ref() - .and_then(|a| a.arguments.clone()) - .or_else(|| { - item.get("arguments") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .unwrap_or_default(); - let input = arguments; - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id, - tool_name, - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - has_tool_calls = true; - } - "reasoning" => { - let id = item - .get("id") + .map(std::string::ToString::to_string) + }) + .unwrap_or_default(); + let tool_call_id = accum + .as_ref() + .and_then(|a| a.tool_call_id.clone()) + .or_else(|| { + item.get("call_id") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningEnd { - id, - provider_metadata: None, - }); - is_active_reasoning = false; - } - "message" => { - let id = item - .get("id") + .map(std::string::ToString::to_string) + }) + .unwrap_or_default(); + let arguments = accum + .as_ref() + .and_then(|a| a.arguments.clone()) + .or_else(|| { + item.get("arguments") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); - } - _ => {} - } + .map(std::string::ToString::to_string) + }) + .unwrap_or_default(); + let input = arguments; + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id, + tool_name, + input, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); + *has_tool_calls = true; } + "reasoning" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::ReasoningEnd { + id, + provider_metadata: None, + }); + *is_active_reasoning = false; + } + "message" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: None, + }); + } + _ => {} } + } + } - // -- Reasoning text delta (LM Studio extension) -- - "response.reasoning_text.delta" => { - let id = chunk - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let delta = chunk - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningDelta { - id, - delta, - provider_metadata: None, - }); - } + // -- Reasoning text delta (LM Studio extension) -- + "response.reasoning_text.delta" => { + let id = chunk + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let delta = chunk + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::ReasoningDelta { + id, + delta, + provider_metadata: None, + }); + } - // -- Text delta -- - "response.output_text.delta" => { - let id = chunk - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let delta = chunk - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextDelta { id, delta, provider_metadata: None}); - } + // -- Text delta -- + "response.output_text.delta" => { + let id = chunk + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let delta = chunk + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::TextDelta { + id, + delta, + provider_metadata: None, + }); + } - // -- Completion events -- - "response.completed" | "response.incomplete" => { - if let Some(response) = chunk.get("response") { - let reason = response - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|r| r.as_str()); - finish_reason = FinishReason { - unified: map_open_responses_finish_reason( - reason, - has_tool_calls, - ), - raw: reason.map(std::string::ToString::to_string), - }; - if let Some(usage_val) = response.get("usage") { - final_usage = extract_usage_from_value(usage_val); - } - } - } - "response.failed" => { - if let Some(response) = chunk.get("response") { - let raw = response - .get("error") - .and_then(|e| e.get("code")) - .and_then(|c| c.as_str()) - .or_else(|| { - response.get("status").and_then(|s| s.as_str()) - }); - finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: raw.map(std::string::ToString::to_string), - }; - if let Some(usage_val) = response.get("usage") { - final_usage = extract_usage_from_value(usage_val); - } - } - } - _ => { - // Ignore unrecognised event types. + // -- Completion events -- + "response.completed" | "response.incomplete" => { + if let Some(response) = chunk.get("response") { + let reason = response + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|r| r.as_str()); + *finish_reason = FinishReason { + unified: map_open_responses_finish_reason(reason, *has_tool_calls), + raw: reason.map(std::string::ToString::to_string), + }; + if let Some(usage_val) = response.get("usage") { + *final_usage = extract_usage_from_value(usage_val); } } } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; + "response.failed" => { + if let Some(response) = chunk.get("response") { + let raw = response + .get("error") + .and_then(|e| e.get("code")) + .and_then(|c| c.as_str()) + .or_else(|| response.get("status").and_then(|s| s.as_str())); + *finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: raw.map(std::string::ToString::to_string), + }; + if let Some(usage_val) = response.get("usage") { + *final_usage = extract_usage_from_value(usage_val); + } } } + _ => { + // Ignore unrecognised event types. + } } } - - // Flush: close any dangling reasoning segment. - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); + Err(error) => { + if error.is_recoverable_stream_error() { + controller.enqueue_error(error); + } else { + controller.error(error); + } } + } + } - yield Ok(StreamPart::Finish { - finish_reason, - usage: final_usage, + fn flush(self, controller: &mut TransformStreamController) { + // Close any dangling reasoning segment. + if self.is_active_reasoning { + controller.enqueue(StreamPart::ReasoningEnd { + id: "reasoning-0".to_string(), provider_metadata: None, }); - }; + } - 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), - }), - }) + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason, + usage: self.final_usage, + provider_metadata: None, + }); } } diff --git a/aimux-providers/src/openai/responses/responses_convert.rs b/aimux-providers/src/openai/responses/responses_convert.rs index eef1f6eb1..d922f63bc 100644 --- a/aimux-providers/src/openai/responses/responses_convert.rs +++ b/aimux-providers/src/openai/responses/responses_convert.rs @@ -17,6 +17,7 @@ use aimux_core::tool::RawToolCall; use std::collections::{HashMap, HashSet, VecDeque}; use std::pin::Pin; +use aimux_provider_utils::{TransformStreamController, Transformer, pipe_through}; use futures::{Stream, StreamExt}; use serde_json::{Value, json}; @@ -617,688 +618,996 @@ where } } - let stream = async_stream::stream! { - // First part: StreamStart. - yield Ok(StreamPart::StreamStart { warnings }); - - // Ongoing tool calls keyed by output_index. - let mut ongoing_tool_calls: HashMap = HashMap::new(); - // Active reasoning items keyed by item_id. - let mut active_reasoning: HashMap = HashMap::new(); + let shell_provider_executed = hosted_shell(&request_body); + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + OpenAIResponsesStream { + warnings, + provider_key, + store_flag, + request_url, + request_body, + response_headers, + include_raw_chunks, + tool_names, + prompt_approval_tool_call_ids, + shell_provider_executed, + ongoing_tool_calls: HashMap::new(), + active_reasoning: HashMap::new(), + has_function_call: false, + final_usage: None, + final_raw_usage: None, + final_service_tier: None, + final_reasoning_context: None, + final_finish_reason: None, + response_id: None, + stream_errored: false, + text_states: HashMap::new(), + output_ids: HashMap::new(), + text_annotations: HashMap::new(), + logprobs: Vec::new(), + provider_items: HashMap::new(), + emitted_calls: HashSet::new(), + ended_inputs: HashSet::new(), + diff_received: HashSet::new(), + hosted_tool_search_ids: VecDeque::new(), + approval_tool_call_ids: HashMap::new(), + }, + ); + Ok(Box::pin(stream)) +} - let mut has_function_call = false; - let mut final_usage: Option = None; - let mut final_raw_usage: Option = None; - let mut final_service_tier: Option = None; - let mut final_reasoning_context: Option = None; - let mut final_finish_reason: Option = None; - let mut response_id: Option = None; - let mut stream_errored = false; - let mut text_states: HashMap = HashMap::new(); - let mut output_ids: HashMap = HashMap::new(); - let mut text_annotations: HashMap> = HashMap::new(); - let mut logprobs = Vec::new(); - let shell_provider_executed = hosted_shell(&request_body); - let mut provider_items: HashMap = HashMap::new(); - let mut emitted_calls = HashSet::new(); - let mut ended_inputs = HashSet::new(); - let mut diff_received = HashSet::new(); - let mut hosted_tool_search_ids = VecDeque::new(); - let mut approval_tool_call_ids: HashMap = HashMap::new(); +/// The `TransformStream` of the OpenAI Responses `doStream`. +struct OpenAIResponsesStream { + warnings: Vec, + provider_key: String, + store_flag: bool, + request_url: String, + request_body: Value, + response_headers: HashMap, + include_raw_chunks: bool, + tool_names: HashMap, + prompt_approval_tool_call_ids: HashMap, + shell_provider_executed: bool, + /// Ongoing tool calls keyed by output_index. + ongoing_tool_calls: HashMap, + /// Active reasoning items keyed by item_id. + active_reasoning: HashMap, + has_function_call: bool, + final_usage: Option, + final_raw_usage: Option, + final_service_tier: Option, + final_reasoning_context: Option, + final_finish_reason: Option, + response_id: Option, + stream_errored: bool, + text_states: HashMap, + output_ids: HashMap, + text_annotations: HashMap>, + logprobs: Vec, + provider_items: HashMap, + emitted_calls: HashSet, + ended_inputs: HashSet, + diff_received: HashSet, + hosted_tool_search_ids: VecDeque, + approval_tool_call_ids: HashMap, +} - let mut event_iter = - futures::stream::iter(first_event.into_iter()).chain(sse_stream); +impl Transformer for OpenAIResponsesStream { + type Input = Result; + type Output = StreamPart; - while let Some(event) = event_iter.next().await { - match event { - Ok(parsed) => { - if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); } - let etype = parsed - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } - match etype.as_str() { - // ── response.created → ResponseMetadata ─────────────────── - "response.created" => { - if let Some(resp_obj) = parsed.get("response") { - response_id = resp_obj - .get("id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let model_id = resp_obj - .get("model") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let timestamp = resp_obj - .get("created_at") - .and_then(serde_json::Value::as_u64) - .and_then(|secs| { - chrono::DateTime::from_timestamp(secs as i64, 0) - }) - .map(|dt| dt.to_rfc3339()); - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: response_id.clone(), - timestamp, - model_id, - })); - } - } + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + let Self { + ref provider_key, + store_flag, + ref request_url, + ref request_body, + ref response_headers, + include_raw_chunks, + ref tool_names, + ref prompt_approval_tool_call_ids, + shell_provider_executed, + ref mut ongoing_tool_calls, + ref mut active_reasoning, + ref mut has_function_call, + ref mut final_usage, + ref mut final_raw_usage, + ref mut final_service_tier, + ref mut final_reasoning_context, + ref mut final_finish_reason, + ref mut response_id, + ref mut stream_errored, + ref mut text_states, + ref mut output_ids, + ref mut text_annotations, + ref mut logprobs, + ref mut provider_items, + ref mut emitted_calls, + ref mut ended_inputs, + ref mut diff_received, + ref mut hosted_tool_search_ids, + ref mut approval_tool_call_ids, + .. + } = *self; + match event { + Ok(parsed) => { + if include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: parsed.clone(), + }); + } + let etype = parsed + .get("type") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); - // ── response.output_item.added ───────────────────────────── - "response.output_item.added" => { - let output_index = parsed - .get("output_index") + match etype.as_str() { + // ── response.created → ResponseMetadata ─────────────────── + "response.created" => { + if let Some(resp_obj) = parsed.get("response") { + *response_id = resp_obj + .get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let model_id = resp_obj + .get("model") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let timestamp = resp_obj + .get("created_at") .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - if let Some(item) = parsed.get("item") { - let item_type = item - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or(""); - match item_type { - + .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) + .map(|dt| dt.to_rfc3339()); + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: response_id.clone(), + timestamp, + model_id, + })); + } + } - "message" => { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - output_ids.insert(output_index, id.clone()); - text_states.insert(id.clone(), item.clone()); - yield Ok(StreamPart::TextStart { id, provider_metadata: Some(provider_namespace(&provider_key, text_metadata(item, None)).expect("provider metadata must be an object"))}); - } - "function_call" => { - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - ongoing_tool_calls.insert( - output_index, - OngoingToolCall { - tool_call_id: call_id.clone(), - }, - ); - yield Ok(StreamPart::ToolInputStart { - id: call_id, - tool_name: name, - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } - "custom_tool_call" => { - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - ongoing_tool_calls.insert( - output_index, - OngoingToolCall { - tool_call_id: call_id.clone(), - }, - ); - yield Ok(StreamPart::ToolInputStart { - id: call_id, - tool_name: name, - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } - "reasoning" => { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let encrypted = item - .get("encrypted_content") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - // Carry itemId + encrypted_content in - // provider_metadata (same shape as the - // non-streaming path) so reasoning can - // be echoed back on the next turn. - let meta = reasoning_stream_metadata( - &provider_key, - &id, - encrypted.as_deref(), - ); - active_reasoning.insert( - id.clone(), - ReasoningState { - encrypted_content: encrypted, - summary_parts: HashMap::from([( - 0usize, - SummaryStatus::Active, - )]), - }, - ); - yield Ok(StreamPart::ReasoningStart { - id: format!("{id}:0"), - provider_metadata: Some(meta), - }); - } - _ => { - provider_items.insert(output_index, item.clone()); - for part in provider_tool_start(item, &tool_names) { - if let StreamPart::ToolCall(call) = &part { emitted_calls.insert(call.tool_call_id.clone()); } - if let StreamPart::ToolInputEnd { id, .. } = &part { ended_inputs.insert(id.clone()); } - yield Ok(part); + // ── response.output_item.added ───────────────────────────── + "response.output_item.added" => { + let output_index = parsed + .get("output_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + if let Some(item) = parsed.get("item") { + let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); + match item_type { + "message" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + output_ids.insert(output_index, id.clone()); + text_states.insert(id.clone(), item.clone()); + controller.enqueue(StreamPart::TextStart { + id, + provider_metadata: Some( + provider_namespace( + provider_key, + text_metadata(item, None), + ) + .expect("provider metadata must be an object"), + ), + }); + } + "function_call" => { + let call_id = item + .get("call_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let name = item + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + ongoing_tool_calls.insert( + output_index, + OngoingToolCall { + tool_call_id: call_id.clone(), + }, + ); + controller.enqueue(StreamPart::ToolInputStart { + id: call_id, + tool_name: name, + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + }); + } + "custom_tool_call" => { + let call_id = item + .get("call_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let name = item + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + ongoing_tool_calls.insert( + output_index, + OngoingToolCall { + tool_call_id: call_id.clone(), + }, + ); + controller.enqueue(StreamPart::ToolInputStart { + id: call_id, + tool_name: name, + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + }); + } + "reasoning" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let encrypted = item + .get("encrypted_content") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + // Carry itemId + encrypted_content in + // provider_metadata (same shape as the + // non-streaming path) so reasoning can + // be echoed back on the next turn. + let meta = reasoning_stream_metadata( + provider_key, + &id, + encrypted.as_deref(), + ); + active_reasoning.insert( + id.clone(), + ReasoningState { + encrypted_content: encrypted, + summary_parts: HashMap::from([( + 0usize, + SummaryStatus::Active, + )]), + }, + ); + controller.enqueue(StreamPart::ReasoningStart { + id: format!("{id}:0"), + provider_metadata: Some(meta), + }); + } + _ => { + provider_items.insert(output_index, item.clone()); + for part in provider_tool_start(item, tool_names) { + if let StreamPart::ToolCall(call) = &part { + emitted_calls.insert(call.tool_call_id.clone()); } + if let StreamPart::ToolInputEnd { id, .. } = &part { + ended_inputs.insert(id.clone()); + } + controller.enqueue(part); } } } } + } - "response.image_generation_call.partial_image" => { - yield Ok(StreamPart::ToolResult(aimux_core::tool::ToolResult { - tool_call_id: parsed["item_id"].as_str().unwrap_or_default().to_owned(), - tool_name: tool_names.get("image_generation").map_or("image_generation", String::as_str).to_owned(), - result: json!({"result": parsed["partial_image_b64"]}), - preliminary: Some(true), is_error: None, dynamic: None, provider_metadata: None, - })); - } - "response.code_interpreter_call_code.delta" | "response.apply_patch_call_operation_diff.delta" => { - let index = parsed["output_index"].as_u64().unwrap_or(0) as usize; - if let Some(item) = provider_items.get(&index) { - let id = item.get("call_id").filter(|v| !v.is_null()).unwrap_or(&item["id"]).as_str().unwrap_or_default().to_owned(); - diff_received.insert(id.clone()); - yield Ok(StreamPart::ToolInputDelta { id, delta: escaped_delta(parsed["delta"].as_str().unwrap_or_default()), provider_metadata: None }); - } + "response.image_generation_call.partial_image" => { + controller.enqueue(StreamPart::ToolResult(aimux_core::tool::ToolResult { + tool_call_id: parsed["item_id"].as_str().unwrap_or_default().to_owned(), + tool_name: tool_names + .get("image_generation") + .map_or("image_generation", String::as_str) + .to_owned(), + result: json!({"result": parsed["partial_image_b64"]}), + preliminary: Some(true), + is_error: None, + dynamic: None, + provider_metadata: None, + })); + } + "response.code_interpreter_call_code.delta" + | "response.apply_patch_call_operation_diff.delta" => { + let index = parsed["output_index"].as_u64().unwrap_or(0) as usize; + if let Some(item) = provider_items.get(&index) { + let id = item + .get("call_id") + .filter(|v| !v.is_null()) + .unwrap_or(&item["id"]) + .as_str() + .unwrap_or_default() + .to_owned(); + diff_received.insert(id.clone()); + controller.enqueue(StreamPart::ToolInputDelta { + id, + delta: escaped_delta(parsed["delta"].as_str().unwrap_or_default()), + provider_metadata: None, + }); } - "response.code_interpreter_call_code.done" | "response.apply_patch_call_operation_diff.done" => { - let index = parsed["output_index"].as_u64().unwrap_or(0) as usize; - if let Some(item) = provider_items.get(&index) { - let id = item.get("call_id").filter(|v| !v.is_null()).unwrap_or(&item["id"]).as_str().unwrap_or_default().to_owned(); - if ended_inputs.insert(id.clone()) { - let code = etype == "response.code_interpreter_call_code.done"; - if !code && !diff_received.contains(&id) { - yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: escaped_delta(parsed["diff"].as_str().unwrap_or_default()), provider_metadata: None }); - } - yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: if code { "\"}" } else { "\"}}" }.to_owned(), provider_metadata: None }); - yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: None }); - if code { - let mut item = item.clone(); item["code"] = parsed["code"].clone(); - for part in super::provider_events::provider_tool_content(&item, &provider_key, &tool_names, shell_provider_executed, None) { - if let GenerateContent::ToolCall(call) = part { emitted_calls.insert(id.clone()); yield Ok(StreamPart::ToolCall(call)); } + } + "response.code_interpreter_call_code.done" + | "response.apply_patch_call_operation_diff.done" => { + let index = parsed["output_index"].as_u64().unwrap_or(0) as usize; + if let Some(item) = provider_items.get(&index) { + let id = item + .get("call_id") + .filter(|v| !v.is_null()) + .unwrap_or(&item["id"]) + .as_str() + .unwrap_or_default() + .to_owned(); + if ended_inputs.insert(id.clone()) { + let code = etype == "response.code_interpreter_call_code.done"; + if !code && !diff_received.contains(&id) { + controller.enqueue(StreamPart::ToolInputDelta { + id: id.clone(), + delta: escaped_delta( + parsed["diff"].as_str().unwrap_or_default(), + ), + provider_metadata: None, + }); + } + controller.enqueue(StreamPart::ToolInputDelta { + id: id.clone(), + delta: if code { "\"}" } else { "\"}}" }.to_owned(), + provider_metadata: None, + }); + controller.enqueue(StreamPart::ToolInputEnd { + id: id.clone(), + provider_metadata: None, + }); + if code { + let mut item = item.clone(); + item["code"] = parsed["code"].clone(); + for part in super::provider_events::provider_tool_content( + &item, + provider_key, + tool_names, + shell_provider_executed, + None, + ) { + if let GenerateContent::ToolCall(call) = part { + emitted_calls.insert(id.clone()); + controller.enqueue(StreamPart::ToolCall(call)); } } } } } + } - // ── response.output_text.delta → TextDelta ──────────────── - "response.output_text.delta" => { - if request_body.get("top_logprobs").and_then(Value::as_u64).is_some_and(|n| n > 0) && let Some(value) = parsed.get("logprobs") { logprobs.push(value.clone()); } - let mut id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - if let Some(index) = parsed.get("output_index").and_then(Value::as_u64) && let Some(original_id) = output_ids.get(&(index as usize)) { id = original_id.clone(); } - let delta = parsed - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextDelta { id, delta, provider_metadata: None}); + // ── response.output_text.delta → TextDelta ──────────────── + "response.output_text.delta" => { + if request_body + .get("top_logprobs") + .and_then(Value::as_u64) + .is_some_and(|n| n > 0) + && let Some(value) = parsed.get("logprobs") + { + logprobs.push(value.clone()); } + let mut id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if let Some(index) = parsed.get("output_index").and_then(Value::as_u64) + && let Some(original_id) = output_ids.get(&(index as usize)) + { + id = original_id.clone(); + } + let delta = parsed + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::TextDelta { + id, + delta, + provider_metadata: None, + }); + } - // ── function_call_arguments.delta → ToolInputDelta ──────── - "response.function_call_arguments.delta" - | "response.custom_tool_call_input.delta" => { - let output_index = parsed - .get("output_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - let delta = parsed - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - if let Some(tc) = ongoing_tool_calls.get(&output_index) { - yield Ok(StreamPart::ToolInputDelta { - id: tc.tool_call_id.clone(), - delta, - provider_metadata: None, + // ── function_call_arguments.delta → ToolInputDelta ──────── + "response.function_call_arguments.delta" + | "response.custom_tool_call_input.delta" => { + let output_index = parsed + .get("output_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + let delta = parsed + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if let Some(tc) = ongoing_tool_calls.get(&output_index) { + controller.enqueue(StreamPart::ToolInputDelta { + id: tc.tool_call_id.clone(), + delta, + provider_metadata: None, + }); + } + } + + // ── reasoning_summary_part.added ────────────────────────── + "response.reasoning_summary_part.added" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let summary_index = parsed + .get("summary_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + + if summary_index > 0 + && let Some(state) = active_reasoning.get_mut(&item_id) + { + let meta = reasoning_stream_metadata( + provider_key, + &item_id, + state.encrypted_content.as_deref(), + ); + // Conclude all 'can-conclude' parts. + let to_conclude: Vec = state + .summary_parts + .iter() + .filter(|(_, s)| **s == SummaryStatus::CanConclude) + .map(|(k, _)| *k) + .collect(); + for idx in to_conclude { + state.summary_parts.insert(idx, SummaryStatus::Concluded); + controller.enqueue(StreamPart::ReasoningEnd { + id: format!("{item_id}:{idx}"), + provider_metadata: Some(meta.clone()), }); } + state + .summary_parts + .insert(summary_index, SummaryStatus::Active); + controller.enqueue(StreamPart::ReasoningStart { + id: format!("{item_id}:{summary_index}"), + provider_metadata: Some(meta), + }); } + } - // ── reasoning_summary_part.added ────────────────────────── - "response.reasoning_summary_part.added" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let summary_index = parsed - .get("summary_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; + // ── reasoning_summary_text.delta → ReasoningDelta ───────── + "response.reasoning_summary_text.delta" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let summary_index = parsed + .get("summary_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + let delta = parsed + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::ReasoningDelta { + id: format!("{item_id}:{summary_index}"), + delta, + provider_metadata: None, + }); + } - if summary_index > 0 - && let Some(state) = active_reasoning.get_mut(&item_id) - { - let meta = reasoning_stream_metadata( - &provider_key, - &item_id, - state.encrypted_content.as_deref(), - ); - // Conclude all 'can-conclude' parts. - let to_conclude: Vec = state - .summary_parts - .iter() - .filter(|(_, s)| **s == SummaryStatus::CanConclude) - .map(|(k, _)| *k) - .collect(); - for idx in to_conclude { - state.summary_parts.insert(idx, SummaryStatus::Concluded); - yield Ok(StreamPart::ReasoningEnd { - id: format!("{item_id}:{idx}"), - provider_metadata: Some(meta.clone()), - }); - } + // ── reasoning_summary_part.done ─────────────────────────── + "response.reasoning_summary_part.done" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let summary_index = parsed + .get("summary_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + if let Some(state) = active_reasoning.get_mut(&item_id) { + if store_flag { state .summary_parts - .insert(summary_index, SummaryStatus::Active); - yield Ok(StreamPart::ReasoningStart { + .insert(summary_index, SummaryStatus::Concluded); + controller.enqueue(StreamPart::ReasoningEnd { id: format!("{item_id}:{summary_index}"), - provider_metadata: Some(meta), + provider_metadata: Some(reasoning_stream_metadata( + provider_key, + &item_id, + state.encrypted_content.as_deref(), + )), }); + } else { + state + .summary_parts + .insert(summary_index, SummaryStatus::CanConclude); } } + } - // ── reasoning_summary_text.delta → ReasoningDelta ───────── - "response.reasoning_summary_text.delta" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let summary_index = parsed - .get("summary_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - let delta = parsed - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningDelta { - id: format!("{item_id}:{summary_index}"), - delta, - provider_metadata: None, - }); - } - - // ── reasoning_summary_part.done ─────────────────────────── - "response.reasoning_summary_part.done" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let summary_index = parsed - .get("summary_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - if let Some(state) = active_reasoning.get_mut(&item_id) { - if store_flag { - state - .summary_parts - .insert(summary_index, SummaryStatus::Concluded); - yield Ok(StreamPart::ReasoningEnd { - id: format!("{item_id}:{summary_index}"), - provider_metadata: Some(reasoning_stream_metadata( - &provider_key, - &item_id, - state.encrypted_content.as_deref(), - )), + // ── response.output_item.done ───────────────────────────── + "response.output_item.done" => { + let output_index = parsed + .get("output_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + if let Some(item) = parsed.get("item") { + let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); + match item_type { + "message" => { + let mut id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if let Some(original_id) = output_ids.remove(&output_index) { + id = original_id; + } + let mut metadata_item = item.clone(); + metadata_item["id"] = json!(id); + if metadata_item.get("phase").is_none() + && let Some(phase) = + text_states.get(&id).and_then(|v| v.get("phase")) + { + metadata_item["phase"] = phase.clone(); + } + let annotations = text_annotations.remove(&id); + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: Some( + provider_namespace( + provider_key, + text_metadata(&metadata_item, annotations.as_ref()), + ) + .expect("provider metadata must be an object"), + ), }); - } else { - state - .summary_parts - .insert(summary_index, SummaryStatus::CanConclude); } - } - } - - // ── response.output_item.done ───────────────────────────── - "response.output_item.done" => { - let output_index = parsed - .get("output_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - if let Some(item) = parsed.get("item") { - let item_type = item - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or(""); - match item_type { - "message" => { - let mut id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - if let Some(original_id) = output_ids.remove(&output_index) { id = original_id; } - let mut metadata_item = item.clone(); - metadata_item["id"] = json!(id); - if metadata_item.get("phase").is_none() && let Some(phase) = text_states.get(&id).and_then(|v| v.get("phase")) { - metadata_item["phase"] = phase.clone(); + "function_call" => { + *has_function_call = true; + ongoing_tool_calls.remove(&output_index); + let call_id = item + .get("call_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let name = item + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let arguments = item + .get("arguments") + .and_then(|v| v.as_str()) + .unwrap_or("{}") + .to_string(); + controller.enqueue(StreamPart::ToolInputEnd { + id: call_id.clone(), + provider_metadata: None, + }); + let input = arguments; + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: call_id, + tool_name: name, + input, + provider_executed: None, + dynamic: None, + provider_metadata: Some( + provider_namespace(provider_key, tool_metadata(item)) + .expect("provider metadata must be an object"), + ), + })); + } + "custom_tool_call" => { + *has_function_call = true; + ongoing_tool_calls.remove(&output_index); + let call_id = item + .get("call_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let name = item + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let input_str = + item.get("input").and_then(|v| v.as_str()).unwrap_or("{}"); + controller.enqueue(StreamPart::ToolInputEnd { + id: call_id.clone(), + provider_metadata: None, + }); + let input = serde_json::to_string(input_str).expect( + "serializing a custom-tool input string cannot fail", + ); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: call_id, + tool_name: name, + input, + provider_executed: None, + dynamic: None, + provider_metadata: Some( + provider_namespace(provider_key, tool_metadata(item)) + .expect("provider metadata must be an object"), + ), + })); + } + "reasoning" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if let Some(state) = active_reasoning.get_mut(&id) { + state.encrypted_content = item + .get("encrypted_content") + .and_then(Value::as_str) + .map(str::to_owned); + let meta = reasoning_stream_metadata( + provider_key, + &id, + state.encrypted_content.as_deref(), + ); + // Conclude all active / can-conclude parts. + let to_conclude: Vec = state + .summary_parts + .iter() + .filter(|(_, s)| { + **s == SummaryStatus::Active + || **s == SummaryStatus::CanConclude + }) + .map(|(k, _)| *k) + .collect(); + for idx in to_conclude { + state + .summary_parts + .insert(idx, SummaryStatus::Concluded); + controller.enqueue(StreamPart::ReasoningEnd { + id: format!("{id}:{idx}"), + provider_metadata: Some(meta.clone()), + }); } - let annotations = text_annotations.remove(&id); - yield Ok(StreamPart::TextEnd { id, provider_metadata: Some(provider_namespace(&provider_key, text_metadata(&metadata_item, annotations.as_ref())).expect("provider metadata must be an object"))}); - } - "function_call" => { - has_function_call = true; - ongoing_tool_calls.remove(&output_index); - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let arguments = item - .get("arguments") - .and_then(|v| v.as_str()) - .unwrap_or("{}") - .to_string(); - yield Ok(StreamPart::ToolInputEnd { - id: call_id.clone(), - provider_metadata: None, - }); - let input = arguments; - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: call_id, - tool_name: name, - input, - provider_executed: None, - dynamic: None, - provider_metadata: Some(provider_namespace(&provider_key, tool_metadata(item)).expect("provider metadata must be an object")), - })); } - "custom_tool_call" => { - has_function_call = true; - ongoing_tool_calls.remove(&output_index); - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let input_str = item - .get("input") - .and_then(|v| v.as_str()) - .unwrap_or("{}"); - yield Ok(StreamPart::ToolInputEnd { - id: call_id.clone(), - provider_metadata: None, - }); - let input = - serde_json::to_string(input_str).expect( - "serializing a custom-tool input string cannot fail", - ); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: call_id, - tool_name: name, - input, - provider_executed: None, - dynamic: None, - provider_metadata: Some(provider_namespace(&provider_key, tool_metadata(item)).expect("provider metadata must be an object")), - })); + active_reasoning.remove(&id); + } + _ => { + let mut item = item.clone(); + if item_type == "local_shell_call" + && let Some(action) = item["action"].as_object_mut() + { + for (from, to) in [ + ("timeout_ms", "timeoutMs"), + ("working_directory", "workingDirectory"), + ] { + if let Some(value) = action.remove(from) { + action.insert(to.into(), value); + } + } } - "reasoning" => { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - if let Some(state) = active_reasoning.get_mut(&id) { - state.encrypted_content = item.get("encrypted_content").and_then(Value::as_str).map(str::to_owned); - let meta = reasoning_stream_metadata( - &provider_key, - &id, - state.encrypted_content.as_deref(), - ); - // Conclude all active / can-conclude parts. - let to_conclude: Vec = state - .summary_parts - .iter() - .filter(|(_, s)| { - **s == SummaryStatus::Active - || **s == SummaryStatus::CanConclude - }) - .map(|(k, _)| *k) - .collect(); - for idx in to_conclude { - state - .summary_parts - .insert(idx, SummaryStatus::Concluded); - yield Ok(StreamPart::ReasoningEnd { - id: format!("{id}:{idx}"), - provider_metadata: Some(meta.clone()), + if item_type == "apply_patch_call" { + let id = + item["call_id"].as_str().unwrap_or_default().to_owned(); + if ended_inputs.insert(id.clone()) + && item["operation"]["type"] != "delete_file" + { + if !diff_received.contains(&id) { + controller.enqueue(StreamPart::ToolInputDelta { + id: id.clone(), + delta: escaped_delta( + item["operation"]["diff"] + .as_str() + .unwrap_or_default(), + ), + provider_metadata: None, }); } + controller.enqueue(StreamPart::ToolInputDelta { + id: id.clone(), + delta: "\"}}".into(), + provider_metadata: None, + }); + controller.enqueue(StreamPart::ToolInputEnd { + id, + provider_metadata: None, + }); } - active_reasoning.remove(&id); } - _ => { - let mut item = item.clone(); - if item_type == "local_shell_call" && let Some(action) = item["action"].as_object_mut() { - for (from, to) in [("timeout_ms", "timeoutMs"), ("working_directory", "workingDirectory")] { - if let Some(value) = action.remove(from) { action.insert(to.into(), value); } - } - } - if item_type == "apply_patch_call" { - let id = item["call_id"].as_str().unwrap_or_default().to_owned(); - if ended_inputs.insert(id.clone()) && item["operation"]["type"] != "delete_file" { - if !diff_received.contains(&id) { - yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: escaped_delta(item["operation"]["diff"].as_str().unwrap_or_default()), provider_metadata: None }); + if item_type == "tool_search_call" + && item["execution"] == "server" + { + item["call_id"] = Value::Null; + } + pair_tool_search_id(&mut item, hosted_tool_search_ids); + let mcp_tool_call_id = item + .get("approval_request_id") + .and_then(Value::as_str) + .and_then(|id| { + approval_tool_call_ids + .get(id) + .or_else(|| prompt_approval_tool_call_ids.get(id)) + }) + .cloned(); + for part in super::provider_events::provider_tool_content( + &item, + provider_key, + tool_names, + shell_provider_executed, + mcp_tool_call_id.as_deref(), + ) { + match part { + GenerateContent::ToolCall(call) => { + if item_type == "code_interpreter_call" + || (item_type == "apply_patch_call" + && item["status"] != "completed") + { + continue; } - yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: "\"}}".into(), provider_metadata: None }); - yield Ok(StreamPart::ToolInputEnd { id, provider_metadata: None }); - } - } - if item_type == "tool_search_call" && item["execution"] == "server" { - item["call_id"] = Value::Null; - } - pair_tool_search_id(&mut item, &mut hosted_tool_search_ids); - let mcp_tool_call_id = item.get("approval_request_id").and_then(Value::as_str) - .and_then(|id| approval_tool_call_ids.get(id).or_else(|| prompt_approval_tool_call_ids.get(id))).cloned(); - for part in super::provider_events::provider_tool_content(&item, &provider_key, &tool_names, shell_provider_executed, mcp_tool_call_id.as_deref()) { - match part { - GenerateContent::ToolCall(call) => { - if item_type == "code_interpreter_call" || (item_type == "apply_patch_call" && item["status"] != "completed") { continue; } - if emitted_calls.insert(call.tool_call_id.clone()) || item_type == "mcp_call" { - has_function_call |= call.provider_executed != Some(true) && matches!(item_type, "computer_call" | "apply_patch_call"); - if matches!(item_type, "computer_call" | "apply_patch_call" | "tool_search_call") { - if item_type == "tool_search_call" && item["execution"] != "server" { - yield Ok(StreamPart::ToolInputStart { id: call.tool_call_id.clone(), tool_name: call.tool_name.clone(), provider_executed: call.provider_executed, dynamic: None, title: None, provider_metadata: None }); - } - if !ended_inputs.contains(&call.tool_call_id) { - if item_type == "computer_call" && !call.input.is_empty() { yield Ok(StreamPart::ToolInputDelta { id: call.tool_call_id.clone(), delta: call.input.clone(), provider_metadata: None }); } - yield Ok(StreamPart::ToolInputEnd { id: call.tool_call_id.clone(), provider_metadata: None }); + if emitted_calls.insert(call.tool_call_id.clone()) + || item_type == "mcp_call" + { + *has_function_call |= call.provider_executed + != Some(true) + && matches!( + item_type, + "computer_call" | "apply_patch_call" + ); + if matches!( + item_type, + "computer_call" + | "apply_patch_call" + | "tool_search_call" + ) { + if item_type == "tool_search_call" + && item["execution"] != "server" + { + controller.enqueue( + StreamPart::ToolInputStart { + id: call.tool_call_id.clone(), + tool_name: call + .tool_name + .clone(), + provider_executed: call + .provider_executed, + dynamic: None, + title: None, + provider_metadata: None, + }, + ); + } + if !ended_inputs + .contains(&call.tool_call_id) + { + if item_type == "computer_call" + && !call.input.is_empty() + { + controller.enqueue( + StreamPart::ToolInputDelta { + id: call + .tool_call_id + .clone(), + delta: call.input.clone(), + provider_metadata: None, + }, + ); } + controller.enqueue( + StreamPart::ToolInputEnd { + id: call.tool_call_id.clone(), + provider_metadata: None, + }, + ); } - yield Ok(StreamPart::ToolCall(call)); } + controller.enqueue(StreamPart::ToolCall(call)); } - GenerateContent::ToolResult(result) => { yield Ok(StreamPart::ToolResult(result)); } - GenerateContent::ToolApprovalRequest(approval) => { - approval_tool_call_ids.insert(approval.approval_id.clone(), approval.tool_call_id.clone()); - yield Ok(StreamPart::ToolApprovalRequest(approval)); - } - GenerateContent::Custom { kind, provider_metadata } => { - yield Ok(StreamPart::Custom { kind, provider_metadata }); - } - _ => {} } + GenerateContent::ToolResult(result) => { + controller.enqueue(StreamPart::ToolResult(result)); + } + GenerateContent::ToolApprovalRequest(approval) => { + approval_tool_call_ids.insert( + approval.approval_id.clone(), + approval.tool_call_id.clone(), + ); + controller.enqueue( + StreamPart::ToolApprovalRequest(approval), + ); + } + GenerateContent::Custom { + kind, + provider_metadata, + } => { + controller.enqueue(StreamPart::Custom { + kind, + provider_metadata, + }); + } + _ => {} } - provider_items.remove(&output_index); } + provider_items.remove(&output_index); } } } + } - // ── response.output_text.annotation.added → Source ───────── - "response.output_text.annotation.added" => { - match annotation_source(parsed.get("annotation").unwrap_or(&Value::Null), &provider_key) { - Ok(source) => { - let id = parsed.get("item_id").and_then(Value::as_str).unwrap_or_default().to_string(); - text_annotations.entry(id).or_default().push(parsed["annotation"].clone()); - yield Ok(StreamPart::Source(source)); - } - Err(error) => { - stream_errored = true; - yield Ok(StreamPart::Error { error }); - } + // ── response.output_text.annotation.added → Source ───────── + "response.output_text.annotation.added" => { + match annotation_source( + parsed.get("annotation").unwrap_or(&Value::Null), + provider_key, + ) { + Ok(source) => { + let id = parsed + .get("item_id") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(); + text_annotations + .entry(id) + .or_default() + .push(parsed["annotation"].clone()); + controller.enqueue(StreamPart::Source(source)); + } + Err(error) => { + *stream_errored = true; + controller.enqueue(StreamPart::Error { error }); } } + } - // ── response.completed / response.incomplete → finish ──── - "response.completed" | "response.incomplete" => { - if let Some(resp_obj) = parsed.get("response") { - let reason = resp_obj - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|v| v.as_str()); - if !stream_errored { - final_finish_reason = Some(map_responses_finish_reason(reason, has_function_call)); - } - final_usage = - resp_obj.get("usage").and_then(parse_usage); - // Keep the raw wire usage object for `usage.raw` - // (RFC-0015 P0-3) and capture the finish-time - // provider metadata fields carried on the - // terminal response object. - final_raw_usage = resp_obj.get("usage").cloned(); - if let Some(st) = resp_obj - .get("service_tier") - .and_then(|v| v.as_str()) - { - final_service_tier = Some(st.to_string()); - } - if let Some(reasoning) = resp_obj.get("reasoning") - && let Some(ctx) = reasoning.get("context") - && !ctx.is_null() - { - final_reasoning_context = Some(ctx.clone()); - } + // ── response.completed / response.incomplete → finish ──── + "response.completed" | "response.incomplete" => { + if let Some(resp_obj) = parsed.get("response") { + let reason = resp_obj + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|v| v.as_str()); + if !*stream_errored { + *final_finish_reason = + Some(map_responses_finish_reason(reason, *has_function_call)); + } + *final_usage = resp_obj.get("usage").and_then(parse_usage); + // Keep the raw wire usage object for `usage.raw` + // (RFC-0015 P0-3) and capture the finish-time + // provider metadata fields carried on the + // terminal response object. + *final_raw_usage = resp_obj.get("usage").cloned(); + if let Some(st) = resp_obj.get("service_tier").and_then(|v| v.as_str()) + { + *final_service_tier = Some(st.to_string()); + } + if let Some(reasoning) = resp_obj.get("reasoning") + && let Some(ctx) = reasoning.get("context") + && !ctx.is_null() + { + *final_reasoning_context = Some(ctx.clone()); } } + } - // ── response.failed ─────────────────────────────────────── - "response.failed" => { - if let Some(resp_obj) = parsed.get("response") { - let reason = resp_obj - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|v| v.as_str()); - final_finish_reason = Some(match reason { - Some(r) => map_responses_finish_reason(Some(r), has_function_call), - None => FinishReason { - unified: FinishReasonUnified::Error, - raw: Some("error".to_string()), - }, + // ── response.failed ─────────────────────────────────────── + "response.failed" => { + if let Some(resp_obj) = parsed.get("response") { + let reason = resp_obj + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|v| v.as_str()); + *final_finish_reason = Some(match reason { + Some(r) => map_responses_finish_reason(Some(r), *has_function_call), + None => FinishReason { + unified: FinishReasonUnified::Error, + raw: Some("error".to_string()), + }, + }); + *final_usage = resp_obj.get("usage").and_then(parse_usage); + *final_raw_usage = resp_obj.get("usage").cloned(); + // Upstream's `response.failed` arm carries + // `reasoningContext` (but not `serviceTier`). + if let Some(reasoning) = resp_obj.get("reasoning") + && let Some(ctx) = reasoning.get("context") + && !ctx.is_null() + { + *final_reasoning_context = Some(ctx.clone()); + } + if !*stream_errored && resp_obj.get("error").is_some() { + *stream_errored = true; + controller.enqueue(StreamPart::Error { + error: super::super::openai_stream_error( + &parsed, + request_url, + request_body.clone(), + response_headers.clone(), + ), }); - final_usage = - resp_obj.get("usage").and_then(parse_usage); - final_raw_usage = resp_obj.get("usage").cloned(); - // Upstream's `response.failed` arm carries - // `reasoningContext` (but not `serviceTier`). - if let Some(reasoning) = resp_obj.get("reasoning") - && let Some(ctx) = reasoning.get("context") - && !ctx.is_null() - { - final_reasoning_context = Some(ctx.clone()); - } - if !stream_errored - && resp_obj.get("error").is_some() - { - stream_errored = true; - yield Ok(StreamPart::Error { error: super::super::openai_stream_error(&parsed, &request_url, request_body.clone(), response_headers.clone()) }); - - } } } + } - // ── error chunk ─────────────────────────────────────────── - "error" => { - stream_errored = true; - final_finish_reason = Some(FinishReason { - unified: FinishReasonUnified::Error, - raw: Some("error".to_string()), - }); - yield Ok(StreamPart::Error { error: super::super::openai_stream_error(&parsed, &request_url, request_body.clone(), response_headers.clone()) }); - } + // ── error chunk ─────────────────────────────────────────── + "error" => { + *stream_errored = true; + *final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::Error, + raw: Some("error".to_string()), + }); + controller.enqueue(StreamPart::Error { + error: super::super::openai_stream_error( + &parsed, + request_url, + request_body.clone(), + response_headers.clone(), + ), + }); + } - _ => { - // Unknown / unhandled chunk types (content_part.added, - // output_text.done, content_part.done, etc.) are - // ignored — the core streaming path is driven by - // output_item.added/done and the delta events. - } + _ => { + // Unknown / unhandled chunk types (content_part.added, + // output_text.done, content_part.done, etc.) are + // ignored — the core streaming path is driven by + // output_item.added/done and the delta events. } } - Err(error) => { - stream_errored = true; - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - yield Ok(StreamPart::Error { error }); - } + } + Err(error) => { + *stream_errored = true; + *final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }); + controller.enqueue(StreamPart::Error { error }); } } + } + fn flush(self, controller: &mut TransformStreamController) { // Build provider metadata for the Finish part: { : { // responseId, serviceTier?, reasoningContext? } }. Mirrors the TS // flush() providerMetadata, which carries `service_tier` and // `reasoning.context` from the terminal response object. - let mut pm = json!({ "responseId": response_id.unwrap_or_default() }); - if let Some(st) = final_service_tier { + let mut pm = json!({ "responseId": self.response_id.unwrap_or_default() }); + if let Some(st) = self.final_service_tier { pm["serviceTier"] = json!(st); } - if let Some(ctx) = final_reasoning_context { + if let Some(ctx) = self.final_reasoning_context { pm["reasoningContext"] = ctx; } - if !logprobs.is_empty() { pm["logprobs"] = json!(logprobs); } - let provider_metadata = Some(provider_namespace(&provider_key, pm).expect("provider metadata must be an object")); + if !self.logprobs.is_empty() { + pm["logprobs"] = json!(self.logprobs); + } + let provider_metadata = Some( + provider_namespace(&self.provider_key, pm) + .expect("provider metadata must be an object"), + ); - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, raw: None, + controller.enqueue(StreamPart::Finish { + finish_reason: self.final_finish_reason.unwrap_or(FinishReason { + unified: FinishReasonUnified::Other, + raw: None, }), - usage: convert_responses_usage(final_usage.as_ref(), final_raw_usage), + usage: convert_responses_usage(self.final_usage.as_ref(), self.final_raw_usage), provider_metadata, }); - }; - Ok(Box::pin(stream)) + } } fn pair_tool_search_id(item: &mut Value, hosted_ids: &mut VecDeque) { diff --git a/aimux-providers/src/xai/responses/mod.rs b/aimux-providers/src/xai/responses/mod.rs index 84a614b78..9555f28f5 100644 --- a/aimux-providers/src/xai/responses/mod.rs +++ b/aimux-providers/src/xai/responses/mod.rs @@ -20,6 +20,7 @@ use aimux_core::tool::ToolResult; use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; +use aimux_provider_utils::{TransformStreamController, Transformer, pipe_through}; use async_trait::async_trait; use futures::StreamExt; use serde_json::{Value, json}; @@ -525,369 +526,529 @@ impl LanguageModel for XaiResponsesModel { let stream_response_headers = response_headers.clone(); let include_raw_chunks = options.include_raw_chunks == Some(true); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); - - let mut final_usage: Option = None; - let mut final_finish_reason: FinishReason = FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }; - let mut cost_in_usd_ticks: Option = None; - let mut service_tier: Option = None; - let mut prompt_cache_key: Option = None; - let mut safety_identifier: Option = None; - let mut has_function_call = false; - let mut is_first_chunk = true; - - // Track content blocks. - let mut content_blocks: Vec = Vec::new(); - let mut seen_tool_calls: std::collections::HashSet = std::collections::HashSet::new(); - let mut active_reasoning: HashMap = HashMap::new(); - - // Track ongoing function calls by output_index. - let mut ongoing_tool_calls: HashMap = HashMap::new(); // output_index -> (tool_call_id, tool_name) - - let mut event_iter = futures::stream::iter(first_event.into_iter()).chain(sse_stream); - while let Some(event) = event_iter.next().await { - match event { - Ok(parsed) => { - if include_raw_chunks { - yield Ok(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 is_first_chunk { - is_first_chunk = false; - 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| { - chrono::DateTime::from_timestamp(c as i64, 0).map(|dt| dt.to_rfc3339()) - }); - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { id, timestamp, model_id: model })); - } - continue; - } + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + XaiResponsesStream { + warnings, + include_raw_chunks, + provider_tool_names, + stream_error_url, + stream_request_body, + stream_response_headers, + final_usage: None, + final_finish_reason: FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }, + cost_in_usd_ticks: None, + service_tier: None, + prompt_cache_key: None, + safety_identifier: None, + has_function_call: false, + is_first_chunk: true, + content_blocks: Vec::new(), + seen_tool_calls: std::collections::HashSet::new(), + active_reasoning: HashMap::new(), + ongoing_tool_calls: HashMap::new(), + finished: false, + }, + ); - // ── 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(), ()); - yield Ok(StreamPart::ReasoningStart { - id: block_id, - provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), - }); - } - continue; - } + 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), + }), + }) + } +} - // ── 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}"); - yield Ok(StreamPart::ReasoningDelta { - id: block_id, - delta: delta.to_string(), - provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), - }); - continue; - } +/// The `TransformStream` of the xAI Responses `doStream`. +struct XaiResponsesStream { + warnings: Vec, + include_raw_chunks: bool, + provider_tool_names: HashMap, + stream_error_url: String, + stream_request_body: Value, + stream_response_headers: HashMap, + final_usage: Option, + final_finish_reason: FinishReason, + cost_in_usd_ticks: Option, + service_tier: Option, + prompt_cache_key: Option, + safety_identifier: Option, + has_function_call: bool, + is_first_chunk: 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, +} - // ── reasoning_summary_text.done ── - if event_type == "response.reasoning_summary_text.done" { - continue; - } +impl Transformer for XaiResponsesStream { + type Input = Result; + type Output = StreamPart; - // ── 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(), ()); - yield Ok(StreamPart::ReasoningStart { - id: block_id.clone(), - provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), - }); - } - yield Ok(StreamPart::ReasoningDelta { - id: block_id, - delta: delta.to_string(), - provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), - }); - continue; - } + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } - // ── reasoning_text.done ── - if event_type == "response.reasoning_text.done" { - continue; - } + fn transform( + &mut self, + event: Result, + controller: &mut TransformStreamController, + ) { + if self.finished { + return; + } + let Self { + include_raw_chunks, + ref provider_tool_names, + ref stream_error_url, + ref stream_request_body, + ref stream_response_headers, + 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 is_first_chunk, + 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 *is_first_chunk { + *is_first_chunk = false; + 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| { + chrono::DateTime::from_timestamp(c as i64, 0).map(|dt| dt.to_rfc3339()) + }); + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id, + timestamp, + model_id: model, + })); + } + 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()); - yield Ok(StreamPart::TextStart { id: block_id.clone(), provider_metadata: None}); - } - yield Ok(StreamPart::TextDelta { - id: block_id, - delta: delta.to_string(), - provider_metadata: None, - }); - continue; - } + // ── 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; + } - // ── 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(); - yield Ok(StreamPart::Source(Source::Url { - id: generate_source_id(), - url: url.to_string(), - title: Some(title), - provider_metadata: None, - })); - } - } - } - continue; - } + // ── 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.enqueue(StreamPart::TextStart { + id: block_id.clone(), + provider_metadata: None, + }); + } + controller.enqueue(StreamPart::TextDelta { + id: block_id, + delta: delta.to_string(), + 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(); - yield Ok(StreamPart::Source(Source::Url { + // ── 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, })); } - continue; } + } + 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; - } + // ── 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; + } - 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 = 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 = FinishReason { - unified: if has_function_call { - FinishReasonUnified::ToolCalls - } else { - map_xai_responses_finish_reason(status) - }, - raw: Some(status.to_string()), - }; - } - continue; - } + // ── 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; + } - // ── 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 = 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()) { - yield Ok(StreamPart::Error { - error: super::xai_stream_error( - &parsed, - &stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ), - }); - } - continue; - } + *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 = 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 = FinishReason { + unified: if *has_function_call { + FinishReasonUnified::ToolCalls + } else { + map_xai_responses_finish_reason(status) + }, + raw: Some(status.to_string()), + }; + } + return; + } - // ── error event ── - if event_type == "error" { - yield Ok(StreamPart::Error { - error: super::xai_stream_error( - &parsed, - &stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ), - }); - final_finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: Some("error".to_string()), - }; - break; - } + // ── 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 = 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; + } - // ── custom_tool_call_input.delta / .done ── - if event_type == "response.custom_tool_call_input.delta" - || event_type == "response.custom_tool_call_input.done" - { - continue; - } + // ── 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 = FinishReason { + unified: FinishReasonUnified::Error, + raw: Some("error".to_string()), + }; + *finished = true; + 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) { - yield Ok(StreamPart::ToolInputDelta { - id: tool_call_id.clone(), - delta: delta.to_string(), - provider_metadata: None, + // ── 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.enqueue(StreamPart::ToolInputDelta { + id: tool_call_id.clone(), + delta: delta.to_string(), + provider_metadata: None, + }); + } + 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 }) + }, + )), }); } - continue; - } - // ── function_call_arguments.done ── - if event_type == "response.function_call_arguments.done" { - continue; + 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; + } - // ── 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(), ()); - yield Ok(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 }) })), - }); - } + // ── 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.enqueue(StreamPart::ToolInputStart { + id: part_id.to_string(), + tool_name: tool_name.clone(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + }); + controller.enqueue(StreamPart::ToolInputDelta { + id: part_id.to_string(), + delta: input.to_string(), + provider_metadata: None, + }); + controller.enqueue(StreamPart::ToolInputEnd { + id: part_id.to_string(), + provider_metadata: None, + }); + 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, + })); + } - let mut meta = json!({}); - if !part_id.is_empty() { - meta["itemId"] = json!(part_id); - } - if let Some(ec) = encrypted { - meta["reasoningEncryptedContent"] = json!(ec); - } - yield Ok(StreamPart::ReasoningEnd { - id: block_id, - provider_metadata: Some(crate::xai::options::xai_metadata(json!(meta))), - }); - active_reasoning.remove(part_id); - } - continue; + if event_type == "response.output_item.done" { + if item_type == "image_generation_call" { + controller.enqueue(StreamPart::ToolResult( + image_generation_result(&item, tool_name), + )); + 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()); - yield Ok(StreamPart::ToolInputStart { - id: part_id.to_string(), - tool_name: tool_name.clone(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - yield Ok(StreamPart::ToolInputDelta { - id: part_id.to_string(), - delta: input.to_string(), - provider_metadata: None, - }); - yield Ok(StreamPart::ToolInputEnd { - id: part_id.to_string(), - provider_metadata: None, - }); - yield Ok(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" { - yield Ok(StreamPart::ToolResult(image_generation_result(&item, tool_name))); - continue; - } - 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| { + 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), @@ -895,204 +1056,222 @@ impl LanguageModel for XaiResponsesModel { "text": r.get("text").cloned().unwrap_or(Value::Null), })).collect::>()) }).unwrap_or(Value::Null); - yield Ok(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, - })); - } - continue; - } + 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(""); + // ── 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.enqueue(StreamPart::ToolInputStart { + id: part_id.to_string(), + tool_name: tool_name.clone(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + }); + controller.enqueue(StreamPart::ToolInputDelta { + id: part_id.to_string(), + delta: tool_input.clone(), + provider_metadata: None, + }); + controller.enqueue(StreamPart::ToolInputEnd { + id: part_id.to_string(), + provider_metadata: None, + }); + 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, + })); + } - let should_emit = if item_type == "custom_tool_call" { - event_type == "response.output_item.done" + 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 { - !seen_tool_calls.contains(part_id) - }; - - if should_emit && !seen_tool_calls.contains(part_id) { - seen_tool_calls.insert(part_id.to_string()); - yield Ok(StreamPart::ToolInputStart { - id: part_id.to_string(), - tool_name: tool_name.clone(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - yield Ok(StreamPart::ToolInputDelta { - id: part_id.to_string(), - delta: tool_input.clone(), + 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.enqueue(StreamPart::TextStart { + id: block_id.clone(), provider_metadata: None, }); - yield Ok(StreamPart::ToolInputEnd { - id: part_id.to_string(), + controller.enqueue(StreamPart::TextDelta { + id: block_id, + delta: text.to_string(), provider_metadata: None, }); - yield Ok(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" { - yield Ok(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, - })); } - continue; } - - // ── 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()); - yield Ok(StreamPart::TextStart { id: block_id.clone(), provider_metadata: None}); - yield Ok(StreamPart::TextDelta { - id: block_id, - delta: 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(); - yield Ok(StreamPart::Source(Source::Url { - id: generate_source_id(), - url: url.to_string(), - title: Some(title), - 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(); + controller.enqueue(StreamPart::Source(Source::Url { + id: generate_source_id(), + url: url.to_string(), + title: Some(title), + provider_metadata: None, + })); } } - continue; - } - - // ── 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())); - yield Ok(StreamPart::ToolInputStart { - id: call_id.to_string(), - tool_name: name.to_string(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } 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(""); - yield Ok(StreamPart::ToolInputEnd { - id: call_id.to_string(), - provider_metadata: None, - }); - let input = arguments.to_string(); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: call_id.to_string(), - tool_name: name.to_string(), - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - } - continue; } } - - // All other event types (web_search_call.in_progress, etc.) are ignored. + return; } - Err(error) => { - if error.is_recoverable_stream_error() { - if include_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: Value::Null }); - } - yield Ok(StreamPart::Error { error }); + + // ── 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.enqueue(StreamPart::ToolInputStart { + id: call_id.to_string(), + tool_name: name.to_string(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + }); } else { - yield Err(error); - return; + // 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.enqueue(StreamPart::ToolInputEnd { + id: call_id.to_string(), + provider_metadata: None, + }); + 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, + })); } } } - } - // Close any remaining open text blocks. - for block_id in &content_blocks { - yield Ok(StreamPart::TextEnd { id: block_id.clone(), 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); + } + } + } + } - // Final part: Finish. - let provider_meta = response_provider_metadata( - cost_in_usd_ticks, - service_tier.as_deref(), - prompt_cache_key.as_deref(), - safety_identifier.as_deref(), - ); - - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason, - usage: final_usage.unwrap_or_else(zero_usage), - provider_metadata: provider_meta, + fn flush(self, controller: &mut TransformStreamController) { + // Close any remaining open text blocks. + for block_id in self.content_blocks { + controller.enqueue(StreamPart::TextEnd { + id: block_id, + provider_metadata: None, }); - }; + } - 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), - }), - }) + // 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, + usage: self.final_usage.unwrap_or_else(zero_usage), + provider_metadata: provider_meta, + }); } } From a1e1148737696edab855895987c2fc74288eac56 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 13:28:33 +0800 Subject: [PATCH 15/16] docs: record data-only provider config and stream structs Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- CHANGELOG.md | 27 ++++++++++++++++++++++++++ docs/api/rust.md | 10 ++++++---- docs/contributing/adding-a-provider.md | 5 +++-- docs/provider-config-manual.md | 19 ++---------------- 4 files changed, 38 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5bfe89771..bfaaa34cb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -165,6 +165,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `StreamingToolCallDelta::provider_metadata`, read from the delta that starts a call. `TrackerError` converts into `AiMuxError::InvalidResponseData`. +- `StreamingToolCallTracker::with_generate_id` and `SigV4Fetch::with_clock` + take plain function pointers (`IdGenerator = fn() -> String`, + `fn() -> DateTime`) instead of closures. +- Gains `Transformer`, `TransformStreamController` and `pipe_through`, the + Rust form of the WHATWG `TransformStream` the AI SDK models build `doStream` + with (`start` / `transform` / `flush`; `enqueue`, `enqueue_error`, `error`, + `terminate`). **Rust (aimux-providers)** @@ -183,6 +190,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `generate_id` instead of `tool-call` / `tool-call-N`, and `DeltaToolCall.index` is required again, so a chunk without it fails to parse. +- Provider configuration is data, not closures. Removed: the + `transform_request_body` setting of every provider and of + `PresetSettings` (use the provider's providerOptions namespace for + vendor fields), and `OpenAICompatibleProviderSettings::metadata_extractor` + with the `MetadataExtractor` / `StreamMetadataExtractor` traits (nothing + implemented them). Changed: `OpenAICompatibleProviderSettings::supported_urls` + is a `SupportedUrls` value and `convert_usage` a `ConvertUsage` enum + (`OpenAICompatible`, `Alibaba`, `MoonshotAI`; registry rows select it with + `profile.convert_usage`); the `generate_id` settings of Anthropic, Google, + Cohere and Mistral take an `IdGenerator` (`fn() -> String`) instead of + `Arc String>`. User-supplied credentials (`Resolvable`) and the + `fetch` transport are unchanged. Requests are unchanged for callers that + did not set a removed field. +- Every language-model `doStream` (OpenAI chat and Responses, OpenAI-compatible + chat and completion, Groq, DeepSeek, Mistral, Anthropic, Bedrock, Cohere, + Google, Vertex, xAI, Hugging Face, Open Responses) is a named stream struct + driven by `pipe_through`; emitted parts are unchanged. After an Anthropic + `message_stop` or a Google/Vertex/xAI/Hugging Face error event, the rest of + the response body is read and ignored before `Finish`, as the AI SDK's + `flush` does, instead of dropping the connection early. **Rust (provider factory, RFC-0036)** diff --git a/docs/api/rust.md b/docs/api/rust.md index 6bacbd6fc..97e227991 100644 --- a/docs/api/rust.md +++ b/docs/api/rust.md @@ -55,10 +55,12 @@ let model = anthropic.messages("claude-sonnet-4-5"); Settings are package-specific. OpenAI, Anthropic and Google take `api_key: Option` and `headers: Option`; these are fixed values, not callbacks. `fetch: Option` supplies a custom -transport. These packages have no `transform_request_body` setting; -OpenAI-compatible settings expose that hook, plus `supported_urls` and -`convert_usage` callbacks for chat models. Cohere and Mistral expose -`generate_id` for generated ids. Settings only expose `name` where the upstream +transport. Settings hold data, not callbacks: there is no +`transform_request_body`; OpenAI-compatible settings take `supported_urls` +as a `SupportedUrls` value and `convert_usage` as a `ConvertUsage` choice +(`OpenAICompatible`, `Alibaba`, `MoonshotAI`) for chat models. Anthropic, +Google, Cohere and Mistral take `generate_id` as an `IdGenerator` +(`fn() -> String`, a plain function, no captured state). Settings only expose `name` where the upstream package does. Every request includes the package user-agent suffix `ai-sdk-/`. OpenAI defaults to Responses through `call` and `language_model`; `chat` diff --git a/docs/contributing/adding-a-provider.md b/docs/contributing/adding-a-provider.md index fdb657d7a..054b72a8f 100644 --- a/docs/contributing/adding-a-provider.md +++ b/docs/contributing/adding-a-provider.md @@ -39,8 +39,9 @@ its own, ...), it is case 2: add its own package next to `groq/` and `default`; a derived parameter takes `derive: {from, map, otherwise}`); `profile.max_tokens_key` (`"max_tokens"` or `"max_completion_tokens"`: the only max-token key the vendor accepts). Per-call body fields can use the - provider's options namespace. The OpenAI-compatible factory also exposes - its upstream `transform_request_body` hook; if the quirk changes auth or + provider's options namespace. A usage shape that differs from the generic + OpenAI one is `profile.convert_usage` (`"alibaba"`, `"moonshotai"`; add a + `ConvertUsage` variant for a new shape). If the quirk changes auth or error shapes, it is case 2. 2. Regenerate what is generated from the registry and commit its output: diff --git a/docs/provider-config-manual.md b/docs/provider-config-manual.md index 69256d5bd..5517c4141 100644 --- a/docs/provider-config-manual.md +++ b/docs/provider-config-manual.md @@ -4,7 +4,7 @@ > - `reasoning`（由 provider 包映射为其 wire 参数；不支持的设置产生 warning） > - `provider_options.`（per-call；OpenAI-compatible 家族里，厂商自己命名空间下 schema 不认识的字段**原样写入请求体**，用户定义一切厂商差异） > -> 声明式 `bodyOverrides`（provider 级 + per-call 的 JSON deep-merge）已删除，`ProviderConfig` / `config_json` 里传 `body_overrides` 会报 `InvalidArgument`。OpenAI-compatible settings 保留上游的 `transform_request_body` 闭包；OpenAI、Anthropic、Google settings 没有这个字段。 +> 声明式 `bodyOverrides`（provider 级 + per-call 的 JSON deep-merge）已删除，`ProviderConfig` / `config_json` 里传 `body_overrides` 会报 `InvalidArgument`。请求体改写闭包 `transform_request_body` 已从所有 provider 的 settings 中删除（配置只接受数据）；厂商差异字段写在 `provider_options.` 里透传。 > > 本手册是"知识"的归宿：各厂商的 wire 参数、配置示例、核实日期。知识会过期——以厂商官方文档为准，本手册仅作参考。 > 数据来源：[model-config-research/](internal/model-config-research/)（2026-08-01 全网调研，250 家）。 @@ -108,22 +108,7 @@ await generateText(model, prompt, { // 通用命名空间 openaiCompatible 只认 schema 字段，未知字段会被丢弃——厂商专属字段请写在厂商自己的命名空间下。 ``` -Rust 的 OpenAI-compatible 包支持对**最终**请求体做整体改写（provider 级、每次请求生效、可删字段），用该包 settings 的 `transform_request_body`： - -```rust -let provider = create_openai_compatible(OpenAICompatibleProviderSettings { - name: "acme".into(), - base_url: "https://api.acme.example/v1".into(), - transform_request_body: Some(Arc::new(|mut body| { - body["enable_thinking"] = json!(false); - body.as_object_mut().unwrap().remove("reasoning_effort"); - body - })), - ..Default::default() -})?; -``` - -JS / Python / Go 等绑定没有闭包入口，provider 级的固定字段请在每次调用的 `provider_options` 里带上。 +provider 级的请求体改写（以前的 `transform_request_body` 闭包）已删除，Rust 与各绑定一致：需要固定写入的厂商字段，请在每次调用的 `provider_options.` 里带上。 ## 6. 核实日期与来源 From 269f347fc547a12a8890a5f4f353468ac9cf3be1 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Sat, 10 Oct 2026 13:50:52 +0800 Subject: [PATCH 16/16] docs(roadmap): #185 is superseded by the RFC-0036 stack Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee Co-authored-by: Amp --- ROADMAP.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ROADMAP.md b/ROADMAP.md index 4ea40c032..4e01a85d4 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -31,7 +31,7 @@ aimux 是 **provider 接入与治理的运行时**： ├─ E 轨 内部清理（#164 已合入，改为 master 上独立 PR）+ 文档三层重组 │ ├── 新增（RFC-0036）：L0 native passthrough | ops 协议 schema | stdio CLI | 性能门禁 P -├── 独立项：#185 ToolInput 类型化 + tracker 对齐上游并移入 provider-utils（#204）；调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系) +├── 独立项：#185 由 RFC-0036 栈取代：tracker 对齐上游并移入 provider-utils（原 #204，已并入 integration），Raw / Parsed 由 provider 侧 `RawToolCall` 与 core 侧 `ToolCall` 两套类型区分；调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系) ├── 独立项：#170 registry 维护自动化 → #171 triage（前半：27 个 base_url 分歧） └── 小项（随时插空）：#181 session 后续 | #180 cache probe §10 @@ -50,7 +50,7 @@ aimux 是 **provider 接入与治理的运行时**： | **P** | 性能回归门禁：单请求开销、流式吞吐、RSS 增长 | | A2–A5 | 删 `ProviderName` 枚举及各语言副本；归档 `provider-inventory/`；aimux-web 类型副本；Flutter example 桌面脚手架 | | E1 | 内部清理 + 文档三层重组 | -| #185 | `ToolInput{Raw,Parsed}`；tracker 不删，改为对齐 `@ai-sdk/provider-utils` 并移入 aimux-provider-utils（#204）；调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系) | +| #185 | 由 RFC-0036 栈取代：tracker 不删，对齐 `@ai-sdk/provider-utils` 并移入 aimux-provider-utils（原 #204，已并入 integration）；`ToolInput{Raw,Parsed}` 不再引入，由 `RawToolCall` / `ToolCall` 两套类型区分；调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系) | | #170 / #171 前半 | sync/probe 脚本 + scheduled Action；27 个 base_url 分歧 triage | | **B1 + C2** | 解锁项：protocol 列 + `from_resolved`；40 个 FFI 构造器转发 shim（B1 形式由 descriptor 数据源 + `create_xxx` 工厂替代，C2 shim 已取消；调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系)） | | **ops 协议 schema** | op 表、消息与错误信封、二进制帧、版本协商——先以文档 + 测试落地（RFC-0036 §4） | @@ -116,7 +116,7 @@ aimux 是 **provider 接入与治理的运行时**： 1. **RFC-0032 / 0033 入库**：B1 与 #174 的前提。 2. **S1 + S2**：一个小 PR——改 profile 名、staticlib job 换 profile、全产物体积阈值。 3. **A2–A5**：四个独立 PR；#166 的 A1 checkbox 补勾（#169 已合入）。 -4. **#185**：先删 tracker，再换 `ToolInput`。 +4. **#185**：由 RFC-0036 栈取代（tracker 保留并对齐上游；Raw / Parsed 由 `RawToolCall` / `ToolCall` 区分），待维护者确认后关闭。 5. **B1 + C2**：解锁项。 ## 5. 版本节奏预估