diff --git a/CHANGELOG.md b/CHANGELOG.md index c7067860a..f99ff1ba5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -710,6 +710,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 labels goes to the only matching call whose structured arguments are not complete. Parallel same-name calls whose ids are missing or reused are no longer merged into the first call. +- `SigV4Fetch::with_clock` takes a plain function (`fn() -> DateTime`) + instead of a closure. **Rust (aimux-providers)** @@ -727,6 +729,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `type` is validated when present, ids missing from the wire come from the provider's `generate_id`, and `DeltaToolCall.index` is required again, so a chunk without it fails to parse. +- The private per-model configs hold data where they held closures: base + URLs (`OpenAIUrl`, `MessagesUrl`, `BaseUrl`), provider headers + (`ProviderHeaders`, `EndpointHeaders`), endpoint sources (Vertex, Bedrock, + Polly) and supported-URL rules. A registry row selects its chat usage + accounting with `profile.convert_usage` (`"alibaba"`, `"moonshotai"`) + instead of a name match. Public settings, requests and errors are + unchanged; the closures callers pass in settings (`transform_request_body`, + `generate_id`, the compatible `supported_urls`, `convert_usage` and + `metadata_extractor`) are kept as given. +- `PresetDescriptor` gains a crate-private field (the row's usage converter), + so code outside the crate can no longer build one with a struct literal; + descriptors come from the registry JSON. **Rust (provider factory, RFC-0036)** diff --git a/aimux-provider-utils/src/sigv4_fetch.rs b/aimux-provider-utils/src/sigv4_fetch.rs index 890234ef3..35fcdea0d 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}; @@ -344,7 +343,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. A request marked [`RequestCredentials::Omit`] is @@ -396,15 +397,15 @@ impl SigV4Fetch { inner, credentials, service: service.into(), - clock: Arc::new(Utc::now), + clock: Utc::now, signing, } } /// 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 } } @@ -467,7 +468,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 2e72a921b..4688fb505 100644 --- a/aimux-providers/src/anthropic/config.rs +++ b/aimux-providers/src/anthropic/config.rs @@ -12,34 +12,76 @@ 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, ResponseHandler, - combine_headers, normalize_headers, + ExchangeContext, FetchFunction, HeaderMapOpt, HttpRequest, ResponseHandler, combine_headers, + normalize_headers, }; use super::convert::RequestProfile; +use crate::shared::Endpoint; pub use crate::shared::TransformRequestBody; -/// Maps an endpoint path (`"/messages"`) to the full request URL. -pub(crate) type UrlFn = Arc String + Send + Sync>; +/// 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, +} -/// `(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>; +/// 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, + /// 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. + Vertex(Arc), + /// Claude Platform on AWS with no explicit base URL: every request goes + /// to `https://aws-external-anthropic.{region}.api.aws/v1` for the region + /// resolved then (`getBaseURL`, anthropic-aws-provider.ts:229-231). + AwsRegional(Arc), + /// An invalid base URL setting (`ANTHROPIC_BASE_URL`): every request + /// fails with this error before it is sent. + Invalid(AiMuxError), +} -/// Rewrites the request body a host-specific way. -pub(crate) type BodyFn = Arc Value + Send + Sync>; +/// Where the provider headers of a request come from. +#[derive(Clone)] +pub(crate) enum AnthropicHeaders { + /// Credential, fixed and user headers (the common case). + Provider(ProviderHeaders), + /// Claude Platform on AWS: `anthropic-version`, the workspace id, the + /// caller's headers and, on the API-key path, `x-api-key` last, all + /// resolved on every request. + Aws(crate::anthropic_aws::AwsHeaders), +} + +impl From for AnthropicHeaders { + fn from(headers: ProviderHeaders) -> Self { + Self::Provider(headers) + } +} -/// 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>; +impl AnthropicHeaders { + async fn resolve(&self) -> Result { + match self { + Self::Provider(headers) => headers.resolve().await, + Self::Aws(headers) => headers.resolve().await, + } + } +} /// Builds the host's error handler for non-2xx responses. pub(crate) type ErrorHandlerFn = fn() -> ResponseHandler; @@ -49,12 +91,11 @@ 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 the provider's own /// `transform_request_body`. - 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. @@ -67,7 +108,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, @@ -83,55 +124,79 @@ 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, + pub(crate) headers: AnthropicHeaders, /// Transport; `None` uses the process default, resolved per request. 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. + /// 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::Invalid(error) => Err(error.clone()), + AnthropicEndpoint::Vertex(resolver) => { + let Endpoint { base_url, headers } = resolver + .endpoint(crate::vertex::Publisher::Anthropic) + .await?; + Ok(Self { + base_url, + headers: ProviderHeaders::bearer(Credential::None, Vec::new(), Some(headers)) + .into(), + endpoint: AnthropicEndpoint::Fixed, + ..self.clone() + }) + } + AnthropicEndpoint::AwsRegional(auth) => { + let region = auth.host_region()?; + Ok(Self { + base_url: format!("https://aws-external-anthropic.{region}.api.aws/v1"), + 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) + } } } @@ -218,7 +283,7 @@ impl AnthropicModelConfig { /// 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 { + if let Some(prepare) = self.hooks.prepare_body { body = prepare(body); } match &self.transform_request_body { diff --git a/aimux-providers/src/anthropic/mod.rs b/aimux-providers/src/anthropic/mod.rs index b2885cf99..f38986348 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::decision_model::DecisionModel; use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; @@ -44,12 +45,12 @@ use aimux_core::language_model::{LanguageModel, SupportedUrls}; use aimux_core::model_catalogue::RuntimeModel; use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_provider_utils::{ - DecisionLanguageModel, FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url, + DecisionLanguageModel, FetchFunction, HeaderMapOpt, Resolvable, validate_base_url, }; -use crate::shared::{AuthScheme, Credential, credential_headers}; +use crate::shared::{AuthScheme, Credential}; -use config::{AnthropicModelConfig, AnthropicModelHooks}; +use config::{AnthropicEndpoint, AnthropicModelConfig, AnthropicModelHooks}; pub(crate) fn anthropic_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -201,13 +202,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, - ), + headers: ProviderHeaders::new( + credential, + scheme, + vec![("anthropic-version".to_string(), API_VERSION.to_string())], + settings.headers, + ) + .with_user_agent( crate::shared::aisdk_pins::ANTHROPIC.name, crate::shared::aisdk_pins::ANTHROPIC.version, ), @@ -238,7 +239,7 @@ pub struct AnthropicProvider { /// The normalized base URL, or the error of an invalid /// `ANTHROPIC_BASE_URL`, which every request reports. base_url: Result, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, supported_urls: SupportedUrls, generate_id: Option String + Send + Sync>>, @@ -248,11 +249,9 @@ impl AnthropicProvider { /// The model configuration reporting `provider` as its identity. fn model_config(&self, provider: String) -> AnthropicModelConfig { let base_url = self.base_url.clone().unwrap_or_default(); - let base = base_url.clone(); AnthropicModelConfig { provider, - url: Arc::new(move |path| format!("{base}{path}")), - headers: self.headers.clone(), + headers: self.headers.clone().into(), fetch: self.fetch.clone(), supported_urls: self.supported_urls.clone(), transform_request_body: None, @@ -261,12 +260,10 @@ impl AnthropicProvider { hooks: AnthropicModelHooks::default(), // Every request resolves its configuration first; an invalid // `ANTHROPIC_BASE_URL` fails it there. - resolve: self.base_url.as_ref().err().map(|error| { - let error = error.clone(); - let resolve: config::ResolveFn = - Arc::new(move || Box::pin(std::future::ready(Err(error.clone())))); - resolve - }), + endpoint: match &self.base_url { + Ok(_) => AnthropicEndpoint::Fixed, + Err(error) => AnthropicEndpoint::Invalid(error.clone()), + }, } } diff --git a/aimux-providers/src/anthropic_aws/mod.rs b/aimux-providers/src/anthropic_aws/mod.rs index 03d467a18..97c4b01ae 100644 --- a/aimux-providers/src/anthropic_aws/mod.rs +++ b/aimux-providers/src/anthropic_aws/mod.rs @@ -19,6 +19,8 @@ mod model; +pub(crate) use model::{Auth, AwsHeaders}; + pub use model::AnthropicAwsModel; use std::sync::{Arc, OnceLock}; @@ -32,12 +34,13 @@ 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, default_fetch, - validate_base_url, + AwsCredentials, FetchFunction, HeaderMapOpt, Resolvable, default_fetch, validate_base_url, }; use crate::anthropic::AnthropicMessagesModel; -use crate::anthropic::config::{AnthropicModelConfig, AnthropicModelHooks}; +use crate::anthropic::config::{ + AnthropicEndpoint, AnthropicHeaders, AnthropicModelConfig, AnthropicModelHooks, +}; use crate::anthropic::options::{CANONICAL, options_name_of}; pub use crate::shared::TransformRequestBody; @@ -193,7 +196,7 @@ pub struct AnthropicAwsProvider { auth: Arc, /// The explicit base URL; `None` builds the regional one per request. base_url: Option, - headers: HeadersFn, + headers: AwsHeaders, fetch: FetchFunction, transform_request_body: Option, } @@ -201,18 +204,16 @@ pub struct AnthropicAwsProvider { impl AnthropicAwsProvider { /// The configuration of requests against `base_url`. fn config_for(&self, base_url: String) -> AnthropicModelConfig { - let base = base_url.clone(); AnthropicModelConfig { provider: self.name.clone(), - url: Arc::new(move |path| format!("{base}{path}")), - headers: self.headers.clone(), + headers: AnthropicHeaders::Aws(self.headers.clone()), fetch: Some(self.fetch.clone()), supported_urls: SupportedUrls::default(), transform_request_body: self.transform_request_body.clone(), base_url, provider_options_name: self.provider_options_name.clone(), hooks: AnthropicModelHooks::default(), - resolve: None, + endpoint: AnthropicEndpoint::Fixed, } } @@ -223,25 +224,10 @@ impl AnthropicAwsProvider { if let Some(base_url) = &self.base_url { return self.config_for(base_url.clone()); } - let mut config = self.config_for(String::new()); - let template = self.config_for(String::new()); - let auth = self.auth.clone(); - config.resolve = Some(Arc::new(move || { - let template = template.clone(); - let base_url = auth - .host_region() - .map(|region| format!("https://aws-external-anthropic.{region}.api.aws/v1")); - Box::pin(async move { - let base_url = base_url?; - let base = base_url.clone(); - Ok(AnthropicModelConfig { - url: Arc::new(move |path| format!("{base}{path}")), - base_url, - ..template - }) - }) - })); - config + AnthropicModelConfig { + endpoint: AnthropicEndpoint::AwsRegional(self.auth.clone()), + ..self.config_for(String::new()) + } } /// A Messages model; `provider()` is the provider name diff --git a/aimux-providers/src/anthropic_aws/model.rs b/aimux-providers/src/anthropic_aws/model.rs index f80dde10d..5f1ff055c 100644 --- a/aimux-providers/src/anthropic_aws/model.rs +++ b/aimux-providers/src/anthropic_aws/model.rs @@ -19,8 +19,7 @@ use async_trait::async_trait; use aimux_core::error::AiMuxError; use aimux_provider_utils::{ AwsCredentials, Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, - HeadersFn, Resolvable, SigV4Fetch, SigV4Signing, combine_headers, load_optional_setting, - load_setting, + Resolvable, SigV4Fetch, SigV4Signing, combine_headers, load_optional_setting, load_setting, }; use crate::anthropic::AnthropicMessagesModel; @@ -45,7 +44,7 @@ const SESSION_TOKEN_ENV_VAR: &str = "AWS_SESSION_TOKEN"; /// Where the region, the workspace id and the credentials of a request come /// from. -pub(super) struct Auth { +pub(crate) struct Auth { pub(super) region: Option, pub(super) workspace_id: Option, pub(super) auth: Option, @@ -70,7 +69,7 @@ impl Auth { /// The region of the default base URL: a single DNS label /// (anthropic-aws-provider.ts:212-227). - pub(super) fn host_region(&self) -> Result { + pub(crate) fn host_region(&self) -> Result { let region = self.region()?; if !is_valid_hostname_part(®ion) { return Err(AiMuxError::InvalidArgument( @@ -156,25 +155,11 @@ pub(super) fn authenticated( auth: std::sync::Arc, user_headers: Option, inner: FetchFunction, -) -> (HeadersFn, FetchFunction) { - let header_auth = auth.clone(); - let headers: HeadersFn = Resolvable::from_async_fn(move || { - let auth = header_auth.clone(); - let user = user_headers.clone().unwrap_or_default(); - async move { - let mut fixed = HeaderMapOpt::new(); - fixed.insert("anthropic-version".into(), Some(API_VERSION.into())); - fixed.insert("anthropic-workspace-id".into(), Some(auth.workspace_id()?)); - let mut key = HeaderMapOpt::new(); - match auth.api_key().await? { - Some(api_key) => { - key.insert("x-api-key".into(), Some(api_key)); - } - None => auth.preflight()?, - } - Ok(combine_headers(&[&fixed, &user, &key])) - } - }); +) -> (AwsHeaders, FetchFunction) { + let headers = AwsHeaders { + auth: auth.clone(), + user: user_headers, + }; let credentials_auth = auth; let fetch: FetchFunction = std::sync::Arc::new(AnthropicAwsFetch { // aws4fetch, like anthropic-aws-fetch.ts:73-83. @@ -192,6 +177,34 @@ pub(super) fn authenticated( (headers, fetch) } +/// The provider headers of Claude Platform on AWS, resolved on every +/// request: `anthropic-version`, the workspace id, the caller's headers and, +/// on the API-key path, `x-api-key` last, as anthropic-aws-fetch.ts:132-134 +/// combines it. +#[derive(Clone)] +pub(crate) struct AwsHeaders { + auth: std::sync::Arc, + user: Option, +} + +impl AwsHeaders { + pub(crate) async fn resolve(&self) -> Result { + let auth = &self.auth; + let user = self.user.clone().unwrap_or_default(); + let mut fixed = HeaderMapOpt::new(); + fixed.insert("anthropic-version".into(), Some(API_VERSION.into())); + fixed.insert("anthropic-workspace-id".into(), Some(auth.workspace_id()?)); + let mut key = HeaderMapOpt::new(); + match auth.api_key().await? { + Some(api_key) => { + key.insert("x-api-key".into(), Some(api_key)); + } + None => auth.preflight()?, + } + Ok(combine_headers(&[&fixed, &user, &key])) + } +} + /// Adds the transport's user-agent part, then signs the request unless it /// carries an `x-api-key`. struct AnthropicAwsFetch { diff --git a/aimux-providers/src/assemblyai.rs b/aimux-providers/src/assemblyai.rs index c5c80110e..d29cc8867 100644 --- a/aimux-providers/src/assemblyai.rs +++ b/aimux-providers/src/assemblyai.rs @@ -28,10 +28,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, }; @@ -141,7 +141,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(), @@ -167,7 +167,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 b715579c0..a5ba67bc0 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, SigV4Signing, + AwsCredentials, FetchFunction, HeaderMapOpt, Resolvable, SigV4Fetch, SigV4Signing, default_fetch, load_optional_setting, load_setting, validate_base_url, }; -use crate::shared::{Endpoint, EndpointConfig}; +use crate::shared::{Endpoint, EndpointConfig, ResolveEndpoint}; /// AWS service name used for SigV4 signing. const SERVICE_NAME: &str = "polly"; @@ -235,17 +235,30 @@ 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)] +struct PollyEndpoint { + auth: Auth, + base_url: Option, + user: Option, +} + +#[async_trait::async_trait] +impl ResolveEndpoint for PollyEndpoint { + /// The base URL and provider headers of one request. + 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 +284,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 || { @@ -285,10 +297,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, }) } @@ -307,32 +321,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?, - }) - }) - }), + self.endpoint.clone(), Some(self.fetch.clone()), ) } diff --git a/aimux-providers/src/azure/mod.rs b/aimux-providers/src/azure/mod.rs index cc85d609d..a3b7c6159 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,18 +40,18 @@ 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, RequestCredentials, Resolvable, default_fetch, load_setting, validate_base_url, + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, HttpRequest, + RequestCredentials, Resolvable, default_fetch, load_setting, validate_base_url, 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, 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 +211,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 +220,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,8 +232,7 @@ 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, + speech_headers: speech_headers.with_user_agent( crate::shared::aisdk_pins::AZURE.name, crate::shared::aisdk_pins::AZURE.version, ), @@ -243,8 +243,7 @@ pub fn create_azure( fetch, crate::shared::aisdk_pins::AZURE, )), - headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( - headers, + headers: headers.with_user_agent( crate::shared::aisdk_pins::AZURE.name, crate::shared::aisdk_pins::AZURE.version, ), @@ -308,19 +307,19 @@ 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, /// The provider tools (`azure.tools` in the AI SDK). pub tools: crate::openai::tools::AzureOpenAITools, } -/// 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, @@ -352,7 +351,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 { @@ -373,6 +372,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. @@ -417,11 +435,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(), @@ -483,24 +502,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/bedrock/mod.rs b/aimux-providers/src/bedrock/mod.rs index b846b5873..a697258ad 100644 --- a/aimux-providers/src/bedrock/mod.rs +++ b/aimux-providers/src/bedrock/mod.rs @@ -46,26 +46,25 @@ 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, SigV4Signing, combine_headers, default_fetch, - load_optional_setting, load_setting, validate_base_url, without_trailing_slash, + Resolvable, SigV4Fetch, SigV4Signing, combine_headers, default_fetch, load_optional_setting, + load_setting, validate_base_url, without_trailing_slash, }; -use crate::shared::{Endpoint, EndpointConfig, TransformRequestBody, is_valid_hostname_part}; +use crate::shared::{ + Endpoint, EndpointConfig, ResolveEndpoint, TransformRequestBody, 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"; @@ -230,7 +229,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, @@ -253,7 +252,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 @@ -336,7 +335,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()); @@ -346,7 +345,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"), @@ -414,36 +413,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()]); - // amazon-bedrock-provider.ts:292-295 (getHeaders) appends the pinned - // package part. The helper then adds `ai-sdk-provider-utils/` and - // the transport (`BedrockAuthFetch`) the second occurrence. - let package = crate::shared::aisdk_pins::AMAZON_BEDROCK; - let suffix = format!("ai-sdk-{}/{}", package.name, package.version); - let agent = user - .get("user-agent") - .and_then(|value| value.as_deref()) - .map_or_else(|| suffix.clone(), |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 { @@ -464,7 +433,7 @@ pub fn create_amazon_bedrock( Ok(AmazonBedrockProvider { auth, base_url: base_url.map(|url| without_trailing_slash(&url)), - headers, + user_headers: settings.headers, fetch, transform_request_body: settings.transform_request_body, }) @@ -487,11 +456,70 @@ pub fn amazon_bedrock() -> &'static AmazonBedrockProvider { pub struct AmazonBedrockProvider { auth: Arc, base_url: Option, - headers: HeadersFn, + user_headers: Option, fetch: FetchFunction, transform_request_body: Option, } +/// 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 package user-agent part, 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)] +struct BedrockEndpoint { + auth: Arc, + base_url: Option, + /// `bedrock-runtime` or `bedrock-agent-runtime`. + service: &'static str, + endpoint_env_var: &'static str, + user: Option, +} + +#[async_trait::async_trait] +impl ResolveEndpoint for BedrockEndpoint { + /// The base URL and provider headers of one request. + 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?, + }) + } +} + +impl BedrockEndpoint { + async fn headers(&self) -> Result { + let mut user = combine_headers(&[&self.user.clone().unwrap_or_default()]); + // amazon-bedrock-provider.ts:292-295 (getHeaders) appends the pinned + // package part. The helper then adds `ai-sdk-provider-utils/` and + // the transport (`BedrockAuthFetch`) the second occurrence. + let package = crate::shared::aisdk_pins::AMAZON_BEDROCK; + let suffix = format!("ai-sdk-{}/{}", package.name, package.version); + let agent = user + .get("user-agent") + .and_then(|value| value.as_deref()) + .map_or_else(|| suffix.clone(), |value| format!("{value} {suffix}")); + user.insert("user-agent".into(), Some(agent)); + match self.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 => { + self.auth.preflight()?; + Ok(user) + } + } + } +} + impl AmazonBedrockProvider { /// The configuration of a model that calls `service` (`bedrock-runtime` /// or `bedrock-agent-runtime`). @@ -500,33 +528,18 @@ 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?, - }) - }) - }), - fetch: Some(self.fetch.clone()), - supported_urls: Arc::new(|_| SupportedUrls::default()), - transform_request_body: self.transform_request_body.clone(), - } + EndpointConfig::new( + PROVIDER.to_string(), + BedrockEndpoint { + auth: self.auth.clone(), + base_url: self.base_url.clone(), + service, + endpoint_env_var, + user: self.user_headers.clone(), + }, + Some(self.fetch.clone()), + ) + .with_transform_request_body(self.transform_request_body.clone()) } fn runtime_config(&self) -> EndpointConfig { @@ -591,11 +604,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 6bc9d2b95..f7ac6fa06 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 68d6dcefa..c70af073d 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 7e202fa45..3d47c0dd2 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/codex.rs b/aimux-providers/src/codex.rs index 7702aceba..cb1be07e4 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; @@ -60,7 +61,7 @@ use crate::openai::responses::{ OpenAIResponsesModel, ResponsesNamespace, build_responses_request_body, }; pub use crate::shared::TransformRequestBody; -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"; @@ -209,7 +210,7 @@ pub fn create_codex(settings: CodexProviderSettings) -> Result, transform_request_body: Option, } diff --git a/aimux-providers/src/cohere/mod.rs b/aimux-providers/src/cohere/mod.rs index 08e77bda1..b9429e8aa 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,8 +118,7 @@ pub fn create_cohere(settings: CohereProviderSettings) -> Result &'static CohereProvider { pub struct CohereProvider { generate_id: Option String + Send + Sync>>, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, } diff --git a/aimux-providers/src/dataforseo.rs b/aimux-providers/src/dataforseo.rs index 5cb670193..7b3cf15e1 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/decision_support.rs b/aimux-providers/src/decision_support.rs index daa181443..8a37aa446 100644 --- a/aimux-providers/src/decision_support.rs +++ b/aimux-providers/src/decision_support.rs @@ -280,9 +280,8 @@ pub(crate) fn options( #[cfg(test)] mod tests { - use std::sync::Arc; - - use aimux_provider_utils::Resolvable; + use crate::openai::config::OpenAIUrl; + use crate::shared::{Credential, ProviderHeaders}; use super::*; @@ -291,17 +290,14 @@ mod tests { let mut config = OpenAIModelConfig::fixed( "openai".into(), "https://example.com/v1".into(), - Resolvable::Value(HashMap::new()), + ProviderHeaders::bearer(Credential::None, Vec::new(), None), None, None, ); - config.url = Arc::new(|path| { - assert_eq!(path, "/decisions"); - Ok("https://example.com/decisions?version=1".into()) - }); + config.url = OpenAIUrl::Base("https://example.com/decisions-host".into()); assert_eq!( DecisionConfig::OpenAI(config).url("/decisions").unwrap(), - "https://example.com/decisions?version=1" + "https://example.com/decisions-host/decisions" ); } diff --git a/aimux-providers/src/deepgram.rs b/aimux-providers/src/deepgram.rs index 0180d83c8..c7a54cfcc 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 148b28836..b07e1a1af 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; pub(crate) const DEFAULT_BASE_URL: &str = "https://api.deepseek.com"; @@ -91,16 +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, } @@ -149,7 +149,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 a1e48beac..69c91455d 100644 --- a/aimux-providers/src/deepseek/model.rs +++ b/aimux-providers/src/deepseek/model.rs @@ -358,7 +358,7 @@ impl LanguageModel for DeepSeekChatLanguageModel { } async 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 524c46107..34a5f4367 100644 --- a/aimux-providers/src/elevenlabs.rs +++ b/aimux-providers/src/elevenlabs.rs @@ -28,12 +28,10 @@ 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, TransformRequestBody, credential_headers, + AuthScheme, Credential, EndpointConfig, ProviderHeaders, TransformRequestBody, }; pub(crate) mod options; @@ -159,7 +157,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(), @@ -191,7 +189,7 @@ pub fn elevenlabs() -> &'static ElevenLabsProvider { pub struct ElevenLabsProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, transform_request_body: Option, #[cfg(feature = "realtime")] diff --git a/aimux-providers/src/exa_ai.rs b/aimux-providers/src/exa_ai.rs index 8b59450a8..a5751af43 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 1b9f4da03..5372a9505 100644 --- a/aimux-providers/src/fal.rs +++ b/aimux-providers/src/fal.rs @@ -25,10 +25,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, }; @@ -145,7 +145,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 3380781eb..a34d3e5c7 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 23d3aea7c..9c90a955f 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 e4571032f..e94835d86 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::{ - DecisionLanguageModel, FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url, + DecisionLanguageModel, 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,17 +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 String + Send + Sync>>, #[cfg(feature = "realtime")] @@ -293,25 +293,17 @@ 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))), - transform_request_body: None, - } + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) + .with_supported_urls(SupportedUrlsSource::ByModel { + base_url: self.base_url.clone(), + resolve: supported_urls, + }) } /// 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 6068468c9..1d7ac9306 100644 --- a/aimux-providers/src/google/model.rs +++ b/aimux-providers/src/google/model.rs @@ -79,7 +79,7 @@ impl LanguageModel for GoogleModel { } async 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 d6ed89bfd..15733284e 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 807379a71..6c3c7af0d 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 488db1dda..82658b134 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}; pub(crate) const DEFAULT_BASE_URL: &str = "https://api.groq.com/openai/v1"; pub(crate) const API_KEY_ENV_VAR: &str = "GROQ_API_KEY"; @@ -95,16 +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 d938001dc..5e1d82e74 100644 --- a/aimux-providers/src/huggingface.rs +++ b/aimux-providers/src/huggingface.rs @@ -35,10 +35,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}; pub(crate) const DEFAULT_BASE_URL: &str = "https://router.huggingface.co/v1"; pub(crate) const API_KEY_ENV_VAR: &str = "HUGGINGFACE_API_KEY"; @@ -113,7 +113,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, @@ -142,7 +142,7 @@ pub fn huggingface() -> &'static HuggingFaceProvider { pub struct HuggingFaceProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, transform_request_body: Option, } diff --git a/aimux-providers/src/hume.rs b/aimux-providers/src/hume.rs index 8ea3e7c95..05ef9c39c 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 e39bd999e..b7ee8b524 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 479834456..5123bdd98 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 437a382e2..ccf42a438 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. /// Failed-response handler of the Linkup endpoint: `message`, and `type` @@ -140,7 +140,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 2f56069be..22ba728f1 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 efa211136..4e33662ab 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 38f0a09b8..c8b4c7ed9 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,16 +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", ), + Vec::new(), + settings.headers, + ) + .with_user_agent( crate::shared::aisdk_pins::MISTRAL.name, crate::shared::aisdk_pins::MISTRAL.version, ); @@ -200,7 +200,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 String + Send + Sync>>, } @@ -222,7 +222,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.clone(), ) } diff --git a/aimux-providers/src/mistral/model.rs b/aimux-providers/src/mistral/model.rs index 7c9a726b4..b33a61dc0 100644 --- a/aimux-providers/src/mistral/model.rs +++ b/aimux-providers/src/mistral/model.rs @@ -165,7 +165,7 @@ impl LanguageModel for MistralModel { } async 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 dd7cb159f..236ec1f91 100644 --- a/aimux-providers/src/open_responses.rs +++ b/aimux-providers/src/open_responses.rs @@ -33,12 +33,11 @@ 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, normalize_headers, - validate_base_url, with_user_agent_suffix, -}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::shared::{Credential, EndpointConfig, TransformRequestBody, provider_headers}; +use crate::shared::{ + Credential, EndpointConfig, EndpointHeaders, ProviderHeaders, TransformRequestBody, +}; fn open_responses_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -188,8 +187,7 @@ 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 provider = ProviderHeaders::bearer(credential, Vec::new(), None); // open-responses-provider.ts:114-127: the key and the user headers, then // `withUserAgentSuffix(…, userAgentSuffix ?? ai-sdk-open-responses/VERSION)`. let package = crate::shared::aisdk_pins::OPEN_RESPONSES; @@ -197,24 +195,11 @@ pub fn create_open_responses( .user_agent_suffix .unwrap_or_else(|| format!("ai-sdk-{}/{}", package.name, package.version)) .into(); - let headers: HeadersFn = Resolvable::from_async_fn(move || { - let provider_layer = provider_layer.clone(); - let user = user.clone(); - let suffix = suffix.clone(); - async move { - let layer = provider_layer.resolve().await?; - let combined = match &user { - Some(user) => combine_headers(&[&layer, &user.resolve().await?]), - None => layer, - }; - let mut headers = normalize_headers(combined).into_iter().collect(); - with_user_agent_suffix(&mut headers, &suffix); - Ok(headers - .into_iter() - .map(|(name, value)| (name, Some(value))) - .collect()) - } - }); + let headers = EndpointHeaders::WithUserHeaders { + provider, + user: settings.headers, + user_agent_suffix: Some(suffix), + }; Ok(OpenResponsesProvider { name: settings.name, base_url, @@ -230,7 +215,7 @@ pub fn create_open_responses( pub struct OpenResponsesProvider { name: String, base_url: String, - headers: HeadersFn, + headers: EndpointHeaders, fetch: Option, transform_request_body: Option, } diff --git a/aimux-providers/src/openai/config.rs b/aimux-providers/src/openai/config.rs index f39d20f0f..80c0fa1f1 100644 --- a/aimux-providers/src/openai/config.rs +++ b/aimux-providers/src/openai/config.rs @@ -9,25 +9,61 @@ //! refuses to send credentialed headers to any other origin after a redirect. use std::collections::HashMap; -use std::sync::Arc; use serde_json::Value; +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; +use crate::azure::UrlRules; -pub use crate::shared::TransformRequestBody; +/// 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, + }, + /// An invalid base URL setting (`OPENAI_BASE_URL`): every endpoint fails + /// with this error, so every request reports it. + Invalid(AiMuxError), +} -/// 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>; +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()), + Self::Invalid(error) => Err(error.clone()), + } + } +} + +pub use crate::shared::TransformRequestBody; /// 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. @@ -36,10 +72,10 @@ 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, + pub(crate) headers: ProviderHeaders, pub(crate) token_provider: Option>, /// Transport; `None` uses the process default, resolved per request. pub(crate) fetch: Option, @@ -58,13 +94,13 @@ impl OpenAIModelConfig { pub(crate) fn fixed( provider: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, transform_request_body: Option, ) -> Self { Self { provider, - url: Arc::new(move |path| Ok(format!("{base_url}{path}"))), + url: OpenAIUrl::Base(base_url), headers, token_provider: None, fetch, @@ -82,7 +118,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/openai/mod.rs b/aimux-providers/src/openai/mod.rs index e062df5b1..99e57b31b 100644 --- a/aimux-providers/src/openai/mod.rs +++ b/aimux-providers/src/openai/mod.rs @@ -42,6 +42,7 @@ use std::sync::{Arc, OnceLock}; use futures::future::BoxFuture; use serde_json::Value; +use crate::shared::ProviderHeaders; use aimux_core::decision_model::DecisionModel; use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; @@ -53,10 +54,10 @@ use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_core::speech_model::SpeechModel; use aimux_core::transcription_model::TranscriptionModel; use aimux_provider_utils::{ - DecisionLanguageModel, FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url, + DecisionLanguageModel, FetchFunction, HeaderMapOpt, Resolvable, validate_base_url, }; -use crate::shared::{Credential, provider_headers}; +use crate::shared::Credential; use config::OpenAIModelConfig; @@ -260,22 +261,22 @@ pub fn create_openai(settings: OpenAIProviderSettings) -> Result, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, #[cfg(feature = "realtime")] web_socket: Option>, @@ -329,8 +330,7 @@ impl OpenAIProvider { config.responses.file_id_prefixes = vec!["file-"]; } if let Err(error) = &self.base_url { - let error = error.clone(); - config.url = Arc::new(move |_| Err(error.clone())); + config.url = config::OpenAIUrl::Invalid(error.clone()); } config } diff --git a/aimux-providers/src/openai_compatible/chat.rs b/aimux-providers/src/openai_compatible/chat.rs index 3bea51d80..adcd06f15 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,10 +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, - supports_multi_part_tool_content: self.config.chat.supports_multi_part_tool_content, - dialect: &self.config.chat.dialect, + chat: &self.config.chat, }, ) } @@ -189,10 +186,7 @@ impl LanguageModel for OpenAICompatibleChatModel { } async fn supported_urls(&self) -> SupportedUrls { - match &self.config.chat.dialect.supported_urls { - Some(urls) => urls().await, - None => self.config.chat.supported_urls.clone(), - } + self.config.chat.supported_urls().await } async fn do_generate(&self, options: &CallOptions) -> Result { @@ -282,15 +276,12 @@ 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 chat = &self.config.chat; + let usage = chat.convert_usage(raw.get("usage")); let mut metadata = aimux_core::shared::SharedProviderMetadata::new(); metadata.insert(metadata_key.clone(), Map::new()); - if let Some(extractor) = &dialect.metadata_extractor + if let Some(extractor) = &chat.metadata_extractor && let Some(extra) = extractor.extract_metadata(&raw) { for (key, value) in extra { @@ -347,8 +338,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 + let chat = self.config.chat.clone(); + let mut metadata_extractor = chat .metadata_extractor .as_ref() .map(|extractor| extractor.create_stream_extractor()); @@ -406,7 +397,7 @@ impl LanguageModel for OpenAICompatibleChatModel { continue; } - let chunk_usage_raw: Option = match &dialect.stream_usage_key { + let chunk_usage_raw: Option = match &chat.stream_usage_key { Some(key) => parsed.get(key).and_then(|v| v.get("usage")).cloned(), None => parsed.get("usage").cloned(), } @@ -636,10 +627,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: chat.convert_usage(final_usage_raw.as_ref()), provider_metadata: Some(metadata), }); }; @@ -772,15 +760,15 @@ mod tests { use aimux_core::options::CallOptions; use aimux_provider_utils::ProviderErrorParts; - use crate::openai_compatible::config::{BaseUrl, ChatDialect}; - use crate::openai_compatible::{Assembly, ChatProfile, OpenAICompatibleProvider}; + use crate::openai_compatible::config::{BaseUrl, ChatSettings}; + use crate::openai_compatible::{Assembly, OpenAICompatibleProvider}; use crate::shared::Credential; fn hello() -> CallOptions { CallOptions::new(vec![LanguageModelMessage::user_text("hi")]) } - fn provider(base_url: BaseUrl, dialect: ChatDialect) -> OpenAICompatibleProvider { + fn provider(base_url: BaseUrl, chat: ChatSettings) -> OpenAICompatibleProvider { OpenAICompatibleProvider::assemble(Assembly { name: "acme".to_string(), base_url, @@ -790,11 +778,7 @@ mod tests { query_params: None, fetch: None, transform_request_body: None, - profile: ChatProfile { - include_usage: false, - supports_structured_outputs: false, - dialect, - }, + chat, package: crate::shared::aisdk_pins::OPENAI_COMPATIBLE, }) .unwrap() @@ -816,12 +800,14 @@ mod tests { ResponseTemplate::new(400).set_body_json(json!({"detail": "bad prompt", "kind": 42})), ) .await; - let mut dialect = ChatDialect::baseline(); - dialect.error_structure = Arc::new(|data: &Value| ProviderErrorParts { - message: data["detail"].as_str().unwrap_or_default().to_string(), - provider_code: data["kind"].as_u64().map(|k| k.to_string()), - }); - let error = provider(BaseUrl::Fixed(server.uri()), dialect) + let chat = ChatSettings { + error_structure: Arc::new(|data: &Value| ProviderErrorParts { + message: data["detail"].as_str().unwrap_or_default().to_string(), + provider_code: data["kind"].as_u64().map(|k| k.to_string()), + }), + ..ChatSettings::default() + }; + let error = provider(BaseUrl::Fixed(server.uri()), chat) .chat("m") .do_generate(&hello()) .await @@ -835,8 +821,15 @@ mod tests { } } + /// A preset base URL is data the request resolves: creating the provider + /// and the model reads nothing, a request that cannot resolve it fails, + /// and the next request resolves it again. #[tokio::test] + #[serial_test::serial] async fn a_lazy_base_url_is_evaluated_per_request_and_its_error_fails_the_request() { + use crate::preset::{AuthMode, ParamSpec, PresetDescriptor}; + + const VAR: &str = "AIMUX_TEST_COMPAT_LAZY_BASE_URL"; let server = MockServer::start().await; serve( &server, @@ -846,29 +839,38 @@ mod tests { })), ) .await; - let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let counter = calls.clone(); - let uri = server.uri(); - let lazy = BaseUrl::Lazy(Arc::new(move || { - if counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { - Err(AiMuxError::InvalidArgument("not ready".into())) - } else { - Ok(uri.clone()) - } + // SAFETY: the test is `#[serial]` and no other test reads `VAR`. + unsafe { std::env::remove_var(VAR) }; + let descriptor: &'static PresetDescriptor = Box::leak(Box::new(PresetDescriptor { + name: "acme".to_string(), + display: "Acme".to_string(), + base_url: "http://{host}/v1".to_string(), + env_var: String::new(), + auth: AuthMode::None, + base_url_env: Some(VAR.to_string()), + max_tokens_key: None, + convert_usage: None, + params: vec![ParamSpec { + name: "host".to_string(), + env: Vec::new(), + default: None, + derive: None, + }], + user_agent: ("openai-compatible", "3.0.62"), })); - let provider = provider(lazy, ChatDialect::baseline()); + let lazy = BaseUrl::Preset { + descriptor, + params: Arc::new(std::collections::HashMap::new()), + }; + let provider = provider(lazy, ChatSettings::default()); let model = provider.chat("m"); - assert_eq!( - calls.load(std::sync::atomic::Ordering::SeqCst), - 0, - "creating a provider or a model resolves nothing" - ); let error = model.do_generate(&hello()).await.unwrap_err(); assert!(matches!(error, AiMuxError::InvalidArgument(_)), "{error:?}"); - model - .do_generate(&hello()) - .await - .expect("resolved on the next request"); - assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2); + // SAFETY: as above. + unsafe { std::env::set_var(VAR, server.uri()) }; + let result = model.do_generate(&hello()).await; + // SAFETY: as above. + unsafe { std::env::remove_var(VAR) }; + result.expect("resolved on the next request"); } } diff --git a/aimux-providers/src/openai_compatible/config.rs b/aimux-providers/src/openai_compatible/config.rs index a25fa7769..579572f90 100644 --- a/aimux-providers/src/openai_compatible/config.rs +++ b/aimux-providers/src/openai_compatible/config.rs @@ -3,8 +3,12 @@ //! 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 `url` and `headers` functions, this +//! config holds data, or a package's own data behind [`ResolveBaseUrl`] / +//! [`ResolveFetchHeaders`]; the functions a caller or a vendor package passes +//! (`supportedUrls`, `convertUsage`, `errorStructure`, `metadataExtractor`, +//! `transformRequestBody`) are kept as given. The shared model never compares +//! a provider name. use futures::future::BoxFuture; use std::collections::HashMap; @@ -17,38 +21,22 @@ use aimux_core::AiMuxError; use aimux_core::language_model::SupportedUrls; use aimux_provider_utils::{ ExchangeContext, Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, - HeadersFn, HttpRequest, ProviderErrorParts, combine_headers, normalize_headers, + HttpRequest, ProviderErrorParts, combine_headers, normalize_headers, }; +use crate::preset::PresetDescriptor; +use crate::shared::ProviderHeaders; + pub use crate::shared::TransformRequestBody; /// Converts the raw optional token usage of a chat response. pub type ConvertUsage = Arc) -> Usage + Send + Sync>; -/// 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)] -pub(crate) enum BaseUrl { - Fixed(String), - Lazy(Arc Result + Send + Sync>), -} - -impl BaseUrl { - pub(crate) fn resolve(&self) -> Result { - match self { - Self::Fixed(url) => Ok(url.clone()), - Self::Lazy(produce) => produce(), - } - } -} - -/// Maps a provider error payload to the message and code of the API error. -pub(crate) type ErrorStructure = Arc ProviderErrorParts + Send + Sync>; +/// Returns the URLs supported by chat models, read each time a chat model's +/// supported URLs are: `supportedUrls?: () => LanguageModelV4['supportedUrls']` +/// (openai-compatible@3.0.62 `openai-compatible-chat-language-model.ts:84`, +/// `:161-163`), whose value may be a `PromiseLike`. +pub type SupportedUrlsFn = Arc BoxFuture<'static, SupportedUrls> + Send + Sync>; /// Captures vendor-specific metadata from responses (the AI SDK's /// `MetadataExtractor`). The returned map is merged into `provider_metadata` @@ -70,46 +58,59 @@ pub trait StreamMetadataExtractor: Send { 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, - pub convert_usage: Option, - pub error_structure: ErrorStructure, +/// A base URL a package builds from settings read when a request is made +/// (a Vertex AI partner sub-provider's project and location). +pub(crate) trait ResolveBaseUrl: std::fmt::Debug + Send + Sync { + /// The base URL of one request. + fn resolve(&self) -> Result; } -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), +/// The headers a package adds in its own fetch wrapper +/// ([`CompatModelConfig::fetch_headers`]). +#[async_trait::async_trait] +pub(crate) trait ResolveFetchHeaders: Send + Sync { + /// The headers of one request. + async fn resolve(&self) -> Result; +} + +/// 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 +/// [`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), + Preset { + descriptor: &'static PresetDescriptor, + /// The explicit template parameters of the settings. + params: Arc>, + }, + /// A package's own base URL, resolved when a request is made. + Dynamic(Arc), +} + +impl BaseUrl { + pub(crate) fn resolve(&self) -> Result { + match self { + Self::Fixed(url) => Ok(url.clone()), + Self::Preset { descriptor, params } => { + crate::preset::resolve_base_url(descriptor, params) + } + Self::Dynamic(resolver) => resolver.resolve(), } } } +/// Maps a provider error payload to the message and code of the API error +/// (the AI SDK's `errorStructure`). +pub(crate) type ErrorStructure = Arc ProviderErrorParts + Send + Sync>; + /// `{ 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 { @@ -128,21 +129,75 @@ pub(crate) fn default_error_structure(data: &Value) -> ProviderErrorParts { } } -/// Returns the URLs supported by chat models, read each time a chat model's -/// supported URLs are: `supportedUrls?: () => LanguageModelV4['supportedUrls']` -/// (openai-compatible@3.0.62 `openai-compatible-chat-language-model.ts:84`, -/// `:161-163`), whose value may be a `PromiseLike`. -pub type SupportedUrlsFn = Arc BoxFuture<'static, SupportedUrls> + Send + Sync>; - -/// Chat-only settings of a provider. +/// The chat settings of a compatible endpoint: the AI SDK's +/// `includeUsage` / `supportsStructuredOutputs` / +/// `supportsMultiPartToolContent` / `supportedUrls` / `convertUsage` / +/// `metadataExtractor`, plus the preset capability data of RFC-0036 section +/// 5. #[derive(Clone)] pub(crate) struct ChatSettings { pub include_usage: bool, pub supports_structured_outputs: bool, - /// `supportsMultiPartToolContent` of the AI SDK compat settings. + /// `supportsMultiPartToolContent`: a `content` tool output is sent as an + /// array of content parts instead of a JSON string. pub supports_multi_part_tool_content: bool, - pub supported_urls: SupportedUrls, - pub dialect: Arc, + /// URL patterns supported by chat models, evaluated when requested. + /// `None`: none. + pub supported_urls: Option, + /// The vendor's or the caller's usage accounting. `None`: the generic + /// OpenAI-shaped usage (`convertOpenAICompatibleChatUsage`). + pub convert_usage: Option, + pub metadata_extractor: Option>, + /// 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, + pub error_structure: ErrorStructure, +} + +impl Default for ChatSettings { + /// The AI SDK baseline. + fn default() -> Self { + Self { + include_usage: false, + supports_structured_outputs: false, + supports_multi_part_tool_content: false, + supported_urls: None, + convert_usage: None, + metadata_extractor: None, + supports_top_k: false, + supports_tools: true, + supports_response_format: true, + max_tokens_key: None, + stream_usage_key: None, + error_structure: Arc::new(default_error_structure), + } + } +} + +impl ChatSettings { + /// The URL patterns chat models fetch themselves. + pub(crate) async fn supported_urls(&self) -> SupportedUrls { + match &self.supported_urls { + Some(supported_urls) => supported_urls().await, + None => SupportedUrls::default(), + } + } + + /// The core usage of a chat response's raw `usage` object. + pub(crate) fn convert_usage(&self, raw: Option<&Value>) -> Usage { + match &self.convert_usage { + Some(convert) => convert(raw), + None => super::chat::usage_from_raw(raw), + } + } } /// What a model needs to talk to the API. @@ -153,14 +208,14 @@ 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, /// Headers layered over the per-call headers, resolved on every request /// after `headers`: what an AI SDK package that builds on /// `createOpenAICompatible` adds in its own fetch wrapper (google-vertex's /// Node MaaS and xAI providers). Their `user-agent` replaces the whole /// header, the helper layer's part included ([`FetchUserAgent`]). - pub fetch_headers: Option, + pub fetch_headers: Option>, /// Transport; `None` uses the process default, resolved per request. pub fetch: Option, pub transform_request_body: Option, @@ -269,7 +324,7 @@ impl CompatModelConfig { pub(crate) fn failed_response_handler( &self, ) -> aimux_provider_utils::ResponseHandler { - let structure = self.chat.dialect.error_structure.clone(); + let structure = self.chat.error_structure.clone(); aimux_provider_utils::create_json_error_response_handler(move |data| structure(data)) } } diff --git a/aimux-providers/src/openai_compatible/convert.rs b/aimux-providers/src/openai_compatible/convert.rs index ef25c09ba..c16fc641b 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,10 +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 supports_multi_part_tool_content: bool, - pub dialect: &'a ChatDialect, + pub chat: &'a ChatSettings, } /// A request body and what building it produced. @@ -191,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 @@ -239,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), ); } @@ -254,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 { @@ -269,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)); @@ -338,12 +335,12 @@ pub(crate) fn build_request_body( &options.prompt, &MessageSpec { metadata_key: &metadata_key, - supports_multi_part_tool_content: spec.supports_multi_part_tool_content, + supports_multi_part_tool_content: chat.supports_multi_part_tool_content, }, )?; 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 { @@ -365,7 +362,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 307630e82..34f6d4e02 100644 --- a/aimux-providers/src/openai_compatible/mod.rs +++ b/aimux-providers/src/openai_compatible/mod.rs @@ -49,16 +49,14 @@ use futures::future::BoxFuture; 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_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, combine_headers, validate_base_url, -}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; use crate::shared::aisdk_pins::Package; -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`). @@ -152,7 +150,7 @@ pub fn create_openai_compatible( .map_or(Credential::None, |key| { Credential::Explicit(Resolvable::Value(key)) }); - let headers = compatible_headers(credential, settings.headers.clone()); + let headers = ProviderHeaders::bearer(credential, Vec::new(), settings.headers.clone()); build( settings, BaseUrl::Fixed(base_url), @@ -172,10 +170,11 @@ pub fn create_openai_compatible( pub(crate) fn build( settings: OpenAICompatibleProviderSettings, base_url: BaseUrl, - headers: HeadersFn, - fetch_headers: Option, + headers: ProviderHeaders, + fetch_headers: Option>, package: Package, ) -> Result { + let defaults = ChatSettings::default(); let mut provider = OpenAICompatibleProvider::assemble(Assembly { name: settings.name, base_url, @@ -186,40 +185,25 @@ pub(crate) fn build( fetch: settings.fetch, transform_request_body: settings.transform_request_body, package, - 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(), + supports_multi_part_tool_content: settings + .supports_multi_part_tool_content + .unwrap_or(false), + supported_urls: settings.supported_urls, + convert_usage: settings.convert_usage, + metadata_extractor: settings.metadata_extractor, + ..defaults }, })?; - provider.chat.supports_multi_part_tool_content = - settings.supports_multi_part_tool_content.unwrap_or(false); - 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; - provider.headers = aimux_provider_utils::headers::with_user_agent_suffix_fn( - headers, - package.name, - package.version, - ); + provider.headers = headers.with_user_agent(package.name, package.version); provider.fetch_headers = fetch_headers; 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, @@ -230,7 +214,7 @@ pub(crate) struct Assembly { pub query_params: Option>, pub fetch: Option, pub transform_request_body: Option, - pub profile: ChatProfile, + pub chat: ChatSettings, /// The package the user-agent names: its `ai-sdk-{name}/{version}` part /// in the provider headers, and its provider-utils in the transport. pub package: Package, @@ -242,64 +226,38 @@ pub struct OpenAICompatibleProvider { name: String, base_url: BaseUrl, query_params: Option>>, - headers: HeadersFn, - fetch_headers: Option, + headers: ProviderHeaders, + fetch_headers: Option>, 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, - ), - assembly.package.name, - assembly.package.version, - ), + headers: ProviderHeaders { + credential: assembly.credential, + scheme: AuthScheme::Bearer, + fixed: assembly.fixed_headers, + user: assembly.headers, + user_agent: Some((assembly.package.name, assembly.package.version)), + }, fetch_headers: None, fetch: Some(crate::shared::package_fetch( assembly.fetch, assembly.package, )), transform_request_body: assembly.transform_request_body, - chat: ChatSettings { - include_usage: assembly.profile.include_usage, - supports_structured_outputs: assembly.profile.supports_structured_outputs, - supports_multi_part_tool_content: false, - supported_urls: SupportedUrls::default(), - dialect: Arc::new(assembly.profile.dialect), - }, + chat: assembly.chat, }) } diff --git a/aimux-providers/src/parallel_ai.rs b/aimux-providers/src/parallel_ai.rs index 6b360e7d9..66f33c9d8 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. /// Failed-response handler of the Parallel AI endpoint: `message`, and `type` @@ -116,7 +116,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(), @@ -142,7 +142,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/preset.rs b/aimux-providers/src/preset.rs index 3e9ceecf0..52feb5744 100644 --- a/aimux-providers/src/preset.rs +++ b/aimux-providers/src/preset.rs @@ -30,13 +30,16 @@ use std::collections::{HashMap, HashSet}; use std::sync::{Arc, OnceLock}; use serde::Deserialize; +use serde_json::Value; use aimux_core::error::AiMuxError; +use aimux_core::types::Usage; use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; -use crate::openai_compatible::config::{BaseUrl, ChatDialect}; +use crate::openai_compatible::chat::usage_from_raw; +use crate::openai_compatible::config::{BaseUrl, ChatSettings}; use crate::openai_compatible::{ - Assembly, ChatProfile, OpenAICompatibleProvider, TransformRequestBody, + Assembly, ConvertUsage, OpenAICompatibleProvider, TransformRequestBody, }; use crate::shared::Credential; @@ -87,6 +90,9 @@ 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, + /// How chat usage is read (`profile.convert_usage`); `None`: the generic + /// OpenAI-shaped usage. + pub(crate) convert_usage: Option) -> Usage>, pub params: Vec, /// `(package, version)` of the `ai-sdk-{package}/{version}` user-agent /// (`profile.user_agent`). `openai-compatible` by default; @@ -244,9 +250,12 @@ fn parse_registry(json: &str) -> Vec { "invalid base_url_env name", ); check( - row.profile - .keys() - .all(|key| matches!(key.as_str(), "max_tokens_key" | "user_agent")), + row.profile.keys().all(|key| { + matches!( + key.as_str(), + "max_tokens_key" | "convert_usage" | "user_agent" + ) + }), "unknown profile key", ); // bedrock-mantle-provider.ts:236-240 (getHeaders) names the @@ -264,6 +273,13 @@ fn parse_registry(json: &str) -> Vec { .is_none_or(|key| matches!(key, "max_tokens" | "max_completion_tokens")), "invalid max_tokens_key", ); + let convert_usage: Option) -> Usage> = + match row.profile.get("convert_usage").map(String::as_str) { + None => None, + Some("alibaba") => Some(alibaba_usage), + Some("moonshotai") => Some(moonshotai_usage), + Some(value) => panic!("registry row '{name}': unknown convert_usage {value:?}"), + }; let mut params = Vec::with_capacity(row.params.len()); for param in row.params { check(identifier.is_match(¶m.name), "bad param name"); @@ -310,6 +326,7 @@ fn parse_registry(json: &str) -> Vec { auth, base_url_env: row.base_url_env, max_tokens_key, + convert_usage, params, user_agent, }); @@ -318,7 +335,7 @@ fn parse_registry(json: &str) -> 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. /// @@ -395,10 +412,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) { @@ -411,42 +428,15 @@ fn assemble( (AuthMode::None, None) => Credential::OptionalEnv(descriptor.env_var.clone()), }; - let mut dialect = ChatDialect::baseline(); - dialect.supports_top_k = true; - dialect.max_tokens_key = descriptor.max_tokens_key.as_deref(); - if matches!(descriptor.name.as_str(), "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.as_deref(), + convert_usage: descriptor + .convert_usage + .map(|convert| Arc::new(convert) as ConvertUsage), + ..ChatSettings::default() }; // The Bedrock Mantle transport appends the package part a second time, @@ -474,7 +464,7 @@ fn assemble( query_params: None, fetch, transform_request_body: settings.transform_request_body, - profile, + chat, package, }) } @@ -526,7 +516,7 @@ fn env_value(names: &[impl AsRef]) -> 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 { @@ -598,6 +588,44 @@ fn resolve_base_url( validate_base_url(&url) } +/// Alibaba's `convertUsage`: cache writes from +/// `prompt_tokens_details.cache_creation_input_tokens` (or +/// `cache_write_tokens`), taken out of the uncached input. +fn alibaba_usage(raw: Option<&Value>) -> Usage { + let mut usage = usage_from_raw(raw); + if let Some(raw) = raw.filter(|raw| !raw.is_null()) { + 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)); + } + usage +} + +/// Moonshot AI's `convertUsage`: cache reads from the top-level +/// `cached_tokens`. +fn moonshotai_usage(raw: Option<&Value>) -> Usage { + let mut usage = usage_from_raw(raw); + if let Some(cached) = raw.and_then(|raw| 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 +} + +/// A token count of a raw usage object. +fn tokens(value: &Value) -> Option { + value.as_u64().and_then(|value| u32::try_from(value).ok()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/aimux-providers/src/prodia.rs b/aimux-providers/src/prodia.rs index f0169fcb9..172f3523c 100644 --- a/aimux-providers/src/prodia.rs +++ b/aimux-providers/src/prodia.rs @@ -23,7 +23,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| { @@ -151,7 +151,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![( @@ -160,7 +160,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/provider.rs b/aimux-providers/src/provider.rs index dede73721..b7ffe9453 100644 --- a/aimux-providers/src/provider.rs +++ b/aimux-providers/src/provider.rs @@ -20,8 +20,8 @@ use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; use aimux_provider_utils::{Resolvable, load_api_key, validate_base_url}; -use crate::openai_compatible::config::{BaseUrl, ChatDialect}; -use crate::openai_compatible::{Assembly, ChatProfile, OpenAICompatibleProvider}; +use crate::openai_compatible::config::{BaseUrl, ChatSettings}; +use crate::openai_compatible::{Assembly, OpenAICompatibleProvider}; use crate::preset::PresetSettings; use crate::shared::Credential; @@ -64,7 +64,7 @@ fn default_true() -> bool { impl ProviderProfile { /// The chat behavior this profile describes. - fn chat_profile(&self) -> Result { + fn chat_settings(&self) -> Result { let max_tokens_key = match self.max_tokens_key.as_deref() { None => None, Some("max_tokens") => Some("max_tokens"), @@ -76,16 +76,15 @@ impl ProviderProfile { ))); } }; - let mut dialect = ChatDialect::baseline(); - dialect.supports_top_k = self.supports_top_k; - dialect.supports_tools = self.supports_tools; - dialect.supports_response_format = self.supports_response_format; - dialect.stream_usage_key = self.stream_usage_key.clone(); - dialect.max_tokens_key = max_tokens_key; - Ok(ChatProfile { + Ok(ChatSettings { include_usage: true, supports_structured_outputs: true, - dialect, + supports_top_k: self.supports_top_k, + supports_tools: self.supports_tools, + supports_response_format: self.supports_response_format, + stream_usage_key: self.stream_usage_key.clone(), + max_tokens_key, + ..ChatSettings::default() }) } } @@ -385,7 +384,7 @@ fn validate_external_entry(entry: &ExternalProviderEntry) -> Result<(), AiMuxErr entry.name, entry.protocol ))); } - entry.profile.chat_profile().map_err(|e| { + entry.profile.chat_settings().map_err(|e| { AiMuxError::InvalidArgument(format!("external provider '{}': {e}", entry.name)) })?; Ok(()) @@ -481,7 +480,7 @@ fn build_external( query_params: None, fetch: None, transform_request_body: None, - profile: entry.profile.chat_profile()?, + chat: entry.profile.chat_settings()?, package: crate::shared::aisdk_pins::OPENAI_COMPATIBLE, }) } diff --git a/aimux-providers/src/provider_registry.json b/aimux-providers/src/provider_registry.json index f48d73019..4e5dc8608 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", @@ -1170,7 +1172,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/recraft.rs b/aimux-providers/src/recraft.rs index 83a74a5e4..32a7cb82e 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 3ab75f89b..66c77f022 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 21745b253..25bd2b897 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 9b5318490..06036f6db 100644 --- a/aimux-providers/src/runwayml.rs +++ b/aimux-providers/src/runwayml.rs @@ -25,9 +25,9 @@ use aimux_core::video_model::{ VideoCallOptions, VideoData, VideoFile, VideoFileData, VideoFrameType, 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}; const PROVIDER_NAME: &str = "runwayml"; const RUNWAY_VERSION: &str = "2024-11-06"; @@ -110,7 +110,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..a55eb2bc4 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, HttpEndpoint, 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?, - }) - }) - }), + HttpEndpoint { + 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 caa15e7de..bc617a791 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}; /// Failed-response handler of the Serper endpoint: `message`, and `type` /// else `code` as the provider code, from the body's `error` object, or from @@ -109,7 +109,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 d6a941890..ed73e5b26 100644 --- a/aimux-providers/src/shared/exchange.rs +++ b/aimux-providers/src/shared/exchange.rs @@ -11,17 +11,18 @@ use std::collections::HashMap; use std::sync::Arc; -use futures::future::BoxFuture; +use async_trait::async_trait; +use base64::Engine as _; use serde_json::Value; 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::TransformRequestBody; +use super::{Credential, ProviderHeaders, TransformRequestBody}; /// What one request is addressed to: the base URL and the provider headers, /// both resolved for this request. @@ -30,12 +31,159 @@ 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), then the user-agent + /// suffix when one is set (the AI SDK's `withUserAgentSuffix`). + WithUserHeaders { + provider: ProviderHeaders, + user: Option, + user_agent_suffix: Option>, + }, +} -/// The URL patterns a model fetches itself, for a model id. -pub(crate) type SupportedUrlsFn = Arc SupportedUrls + Send + Sync>; +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, + user_agent_suffix, + } => { + let layer = provider.resolve().await?; + let combined = match user { + Some(user) => combine_headers(&[&layer, &user.resolve().await?]), + None => layer, + }; + let Some(suffix) = user_agent_suffix else { + return Ok(combined); + }; + let mut headers = normalize_headers(combined).into_iter().collect(); + aimux_provider_utils::with_user_agent_suffix(&mut headers, suffix); + Ok(headers + .into_iter() + .map(|(name, value)| (name, Some(value))) + .collect()) + } + } + } +} + +/// Where the endpoint (base URL and provider headers) of a request comes +/// from, resolved in [`EndpointConfig::endpoint`] when a request is made. +/// +/// Most providers use [`HttpEndpoint`]. A package whose mode, credential and +/// region decide the URL and the headers together (Vertex AI, Bedrock, +/// Polly) implements it on its own plain-data struct. +#[async_trait] +pub(crate) trait ResolveEndpoint: Send + Sync { + /// The base URL and provider headers of one request. + async fn resolve(&self) -> Result; +} + +/// A base URL and headers, each resolved on its own. +pub(crate) struct HttpEndpoint { + pub(crate) base_url: BaseUrl, + pub(crate) headers: EndpointHeaders, +} + +#[async_trait] +impl ResolveEndpoint for HttpEndpoint { + async fn resolve(&self) -> Result { + Ok(Endpoint { + base_url: self.base_url.resolve()?, + headers: self.headers.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), + /// Patterns that depend on the base URL and the model id (the Gemini + /// API's). + ByModel { + base_url: String, + resolve: fn(&str, Option<&str>) -> SupportedUrls, + }, +} + +impl SupportedUrlsSource { + fn for_model(&self, model_id: &str) -> SupportedUrls { + match self { + Self::None => SupportedUrls::default(), + Self::Fixed(urls) => urls.clone(), + Self::ByModel { base_url, resolve } => resolve(base_url, Some(model_id)), + } + } +} /// What a model needs to talk to the API. Built by the provider that owns the /// model. @@ -44,12 +192,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: Arc, /// 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, /// Provider-level request-body rewrite. pub(crate) transform_request_body: Option, } @@ -60,51 +208,67 @@ impl EndpointConfig { pub(crate) fn fixed( provider: String, base_url: String, - headers: HeadersFn, + headers: impl Into, fetch: Option, transform_request_body: 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?, - }) - }) - }), + HttpEndpoint { + base_url: BaseUrl::Fixed(base_url), + headers: headers.into(), + }, fetch, - supported_urls: Arc::new(|_| SupportedUrls::default()), - transform_request_body, - } + ) + .with_transform_request_body(transform_request_body) } - /// 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: impl ResolveEndpoint + 'static, fetch: Option, ) -> Self { Self { provider, - endpoint, + endpoint: Arc::new(endpoint), fetch, - supported_urls: Arc::new(|_| SupportedUrls::default()), + supported_urls: SupportedUrlsSource::None, transform_request_body: 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 same config with a provider-level request-body rewrite. + #[must_use] + pub(crate) fn with_transform_request_body( + mut self, + transform_request_body: Option, + ) -> Self { + self.transform_request_body = transform_request_body; + 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). /// @@ -117,7 +281,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 021e65f1d..15bc9edd9 100644 --- a/aimux-providers/src/shared/mod.rs +++ b/aimux-providers/src/shared/mod.rs @@ -1,7 +1,9 @@ //! 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. pub(crate) mod aisdk_pins; mod discovery; @@ -9,7 +11,10 @@ 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, HttpEndpoint, ResolveEndpoint, + 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, @@ -21,7 +26,8 @@ use serde_json::Value; use aimux_core::AiMuxError; use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, combine_headers, load_api_key, + FetchFunction, HeaderMapOpt, 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 @@ -126,60 +132,83 @@ 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 +/// 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)>, } -/// 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). -pub(crate) fn credential_headers( - credential: Credential, - scheme: AuthScheme, - fixed: Vec<(String, String)>, - user: Option, -) -> HeadersFn { - 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, - }) +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 + } + + 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()) + } } /// The transport of a pinned package: `fetch` (the process default, resolved diff --git a/aimux-providers/src/stability.rs b/aimux-providers/src/stability.rs index d150ab575..3de3d8d1f 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 9bc22c5dd..1e4fbac95 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 6647e256b..2945c5de0 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. /// Failed-response handler of the TinyFish endpoint: `message`, and `type` @@ -118,7 +118,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/anthropic_model.rs b/aimux-providers/src/vertex/anthropic_model.rs index 5f5511e10..3b365fa54 100644 --- a/aimux-providers/src/vertex/anthropic_model.rs +++ b/aimux-providers/src/vertex/anthropic_model.rs @@ -20,17 +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 crate::shared::{Credential, ProviderHeaders}; + +use super::Resolver; use crate::anthropic::AnthropicMessagesModel; -use crate::anthropic::config::{AnthropicModelConfig, AnthropicModelHooks, TransformRequestBody}; +use crate::anthropic::config::{ + AnthropicEndpoint, AnthropicModelConfig, AnthropicModelHooks, MessagesUrl, TransformRequestBody, +}; use crate::anthropic::options::options_name_of; -use crate::shared::Endpoint; /// `anthropic_version` envelope value required by the Vertex AI `rawPredict` / /// `streamRawPredict` endpoints. @@ -42,11 +44,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,67 +58,32 @@ fn raw_predict_envelope(body: Value) -> Value { Value::Object(envelope) } -/// The configuration of one request against `endpoint`. -fn request_config( - endpoint: Endpoint, +/// The model `model_id` on the Vertex endpoint `resolver` resolves for each +/// request. +pub(super) fn model( + model_id: &str, + resolver: Arc, fetch: Option, transform_request_body: Option, -) -> AnthropicModelConfig { - let Endpoint { base_url, headers } = endpoint; - let models = base_url.clone(); - let url_base = base_url.clone(); - AnthropicModelConfig { +) -> VertexAnthropicModel { + // 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(), - url: Arc::new(move |path| format!("{url_base}{path}")), - headers: Resolvable::Value(headers), + headers: ProviderHeaders::bearer(Credential::None, Vec::new(), None).into(), fetch, supported_urls: SupportedUrls::default(), transform_request_body, - base_url, + base_url: String::new(), 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)), + 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, }, - resolve: None, - } -} - -/// The model `model_id` on the Vertex endpoint `endpoint` resolves to for each -/// request. -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)) }) - })); + endpoint: AnthropicEndpoint::Vertex(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 ff0c61ce2..d8ee133cb 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(), self.0.transform_request_body.clone(), )) diff --git a/aimux-providers/src/vertex/gemini_transcription.rs b/aimux-providers/src/vertex/gemini_transcription.rs index 8afe88d0f..d5364c634 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 = aimux_core::Timestamp::now(); let config = crate::google::transcription::TranscriptionOptions::parse( Namespace::Vertex.read(options.provider_options.as_ref()), @@ -226,7 +226,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 53f25ebbd..494c83ee3 100644 --- a/aimux-providers/src/vertex/mod.rs +++ b/aimux-providers/src/vertex/mod.rs @@ -17,12 +17,15 @@ 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, TransformRequestBody}; +use crate::shared::{ + Endpoint, EndpointConfig, ResolveEndpoint, SupportedUrlsSource, TransformRequestBody, +}; mod anthropic_model; mod anthropic_provider; @@ -142,7 +145,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, } @@ -153,14 +156,67 @@ 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)] +struct GeminiEndpoint { + resolver: Arc, + tuned: bool, +} + +#[async_trait::async_trait] +impl ResolveEndpoint for GeminiEndpoint { + /// The base URL and provider headers of one request. + 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. -struct Resolver { +pub(crate) struct Resolver { api_key: Option>, express_mode: bool, location: Option, @@ -253,7 +309,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, @@ -263,13 +319,16 @@ impl Resolver { let headers = match publisher { // google-vertex-provider-base.ts:270-276 appends the package part. Publisher::Google => { - aimux_provider_utils::headers::with_user_agent_suffix_fn( - Resolvable::Value(headers), - crate::shared::aisdk_pins::GOOGLE_VERTEX.name, - crate::shared::aisdk_pins::GOOGLE_VERTEX.version, - ) - .resolve() - .await? + let package = crate::shared::aisdk_pins::GOOGLE_VERTEX; + let mut headers = normalize_headers(headers).into_iter().collect(); + with_user_agent_suffix( + &mut headers, + &format!("ai-sdk-{}/{}", package.name, package.version), + ); + headers + .into_iter() + .map(|(name, value)| (name, Some(value))) + .collect() } // The Anthropic sub-provider passes the caller's headers as they // are (google-vertex-anthropic-provider.ts:206-210). @@ -365,54 +424,24 @@ 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) - }) - }), - fetch: self.fetch.clone(), - supported_urls: Arc::new(|_| supported_urls()), - transform_request_body: self.transform_request_body.clone(), - } + EndpointConfig::new( + provider.to_string(), + GeminiEndpoint { + resolver: self.resolver.clone(), + tuned, + }, + self.fetch.clone(), + ) + .with_supported_urls(SupportedUrlsSource::Fixed(supported_urls())) + .with_transform_request_body(self.transform_request_body.clone()) } /// 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 6df1aa0e5..1876d0b97 100644 --- a/aimux-providers/src/vertex/model.rs +++ b/aimux-providers/src/vertex/model.rs @@ -81,7 +81,7 @@ impl LanguageModel for VertexModel { } async 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/partner_provider.rs b/aimux-providers/src/vertex/partner_provider.rs index 67af4ea13..e9973c90c 100644 --- a/aimux-providers/src/vertex/partner_provider.rs +++ b/aimux-providers/src/vertex/partner_provider.rs @@ -26,7 +26,7 @@ use aimux_provider_utils::{ load_setting, validate_base_url, }; -use crate::openai_compatible::config::BaseUrl; +use crate::openai_compatible::config::{BaseUrl, ResolveBaseUrl, ResolveFetchHeaders}; use crate::openai_compatible::{ OpenAICompatibleChatModel, OpenAICompatibleProvider, OpenAICompatibleProviderSettings, }; @@ -82,13 +82,77 @@ impl std::fmt::Debug for VertexPartnerProviderSettings { } } +/// The base URL of a partner sub-provider with no explicit `base_url`, built +/// from the project and location read when a request is made. +#[derive(Clone, Debug)] +struct PartnerBaseUrl { + project: Option, + location: Option, + /// MaaS picks the location's host (google-vertex-maas-provider.ts:91-106); + /// xAI always calls the global host (google-vertex-xai-provider.ts:159-164). + regional_host: bool, +} + +impl ResolveBaseUrl for PartnerBaseUrl { + /// The base URL of one request. + /// + /// # Errors + /// + /// `AiMuxError::LoadSetting` when no project is set. + fn resolve(&self) -> Result { + let project = load_setting(self.project.as_deref(), PROJECT_ENV_VAR, "project")?; + let location = load_optional_setting(self.location.as_deref(), LOCATION_ENV_VAR) + .unwrap_or_else(|| "global".to_string()); + let host = if self.regional_host { + location_host(&location) + } else { + "aiplatform.googleapis.com".to_string() + }; + Ok(format!( + "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi" + )) + } +} + +/// The headers the Node wrapper's fetch adds over the request's own: the +/// provider headers, then `Authorization` with the Google OAuth token. The +/// Node wrapper passes no headers to `createOpenAICompatible` and spreads the +/// resolved headers, then `Authorization`, over the request's headers in its +/// fetch, so they win over the per-call headers, and a `user-agent` among +/// them replaces the whole header (google-vertex-maas-provider-node.ts:34-58, +/// google-vertex-xai-provider-node.ts:33-55). +#[derive(Clone)] +struct PartnerHeaders { + auth: Arc, + user: Option, +} + +#[async_trait::async_trait] +impl ResolveFetchHeaders for PartnerHeaders { + /// The headers of one request. + /// + /// # Errors + /// + /// The token or the caller's header producer error. + async fn resolve(&self) -> Result { + let token = self.auth.access_token().await?; + let user = match &self.user { + Some(headers) => headers.resolve().await?, + None => HeaderMapOpt::new(), + }; + let mut bearer = HeaderMapOpt::new(); + bearer.insert("Authorization".to_string(), Some(format!("Bearer {token}"))); + Ok(combine_headers(&[&user, &bearer])) + } +} + /// The parts every partner sub-provider shares: the base URL source and the /// headers the Node wrapper's fetch adds over the request's own (the /// provider headers, then `Authorization`). fn partner_parts( settings: &VertexPartnerProviderSettings, regional_host: bool, -) -> Result<(BaseUrl, HeadersFn), AiMuxError> { +) -> Result<(BaseUrl, Arc), AiMuxError> { let explicit = settings .base_url .as_deref() @@ -96,25 +160,11 @@ fn partner_parts( .filter(|url| !url.is_empty()); let base_url = match explicit { Some(url) => BaseUrl::Fixed(validate_base_url(url)?), - None => { - let project = settings.project.clone(); - let location = settings.location.clone(); - BaseUrl::Lazy(Arc::new(move || { - let project = load_setting(project.as_deref(), PROJECT_ENV_VAR, "project")?; - let location = load_optional_setting(location.as_deref(), LOCATION_ENV_VAR) - .unwrap_or_else(|| "global".to_string()); - // MaaS picks the location's host (google-vertex-maas-provider.ts:91-106); - // xAI always calls the global host (google-vertex-xai-provider.ts:159-164). - let host = if regional_host { - location_host(&location) - } else { - "aiplatform.googleapis.com".to_string() - }; - Ok(format!( - "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi" - )) - })) - } + None => BaseUrl::Dynamic(Arc::new(PartnerBaseUrl { + project: settings.project.clone(), + location: settings.location.clone(), + regional_host, + })), }; let mut auth_options = settings.google_auth_options.clone().unwrap_or_default(); @@ -124,27 +174,9 @@ fn partner_parts( ..Default::default() }; } - let auth = Arc::new(GoogleAuth::new(auth_options, settings.fetch.clone())); - let user_headers = settings.headers.clone(); - // The Node wrapper passes no headers to `createOpenAICompatible` and - // spreads the resolved headers, then `Authorization`, over the request's - // headers in its fetch, so they win over the per-call headers, and a - // `user-agent` among them replaces the whole header - // (google-vertex-maas-provider-node.ts:34-58, - // google-vertex-xai-provider-node.ts:33-55). - let headers = Resolvable::from_async_fn(move || { - let auth = auth.clone(); - let user_headers = user_headers.clone(); - async move { - let token = auth.access_token().await?; - let user = match user_headers { - Some(headers) => headers.resolve().await?, - None => HeaderMapOpt::new(), - }; - let mut bearer = HeaderMapOpt::new(); - bearer.insert("Authorization".to_string(), Some(format!("Bearer {token}"))); - Ok(combine_headers(&[&user, &bearer])) - } + let headers: Arc = Arc::new(PartnerHeaders { + auth: Arc::new(GoogleAuth::new(auth_options, settings.fetch.clone())), + user: settings.headers.clone(), }); Ok((base_url, headers)) } @@ -188,7 +220,12 @@ pub fn create_google_vertex_maas( ..Default::default() }, base_url, - Resolvable::Value(HeaderMapOpt::new()), + crate::shared::ProviderHeaders::new( + crate::shared::Credential::None, + crate::shared::AuthScheme::Bearer, + Vec::new(), + None, + ), Some(fetch_headers), crate::shared::aisdk_pins::GOOGLE_VERTEX_OPENAI_COMPATIBLE, ) @@ -275,7 +312,12 @@ pub fn create_google_vertex_xai( ..Default::default() }, base_url, - Resolvable::Value(HeaderMapOpt::new()), + crate::shared::ProviderHeaders::new( + crate::shared::Credential::None, + crate::shared::AuthScheme::Bearer, + Vec::new(), + None, + ), Some(fetch_headers), crate::shared::aisdk_pins::GOOGLE_VERTEX_OPENAI_COMPATIBLE, )?; diff --git a/aimux-providers/src/vertex/transcription.rs b/aimux-providers/src/vertex/transcription.rs index 97262ebf8..93909b576 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 ───────────────────────────────────────────────────────── @@ -110,14 +110,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, @@ -161,7 +161,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 eed86624b..bbdc90bc4 100644 --- a/aimux-providers/src/voyage/mod.rs +++ b/aimux-providers/src/voyage/mod.rs @@ -26,9 +26,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 { @@ -114,7 +114,7 @@ pub fn create_voyage(settings: VoyageProviderSettings) -> Result &'static VoyageProvider { pub struct VoyageProvider { name: String, base_url: String, - headers: HeadersFn, + headers: ProviderHeaders, fetch: Option, transform_request_body: Option, } diff --git a/aimux-providers/src/xai/mod.rs b/aimux-providers/src/xai/mod.rs index 3bb8a49e4..29f81edaf 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}; pub(crate) const DEFAULT_BASE_URL: &str = "https://api.x.ai/v1"; pub(crate) const API_KEY_ENV_VAR: &str = "XAI_API_KEY"; @@ -323,16 +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 3a421776a..525327596 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. /// Failed-response handler of the You.com endpoint: `message`, and `type` @@ -119,7 +119,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, } diff --git a/aimux-providers/tests/dataforseo_test.rs b/aimux-providers/tests/dataforseo_test.rs index bb8c09422..908f2de21 100644 --- a/aimux-providers/tests/dataforseo_test.rs +++ b/aimux-providers/tests/dataforseo_test.rs @@ -198,6 +198,49 @@ async fn do_search_sends_correct_basic_auth_header() { ); } +#[serial_test::serial] +#[tokio::test] +async fn login_comes_from_dataforseo_login_per_request() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v3/serp/google/organic/live/advanced")) + .and(header("authorization", expected_basic_auth())) + .respond_with(ResponseTemplate::new(200).set_body_json(search_response_body())) + .mount(&server) + .await; + let saved = std::env::var("DATAFORSEO_LOGIN").ok(); + let provider = create_dataforseo(DataforseoProviderSettings { + login: None, + password: Some(Resolvable::Value(PASSWORD.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }) + .unwrap(); + let model = provider.search_model(); + + unsafe { + std::env::set_var("DATAFORSEO_LOGIN", LOGIN); + } + let found = model.do_search(&opts("rust", Some(10))).await; + + unsafe { + std::env::remove_var("DATAFORSEO_LOGIN"); + } + let missing = model.do_search(&opts("rust", Some(10))).await; + + unsafe { + if let Some(value) = saved { + std::env::set_var("DATAFORSEO_LOGIN", value); + } + } + assert_eq!(found.unwrap().results.len(), 2); + assert!( + matches!(missing, Err(AiMuxError::LoadApiKey { .. })), + "{missing:?}" + ); + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + // -- error mapping ----------------------------------------------------------- #[tokio::test] diff --git a/aimux-providers/tests/searxng_test.rs b/aimux-providers/tests/searxng_test.rs index 8cf8b2303..a398a32df 100644 --- a/aimux-providers/tests/searxng_test.rs +++ b/aimux-providers/tests/searxng_test.rs @@ -185,3 +185,40 @@ fn model_id_is_searxng_search() { assert_eq!(model.model_id(), "searxng-search"); assert_eq!(model.provider(), "searxng.search"); } + +// -- instance URL from the environment --------------------------------------- + +/// Without `base_url`, the instance URL is `SEARXNG_URL`, read for each +/// request: set, the request goes there; unset, the request fails with +/// `LoadSetting` and nothing is sent. +#[serial_test::serial] +#[tokio::test] +async fn instance_url_comes_from_searxng_url_per_request() { + let server = MockServer::start().await; + mount_search_mock(&server).await; + let saved = std::env::var("SEARXNG_URL").ok(); + let provider = create_searxng(SearxngProviderSettings::default()).unwrap(); + let model = provider.search_model(); + + unsafe { + std::env::set_var("SEARXNG_URL", server.uri()); + } + let found = model.do_search(&opts("rust language")).await; + + unsafe { + std::env::remove_var("SEARXNG_URL"); + } + let missing = model.do_search(&opts("rust language")).await; + + unsafe { + if let Some(value) = saved { + std::env::set_var("SEARXNG_URL", value); + } + } + assert_eq!(found.unwrap().results.len(), 2); + assert!( + matches!(missing, Err(AiMuxError::LoadSetting { .. })), + "{missing:?}" + ); + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} diff --git a/docs/contributing/adding-a-provider.md b/docs/contributing/adding-a-provider.md index fdb657d7a..e5cf36214 100644 --- a/docs/contributing/adding-a-provider.md +++ b/docs/contributing/adding-a-provider.md @@ -38,7 +38,10 @@ its own, ...), it is case 2: add its own package next to `groq/` and (`{param}` placeholders, each declared with an optional `env` list 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 + only max-token key the vendor accepts); `profile.convert_usage` + (`"alibaba"` or `"moonshotai"`: a usage shape that differs from the + generic OpenAI one; a new shape is a new converter in `preset.rs`). 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 error shapes, it is case 2. diff --git a/scripts/check_provider_boundaries.sh b/scripts/check_provider_boundaries.sh index add64bb2e..ee06250fe 100755 --- a/scripts/check_provider_boundaries.sh +++ b/scripts/check_provider_boundaries.sh @@ -149,6 +149,20 @@ rule4b_hits="$( )" violation 'groq / deepseek literal in the shared OpenAI or compatible package' "$rule4b_hits" +# `shared/` and `openai_compatible/` never reach into a vendor package: a +# vendor's endpoint, base URL, headers or usage accounting is its own data +# behind a trait or hook of the shared layer. Their code (doc links aside) +# names only the shared layers, the preset registry and `runtime_decision`. +rule4c_hits="$( + grep -rnE 'crate::[a-z_]+' aimux-providers/src/shared \ + aimux-providers/src/openai_compatible --include='*.rs' | + awk '{ code = $0; sub(/^[^:]*:[0-9]+:/, "", code) } + code ~ /^[[:space:]]*\/\// { next } + { gsub(/crate::(shared|openai|openai_compatible|preset|runtime_decision)([^a-z_]|$)/, "", code) } + code ~ /crate::[a-z_]/' || true +)" +violation 'shared or openai_compatible names a vendor package (implement its trait in the vendor module)' "$rule4c_hits" + # ── Rule 5: the builder-era compat names stay removed in the providers crate ── rule5_hits="$( grep -rnE 'OpenAIConfig\b|OpenAIConfigProvider|OpenAICompatProfile|PLACEHOLDER_API_KEY|openai_legacy' \