diff --git a/aimux-providers/layout.toml b/aimux-providers/layout.toml index dec8d7b32..bc2bafced 100644 --- a/aimux-providers/layout.toml +++ b/aimux-providers/layout.toml @@ -131,3 +131,370 @@ status = "mapped" upstream = "src/version.ts" status = "not-applicable" reason = "the crate version is used instead (RFC-0046 §4.2)" + +[vertex] +package = "@ai-sdk/google-vertex" +# One aimux entry point: the edge entries (edge/, */edge/) and the *-node entries map to +# the same Rust file as the provider they wrap. google-vertex-language-model.test.ts and +# google-vertex-location-validation.test.ts have no source file of their own name; every +# case in them is ported, so they need no exclusions. + +[[vertex.files]] +upstream = "src/anthropic/edge/google-vertex-anthropic-provider-edge.ts" +rust = "vertex/anthropic_provider.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/anthropic/edge/index.ts" +rust = "vertex/anthropic_provider.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/anthropic/google-vertex-anthropic-options.ts" +rust = "vertex/anthropic_provider.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/anthropic/google-vertex-anthropic-provider-node.ts" +rust = "vertex/anthropic_provider.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/anthropic/google-vertex-anthropic-provider.ts" +rust = "vertex/anthropic_provider.rs" +status = "mapped" + +[vertex.files.tests] +sub-feature = [ + { case = "google-vertex-anthropic-provider › should include googleVertexAnthropicTools (subset of anthropicTools)", issue = 275, reason = "VertexAnthropicProvider has no tools (no googleVertexAnthropicTools set)" }, +] +ts-only = [ + { case = "google-vertex-anthropic-provider › should throw an error when using new keyword", reason = "JavaScript `new` on a callable provider" }, +] + +[[vertex.files]] +upstream = "src/anthropic/index.ts" +rust = "vertex/anthropic_provider.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/edge/google-vertex-auth-edge.ts" +rust = "vertex/auth.rs" +status = "mapped" + +[vertex.files.tests] +sub-feature = [ + { case = "Google Vertex Edge Auth › should load credentials from environment variables", issue = 275, reason = "vertex/auth.rs does not read GOOGLE_CLIENT_EMAIL / GOOGLE_PRIVATE_KEY / GOOGLE_PRIVATE_KEY_ID" }, + { case = "Google Vertex Edge Auth › should throw error when client email is missing", issue = 275, reason = "vertex/auth.rs does not read GOOGLE_CLIENT_EMAIL / GOOGLE_PRIVATE_KEY / GOOGLE_PRIVATE_KEY_ID" }, + { case = "Google Vertex Edge Auth › should throw error when private key is missing", issue = 275, reason = "vertex/auth.rs does not read GOOGLE_CLIENT_EMAIL / GOOGLE_PRIVATE_KEY / GOOGLE_PRIVATE_KEY_ID" }, + { case = "Google Vertex Edge Auth › should work with or without private key ID", issue = 275, reason = "vertex/auth.rs does not read GOOGLE_CLIENT_EMAIL / GOOGLE_PRIVATE_KEY / GOOGLE_PRIVATE_KEY_ID" }, + { case = "Google Vertex Edge Auth › should handle newlines in private key from env vars", issue = 275, reason = "vertex/auth.rs does not read GOOGLE_CLIENT_EMAIL / GOOGLE_PRIVATE_KEY / GOOGLE_PRIVATE_KEY_ID" }, +] +ts-only = [ + { case = "Google Vertex Edge Auth › should throw error with invalid credentials", reason = "the edge auth module's own error text (Token request failed: …); aimux has one auth path, which follows google-auth-library" }, + { case = "Google Vertex Edge Auth › should throw error when token request fails", reason = "the edge auth module's own error text (Token request failed: …); aimux has one auth path, which follows google-auth-library" }, + { case = "Google Vertex Edge Auth › should include correct user-agent header", reason = "the edge auth module's own token-request user agent; aimux has one auth path, which follows google-auth-library" }, +] + +[[vertex.files]] +upstream = "src/edge/google-vertex-provider-edge.ts" +rust = "vertex/mod.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/edge/index.ts" +rust = "vertex/mod.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/gemini-transcription/google-vertex-gemini-transcription-model-options.ts" +rust = "vertex/gemini_transcription.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/gemini-transcription/google-vertex-gemini-transcription-model.ts" +rust = "vertex/gemini_transcription.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-auth-google-auth-library.ts" +rust = "vertex/auth.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-cloud-tts-speech-model.ts" +status = "missing" +issue = 275 + +[[vertex.files]] +upstream = "src/google-vertex-config.ts" +rust = "vertex/mod.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-embedding-model-options.ts" +rust = "vertex/embedding.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-embedding-model.ts" +rust = "vertex/embedding.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-error.ts" +rust = "google/mod.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-image-model-options.ts" +rust = "vertex/image.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-image-model.ts" +rust = "vertex/image.rs" +status = "mapped" + +[vertex.files.tests] +sub-feature = [ + { case = "GoogleVertexImageModel › doGenerate › should classify prompt blocks as terminal", issue = 275, reason = "ImageResult has no is_retryable" }, +] + +[[vertex.files]] +upstream = "src/google-vertex-image-settings.ts" +rust = "vertex/image.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-options.ts" +rust = "vertex/mod.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-provider-base.ts" +rust = "vertex/mod.rs" +status = "mapped" + +[vertex.files.tests] +sub-feature = [ + { case = "google-vertex-provider-base › should configure the tool result download size limit", issue = 275, reason = "VertexProviderSettings has no tool result download setting (downloadToolResultFiles)" }, + { case = "google-vertex-provider-base › should create an interactions model targeting the location-scoped interactions resource", issue = 275, reason = "no VertexProvider::interactions (Google interactions model)" }, + { case = "google-vertex-provider-base › should throw for interactions models when an Express Mode API key is set", issue = 275, reason = "no VertexProvider::interactions (Google interactions model)" }, + { case = "google-vertex-provider-base › should route chirp speech models to the Cloud Text-to-Speech model", issue = 275, reason = "no Vertex Cloud Text-to-Speech model; speech(\"chirp…\") returns UnsupportedFunctionality" }, + { case = "google-vertex-provider-base › should throw for chirp speech models when an Express Mode API key is set", issue = 275, reason = "no Vertex Cloud Text-to-Speech model; speech(\"chirp…\") returns UnsupportedFunctionality" }, + { case = "google-vertex-provider-base › should expose tools", issue = 275, reason = "no VertexProvider::tools (googleSearch, urlContext, codeExecution builders)" }, +] +ts-only = [ + { case = "google-vertex-provider-base › should throw an error when using new keyword", reason = "JavaScript `new` on a callable provider" }, +] + +[[vertex.files]] +upstream = "src/google-vertex-provider.ts" +rust = "vertex/mod.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-speech-model-options.ts" +rust = "google/speech.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-tools.ts" +status = "missing" +issue = 275 + +[[vertex.files]] +upstream = "src/google-vertex-transcription-model-options.ts" +rust = "vertex/transcription.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-transcription-model.ts" +rust = "vertex/transcription.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-video-model-options.ts" +rust = "vertex/video.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/google-vertex-video-model.ts" +rust = "vertex/video.rs" +status = "mapped" + +[vertex.files.tests] +sub-feature = [ + { case = "GoogleVertexVideoModel › doStatus › should return error status on operation error", issue = 275, reason = "VideoOperationStatus has no Error variant; a failed operation is Err(ApiCall)" }, +] + +[[vertex.files]] +upstream = "src/google-vertex-video-settings.ts" +rust = "vertex/video.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/index.ts" +rust = "vertex/mod.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/maas/edge/google-vertex-maas-provider-edge.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] +ts-only = [ + { case = "google-vertex-maas-provider-edge › should create provider with auth wrapper", reason = "checks the arguments handed to the mocked base provider (fetch: Function, headers: undefined); aimux has no separate base/wrapper layer" }, + { case = "google-vertex-maas-provider-edge › should set headers to undefined in base provider call", reason = "checks the arguments handed to the mocked base provider; aimux has no separate base/wrapper layer" }, +] + +[[vertex.files]] +upstream = "src/maas/edge/index.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/maas/google-vertex-maas-options.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/maas/google-vertex-maas-provider-node.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] +ts-only = [ + { case = "google-vertex-maas-provider-node › should create provider with auth wrapper", reason = "checks the arguments handed to the mocked base provider (fetch: Function, headers: undefined); aimux has no separate base/wrapper layer" }, + { case = "google-vertex-maas-provider-node › should set headers to undefined in base provider call", reason = "checks the arguments handed to the mocked base provider; aimux has no separate base/wrapper layer" }, +] + +[[vertex.files]] +upstream = "src/maas/google-vertex-maas-provider.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] +ts-only = [ + { case = "google-vertex-maas-provider › should not call createOpenAICompatible at provider creation time", reason = "counts calls to the mocked createOpenAICompatible; aimux builds the OpenAI-compatible provider eagerly" }, + { case = "google-vertex-maas-provider › should cache the provider after first access", reason = "counts calls to the mocked createOpenAICompatible; aimux builds the OpenAI-compatible provider eagerly" }, + { case = "google-vertex-maas-provider › should not pass headers to openai-compatible provider", reason = "checks the arguments handed to the mocked createOpenAICompatible; aimux applies the headers in its fetch layer" }, +] + +[[vertex.files]] +upstream = "src/maas/index.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/version.ts" +status = "not-applicable" + +[[vertex.files]] +upstream = "src/xai/edge/google-vertex-xai-provider-edge.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] +ts-only = [ + { case = "google-vertex-xai-provider-edge › should create provider with auth wrapper", reason = "checks the arguments handed to the mocked base provider (fetch: Function, headers: undefined); aimux has no separate base/wrapper layer" }, +] + +[[vertex.files]] +upstream = "src/xai/edge/index.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/xai/google-vertex-xai-options.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] + +[[vertex.files]] +upstream = "src/xai/google-vertex-xai-provider-node.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] +ts-only = [ + { case = "google-vertex-xai-provider-node › should create provider with auth wrapper", reason = "checks the arguments handed to the mocked base provider (fetch: Function, headers: undefined); aimux has no separate base/wrapper layer" }, +] + +[[vertex.files]] +upstream = "src/xai/google-vertex-xai-provider.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] +ts-only = [ + { case = "google-vertex-xai-provider › should not call createOpenAICompatible at provider creation time", reason = "counts calls to the mocked createOpenAICompatible; aimux builds the OpenAI-compatible provider eagerly" }, + { case = "google-vertex-xai-provider › should cache the OpenAI-compatible provider after first access", reason = "counts calls to the mocked createOpenAICompatible; aimux builds the OpenAI-compatible provider eagerly" }, + { case = "google-vertex-xai-provider › should throw an error when using new keyword", reason = "JavaScript `new` on a callable provider" }, +] + +[[vertex.files]] +upstream = "src/xai/index.ts" +rust = "vertex/partner_provider.rs" +status = "mapped" + +[vertex.files.tests] diff --git a/aimux-providers/src/vertex/auth.rs b/aimux-providers/src/vertex/auth.rs index f18d18ec4..675156c34 100644 --- a/aimux-providers/src/vertex/auth.rs +++ b/aimux-providers/src/vertex/auth.rs @@ -100,7 +100,7 @@ struct CachedToken { } pub(super) struct GoogleAuth { - options: GoogleAuthOptions, + pub(super) options: GoogleAuthOptions, token: Mutex>, fetch: Option, } diff --git a/aimux-providers/src/vertex/mod.rs b/aimux-providers/src/vertex/mod.rs index 8bcb36870..93f015c52 100644 --- a/aimux-providers/src/vertex/mod.rs +++ b/aimux-providers/src/vertex/mod.rs @@ -608,3 +608,6 @@ impl ProviderDiscovery for VertexProvider { Box::pin(async move { crate::google::list_models_once(&config).await }) } } + +#[cfg(test)] +mod tests; diff --git a/aimux-providers/src/vertex/tests/anthropic/anthropic_provider.rs b/aimux-providers/src/vertex/tests/anthropic/anthropic_provider.rs new file mode 100644 index 000000000..36b8e5744 --- /dev/null +++ b/aimux-providers/src/vertex/tests/anthropic/anthropic_provider.rs @@ -0,0 +1,315 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 +//! `src/anthropic/google-vertex-anthropic-provider.test.ts`. +//! +//! Upstream mocks the `AnthropicLanguageModel` constructor and inspects the +//! config handed to it. aimux builds the model itself, so each config value +//! is observed through what it does: `baseURL` and `buildRequestUrl` through +//! the request URL, `transformRequestBody` through the request body, `headers` +//! through the request headers, `supportedUrls` through +//! `LanguageModel::supported_urls`. Requests need a token, which the +//! `generate_auth_token` setting supplies (upstream's base provider has no +//! auth); the transport is a mock (upstream's `fetch` is `undefined`). + +use std::collections::BTreeMap; +use std::sync::Arc; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_core::provider::Provider; +use aimux_provider_utils::Resolvable; +use serde_json::json; + +use crate::test_support::mock_fetch::{Canned, MockFetch, Seen}; +use crate::vertex::{ + VertexAnthropicProvider, VertexAnthropicProviderSettings, create_google_vertex_anthropic, +}; + +/// A minimal Messages API response. +fn messages_response() -> Canned { + Canned::json(&json!({ + "id": "msg_1", "type": "message", "role": "assistant", "model": "claude", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", "stop_sequence": null, + "usage": {"input_tokens": 1, "output_tokens": 1} + })) +} + +fn prompt() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +/// Settings with the mock transport and a fixed token. +fn settings(fetch: &Arc) -> VertexAnthropicProviderSettings { + VertexAnthropicProviderSettings { + fetch: Some(fetch.transport()), + generate_auth_token: Some(Resolvable::Value(Some("test-token".into()))), + ..Default::default() + } +} + +/// Generate once with `model_id` and return the request the transport saw. +async fn request( + provider: &VertexAnthropicProvider, + fetch: &Arc, + model_id: &str, +) -> Seen { + provider + .call(model_id) + .do_generate(&prompt()) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + seen[0].clone() +} + +/// The request headers other than the token (supplied by the harness) and +/// the transport's own `content-type` and `user-agent`: upstream's `headers` +/// config value. +fn provider_headers(seen: &Seen) -> BTreeMap { + seen.headers + .iter() + .filter(|(name, _)| { + !matches!( + name.as_str(), + "authorization" | "content-type" | "user-agent" + ) + }) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +mod google_vertex_anthropic_provider { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider.test.ts › google-vertex-anthropic-provider › should create a language model with default settings + #[tokio::test] + async fn should_create_a_language_model_with_default_settings() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("test-project".into()), + location: Some("test-location".into()), + ..settings(&fetch) + }) + .unwrap(); + let model = provider.call("test-model-id"); + assert_eq!(model.model_id(), "test-model-id"); + assert_eq!(model.provider(), "googleVertex.anthropic.messages"); + + let seen = request(&provider, &fetch, "test-model-id").await; + // baseURL + assert!( + seen.url.contains( + "/projects/test-project/locations/test-location/publishers/anthropic/models" + ), + "{}", + seen.url + ); + // buildRequestUrl + assert!( + seen.url.ends_with("/test-model-id:rawPredict"), + "{}", + seen.url + ); + // transformRequestBody + let body = seen.json_body(); + assert!(body.get("model").is_none(), "{body}"); + assert_eq!(body["anthropic_version"], "vertex-2023-10-16"); + // headers: an object (only the token here) + assert_eq!(seen.headers["authorization"], "Bearer test-token"); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider.test.ts › google-vertex-anthropic-provider › should pass baseURL to the model when created + #[tokio::test] + async fn should_pass_base_url_to_the_model_when_created() { + let custom_base_url = "https://custom-url.com"; + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("test-project".into()), + base_url: Some(custom_base_url.into()), + ..settings(&fetch) + }) + .unwrap(); + + let seen = request(&provider, &fetch, "test-model-id").await; + assert_eq!( + seen.url, + format!("{custom_base_url}/test-model-id:rawPredict") + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider.test.ts › google-vertex-anthropic-provider › should throw NoSuchModelError for textEmbeddingModel + #[test] + fn should_throw_no_such_model_error_for_text_embedding_model() { + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("test-project".into()), + ..Default::default() + }) + .unwrap(); + + assert!(matches!( + provider.embedding_model("invalid-model-id"), + Err(AiMuxError::NoSuchModel { .. }) + )); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider.test.ts › google-vertex-anthropic-provider › should pass custom headers to the model constructor + #[tokio::test] + async fn should_pass_custom_headers_to_the_model_constructor() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("test-project".into()), + location: Some("us-east5".into()), + headers: Some(Resolvable::Value( + [( + "Custom-Header".to_string(), + Some("custom-value".to_string()), + )] + .into(), + )), + ..settings(&fetch) + }) + .unwrap(); + + let seen = request(&provider, &fetch, "test-model-id").await; + assert_eq!( + provider_headers(&seen), + BTreeMap::from([("custom-header".to_string(), "custom-value".to_string())]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider.test.ts › google-vertex-anthropic-provider › should create a Google Vertex Anthropic provider instance with custom settings + #[test] + fn should_create_a_google_vertex_anthropic_provider_instance_with_custom_settings() { + let custom_provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("custom-project".into()), + location: Some("custom-location".into()), + base_url: Some("https://custom.base.url".into()), + headers: Some(Resolvable::Value( + [("Custom-Header".to_string(), Some("value".to_string()))].into(), + )), + ..Default::default() + }); + + let custom_provider = custom_provider.unwrap(); + // Callable (`call`), with a language model factory. + let _ = custom_provider.call("test-model-id"); + assert!(custom_provider.language_model("test-model-id").is_ok()); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider.test.ts › google-vertex-anthropic-provider › should not support URL sources to force base64 conversion + #[tokio::test] + async fn should_not_support_url_sources_to_force_base64_conversion() { + let provider = + create_google_vertex_anthropic(VertexAnthropicProviderSettings::default()).unwrap(); + let model = provider.call("test-model-id"); + assert_eq!(model.model_id(), "test-model-id"); + + assert!(model.supported_urls().await.0.is_empty()); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider.test.ts › google-vertex-anthropic-provider › should use correct URL for global location + #[tokio::test] + async fn should_use_correct_url_for_global_location() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + ..settings(&fetch) + }) + .unwrap(); + assert_eq!( + provider.call("test-model-id").provider(), + "googleVertex.anthropic.messages" + ); + + let seen = request(&provider, &fetch, "test-model-id").await; + assert_eq!( + seen.url, + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/publishers/anthropic/models/test-model-id:rawPredict" + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider.test.ts › google-vertex-anthropic-provider › should use region-prefixed URL for non-global locations + #[tokio::test] + async fn should_use_region_prefixed_url_for_non_global_locations() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("test-project".into()), + location: Some("us-east5".into()), + ..settings(&fetch) + }) + .unwrap(); + assert_eq!( + provider.call("test-model-id").provider(), + "googleVertex.anthropic.messages" + ); + + let seen = request(&provider, &fetch, "test-model-id").await; + assert_eq!( + seen.url, + "https://us-east5-aiplatform.googleapis.com/v1/projects/test-project/locations/us-east5/publishers/anthropic/models/test-model-id:rawPredict" + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider.test.ts › google-vertex-anthropic-provider › should use multi-region URL for eu location + #[tokio::test] + async fn should_use_multi_region_url_for_eu_location() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("test-project".into()), + location: Some("eu".into()), + ..settings(&fetch) + }) + .unwrap(); + let model = provider.call("test-model-id"); + assert_eq!(model.provider(), "googleVertex.anthropic.messages"); + assert!(model.supported_urls().await.0.is_empty()); + + let seen = request(&provider, &fetch, "test-model-id").await; + assert_eq!( + seen.url, + "https://aiplatform.eu.rep.googleapis.com/v1/projects/test-project/locations/eu/publishers/anthropic/models/test-model-id:rawPredict" + ); + assert_eq!(provider_headers(&seen), BTreeMap::new()); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider.test.ts › google-vertex-anthropic-provider › should use multi-region URL for us location + #[tokio::test] + async fn should_use_multi_region_url_for_us_location() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("test-project".into()), + location: Some("us".into()), + ..settings(&fetch) + }) + .unwrap(); + let model = provider.call("test-model-id"); + assert_eq!(model.provider(), "googleVertex.anthropic.messages"); + assert!(model.supported_urls().await.0.is_empty()); + + let seen = request(&provider, &fetch, "test-model-id").await; + assert_eq!( + seen.url, + "https://aiplatform.us.rep.googleapis.com/v1/projects/test-project/locations/us/publishers/anthropic/models/test-model-id:rawPredict" + ); + assert_eq!(provider_headers(&seen), BTreeMap::new()); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider.test.ts › google-vertex-anthropic-provider › should support combining tools with structured outputs (inherited from Anthropic) + #[test] + fn should_support_combining_tools_with_structured_outputs_inherited_from_anthropic() { + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("test-project".into()), + location: Some("us-east5".into()), + ..Default::default() + }) + .unwrap(); + + let model = provider.call("claude-3-5-sonnet-v2@20241022"); + + assert_eq!(model.model_id(), "claude-3-5-sonnet-v2@20241022"); + assert_eq!(model.provider(), "googleVertex.anthropic.messages"); + } +} diff --git a/aimux-providers/src/vertex/tests/anthropic/anthropic_provider_node.rs b/aimux-providers/src/vertex/tests/anthropic/anthropic_provider_node.rs new file mode 100644 index 000000000..96cb2c4c7 --- /dev/null +++ b/aimux-providers/src/vertex/tests/anthropic/anthropic_provider_node.rs @@ -0,0 +1,387 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 +//! `src/anthropic/google-vertex-anthropic-provider-node.test.ts`. +//! +//! Upstream mocks the base provider and resolves the `headers` function the +//! Node wrapper hands it. aimux has one entry point, whose provider headers +//! are those of every request: each case makes a request through a mock +//! transport and reads the request's headers. Upstream's mocked +//! `createAuthTokenGenerator` (returning `mock-auth-token`) is replaced by the +//! `generate_auth_token` setting, except in the cases about the default +//! generator, which run Google auth with inline credentials over the mock. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_provider_utils::{HeaderMapOpt, Resolvable}; +use serde_json::json; + +use crate::test_support::mock_fetch::{Canned, MockFetch, RouteFetch, Seen}; +use crate::vertex::{ + GoogleAuthOptions, GoogleAuthScopes, VertexAnthropicProviderSettings, + create_google_vertex_anthropic, +}; + +/// A minimal Messages API response. +fn messages_response() -> Canned { + Canned::json(&json!({ + "id": "msg_1", "type": "message", "role": "assistant", "model": "claude", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", "stop_sequence": null, + "usage": {"input_tokens": 1, "output_tokens": 1} + })) +} + +fn prompt() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +/// Settings of a provider for `test-project` over `fetch`. +fn settings(fetch: &Arc) -> VertexAnthropicProviderSettings { + VertexAnthropicProviderSettings { + project: Some("test-project".into()), + location: Some("us-east5".into()), + fetch: Some(fetch.transport()), + ..Default::default() + } +} + +/// A token generator returning `token(n)` on its n-th call (from 1), and its +/// call counter. +fn counting_generator( + token: impl Fn(usize) -> Option + Send + Sync + 'static, +) -> (Resolvable>, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let token = Arc::new(token); + let generator = Resolvable::from_async_fn(move || { + let call = counter.fetch_add(1, Ordering::SeqCst) + 1; + let token = token.clone(); + async move { Ok(token(call)) } + }); + (generator, calls) +} + +/// Async caller headers, like upstream's `async () => ({ ... })`. +fn async_headers(headers: &[(&str, &str)]) -> Resolvable { + let headers: HeaderMapOpt = headers + .iter() + .map(|(name, value)| (name.to_string(), Some(value.to_string()))) + .collect(); + Resolvable::from_async_fn(move || { + let headers = headers.clone(); + async move { Ok(headers) } + }) +} + +/// The headers the provider adds to a request: everything but the +/// transport's own `content-type` and `user-agent` (upstream resolves the +/// `headers` function alone). +fn provider_headers(seen: &Seen) -> BTreeMap { + seen.headers + .iter() + .filter(|(name, _)| !matches!(name.as_str(), "content-type" | "user-agent")) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +fn headers(entries: &[(&str, &str)]) -> BTreeMap { + entries + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() +} + +/// Inline authorized-user credentials for the default Google auth. +fn authorized_user() -> serde_json::Value { + json!({ + "type": "authorized_user", + "client_id": "client-id", + "client_secret": "client-secret", + "refresh_token": "refresh-token" + }) +} + +mod google_vertex_anthropic_provider_node { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider-node.test.ts › google-vertex-anthropic-provider-node › default headers function should return auth token + #[tokio::test] + async fn default_headers_function_should_return_auth_token() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(Resolvable::Value(Some("mock-auth-token".into()))), + ..settings(&fetch) + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!( + provider_headers(&fetch.seen()[0]), + headers(&[("authorization", "Bearer mock-auth-token")]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider-node.test.ts › google-vertex-anthropic-provider-node › should use custom headers in addition to auth token when provided + #[tokio::test] + async fn should_use_custom_headers_in_addition_to_auth_token_when_provided() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + headers: Some(async_headers(&[("Custom-Header", "custom-value")])), + generate_auth_token: Some(Resolvable::Value(Some("mock-auth-token".into()))), + ..settings(&fetch) + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!( + provider_headers(&fetch.seen()[0]), + headers(&[ + ("authorization", "Bearer mock-auth-token"), + ("custom-header", "custom-value"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider-node.test.ts › google-vertex-anthropic-provider-node › passes googleAuthOptions to createAuthTokenGenerator + /// + /// No seam replaces the generator here, so the options reach the real + /// Google auth: it reads the given (missing) key file, and no request is + /// sent. + #[tokio::test] + async fn passes_google_auth_options_to_create_auth_token_generator() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + google_auth_options: Some(GoogleAuthOptions { + scopes: Some(GoogleAuthScopes::Array(vec![ + "https://www.googleapis.com/auth/cloud-platform".into(), + ])), + key_file: Some("path/to/key.json".into()), + ..Default::default() + }), + ..settings(&fetch) + }) + .unwrap(); + + let error = provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap_err(); + + assert!( + error + .to_string() + .contains("Could not read Google application credentials"), + "{error}" + ); + assert!(fetch.seen().is_empty()); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider-node.test.ts › google-vertex-anthropic-provider-node › uses custom generateAuthToken when provided and skips the default + #[tokio::test] + async fn uses_custom_generate_auth_token_when_provided_and_skips_the_default() { + let fetch = MockFetch::new(vec![messages_response()]); + let (custom_generate, calls) = counting_generator(|_| Some("custom-token".into())); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(custom_generate), + ..settings(&fetch) + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + let seen = fetch.seen(); + assert_eq!( + provider_headers(&seen[0]), + headers(&[("authorization", "Bearer custom-token")]) + ); + assert_eq!(calls.load(Ordering::SeqCst), 1); + // The default generator made no token request: the model request is + // the only one. + assert_eq!(seen.len(), 1); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider-node.test.ts › google-vertex-anthropic-provider-node › merges custom generateAuthToken with user-provided headers + #[tokio::test] + async fn merges_custom_generate_auth_token_with_user_provided_headers() { + let fetch = MockFetch::new(vec![messages_response()]); + let (custom_generate, _) = counting_generator(|_| Some("custom-token".into())); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(custom_generate), + headers: Some(async_headers(&[("Custom-Header", "custom-value")])), + ..settings(&fetch) + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!( + provider_headers(&fetch.seen()[0]), + headers(&[ + ("authorization", "Bearer custom-token"), + ("custom-header", "custom-value"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider-node.test.ts › google-vertex-anthropic-provider-node › invokes custom generateAuthToken on each headers resolution + #[tokio::test] + async fn invokes_custom_generate_auth_token_on_each_headers_resolution() { + let fetch = MockFetch::new(vec![messages_response(), messages_response()]); + let (custom_generate, calls) = counting_generator(|call| Some(format!("token-{call}"))); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(custom_generate), + ..settings(&fetch) + }) + .unwrap(); + let model = provider.call("test-model-id"); + + model.do_generate(&prompt()).await.unwrap(); + model.do_generate(&prompt()).await.unwrap(); + + let seen = fetch.seen(); + assert_eq!( + provider_headers(&seen[0]), + headers(&[("authorization", "Bearer token-1")]) + ); + assert_eq!( + provider_headers(&seen[1]), + headers(&[("authorization", "Bearer token-2")]) + ); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider-node.test.ts › google-vertex-anthropic-provider-node › propagates errors thrown from custom generateAuthToken + #[tokio::test] + async fn propagates_errors_thrown_from_custom_generate_auth_token() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(Resolvable::from_async_fn(|| async { + Err(AiMuxError::other("token mint failed")) + })), + ..settings(&fetch) + }) + .unwrap(); + + let error = provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap_err(); + + assert!(error.to_string().contains("token mint failed"), "{error}"); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider-node.test.ts › google-vertex-anthropic-provider-node › user-provided Authorization in headers overrides the generated token + #[tokio::test] + async fn user_provided_authorization_in_headers_overrides_the_generated_token() { + let fetch = MockFetch::new(vec![messages_response()]); + let (custom_generate, _) = counting_generator(|_| Some("custom-token".into())); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(custom_generate), + headers: Some(async_headers(&[("Authorization", "Bearer user-override")])), + ..settings(&fetch) + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!( + provider_headers(&fetch.seen()[0]), + headers(&[("authorization", "Bearer user-override")]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider-node.test.ts › google-vertex-anthropic-provider-node › accepts a generateAuthToken that returns null (matches default signature) + #[tokio::test] + async fn accepts_a_generate_auth_token_that_returns_null_matches_default_signature() { + let fetch = MockFetch::new(vec![messages_response()]); + let (custom_generate, _) = counting_generator(|_| None); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(custom_generate), + ..settings(&fetch) + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!( + provider_headers(&fetch.seen()[0]), + headers(&[("authorization", "Bearer null")]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/google-vertex-anthropic-provider-node.test.ts › google-vertex-anthropic-provider-node › creates the auth token generator once per provider instance + /// + /// The generator is aimux's Google auth of the provider. It is one + /// instance per provider when its token, cached by the instance, serves + /// both requests: one token exchange for two requests. + #[tokio::test] + async fn creates_the_auth_token_generator_once_per_provider_instance() { + let fetch = RouteFetch::new(|seen| { + if seen.url.starts_with("https://oauth2.googleapis.com/token") { + Canned::json(&json!({ + "access_token": "mock-auth-token", "expires_in": 3600, "token_type": "Bearer" + })) + } else { + messages_response() + } + }); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("test-project".into()), + location: Some("us-east5".into()), + fetch: Some(fetch.transport()), + google_auth_options: Some(GoogleAuthOptions { + credentials: Some(authorized_user()), + ..Default::default() + }), + ..Default::default() + }) + .unwrap(); + let model = provider.call("test-model-id"); + + model.do_generate(&prompt()).await.unwrap(); + model.do_generate(&prompt()).await.unwrap(); + + assert_eq!(fetch.count("POST", "/token"), 1); + let model_requests: Vec<_> = fetch + .seen() + .into_iter() + .filter(|seen| seen.url.contains(":rawPredict")) + .collect(); + assert_eq!(model_requests.len(), 2); + for seen in model_requests { + assert_eq!(seen.headers["authorization"], "Bearer mock-auth-token"); + } + } +} diff --git a/aimux-providers/src/vertex/tests/anthropic/edge/anthropic_provider_edge.rs b/aimux-providers/src/vertex/tests/anthropic/edge/anthropic_provider_edge.rs new file mode 100644 index 000000000..e08314b7a --- /dev/null +++ b/aimux-providers/src/vertex/tests/anthropic/edge/anthropic_provider_edge.rs @@ -0,0 +1,380 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 +//! `src/anthropic/edge/google-vertex-anthropic-provider-edge.test.ts`. +//! +//! aimux has one entry point and no separate Web-Crypto edge auth; these +//! cases exercise the same behaviour through it. Upstream mocks the base +//! provider and resolves the `headers` function the edge wrapper hands it; +//! here each case makes a request through a mock transport and reads the +//! request's headers. Upstream's mocked `generateAuthToken` (returning +//! `mock-auth-token`) is replaced by the `generate_auth_token` setting, +//! except in the cases about the default generator and `googleCredentials` +//! (aimux: `GoogleAuthOptions::credentials`), which run aimux's Google auth. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_provider_utils::{HeaderMapOpt, Resolvable}; +use serde_json::json; + +use crate::test_support::mock_fetch::{Canned, EnvVar, MockFetch, RouteFetch, Seen}; +use crate::vertex::{ + GoogleAuthOptions, VertexAnthropicProviderSettings, create_google_vertex_anthropic, +}; + +/// A minimal Messages API response. +fn messages_response() -> Canned { + Canned::json(&json!({ + "id": "msg_1", "type": "message", "role": "assistant", "model": "claude", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", "stop_sequence": null, + "usage": {"input_tokens": 1, "output_tokens": 1} + })) +} + +fn prompt() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +/// Settings of a provider for `test-project` over `fetch`. +fn settings(fetch: &Arc) -> VertexAnthropicProviderSettings { + VertexAnthropicProviderSettings { + project: Some("test-project".into()), + location: Some("us-east5".into()), + fetch: Some(fetch.transport()), + ..Default::default() + } +} + +/// A token generator returning `token(n)` on its n-th call (from 1), and its +/// call counter. +fn counting_generator( + token: impl Fn(usize) -> Option + Send + Sync + 'static, +) -> (Resolvable>, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let token = Arc::new(token); + let generator = Resolvable::from_async_fn(move || { + let call = counter.fetch_add(1, Ordering::SeqCst) + 1; + let token = token.clone(); + async move { Ok(token(call)) } + }); + (generator, calls) +} + +/// Async caller headers, like upstream's `async () => ({ ... })`. +fn async_headers(headers: &[(&str, &str)]) -> Resolvable { + let headers: HeaderMapOpt = headers + .iter() + .map(|(name, value)| (name.to_string(), Some(value.to_string()))) + .collect(); + Resolvable::from_async_fn(move || { + let headers = headers.clone(); + async move { Ok(headers) } + }) +} + +/// The headers the provider adds to a request: everything but the +/// transport's own `content-type` and `user-agent` (upstream resolves the +/// `headers` function alone). +fn provider_headers(seen: &Seen) -> BTreeMap { + seen.headers + .iter() + .filter(|(name, _)| !matches!(name.as_str(), "content-type" | "user-agent")) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +fn headers(entries: &[(&str, &str)]) -> BTreeMap { + entries + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() +} + +/// No credential file or gcloud configuration and the default metadata +/// host: the default Google auth asks the metadata server for a token. +fn metadata_only_env() -> Vec { + vec![ + EnvVar::set("GOOGLE_APPLICATION_CREDENTIALS", None), + EnvVar::set("google_application_credentials", None), + EnvVar::set( + "CLOUDSDK_CONFIG", + Some("/nonexistent/aimux-vertex-test-gcloud"), + ), + EnvVar::set("GCE_METADATA_IP", None), + EnvVar::set("GCE_METADATA_HOST", None), + ] +} + +/// The metadata server's token response. +fn metadata_token(token: &str) -> Canned { + let mut canned = Canned::json(&json!({ + "access_token": token, "expires_in": 3600, "token_type": "Bearer" + })); + canned + .headers + .push(("Metadata-Flavor".into(), "Google".into())); + canned +} +mod google_vertex_anthropic_provider_edge { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/anthropic/edge/google-vertex-anthropic-provider-edge.test.ts › google-vertex-anthropic-provider-edge › default headers function should return auth token + #[tokio::test] + async fn default_headers_function_should_return_auth_token() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(Resolvable::Value(Some("mock-auth-token".into()))), + ..settings(&fetch) + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!( + provider_headers(&fetch.seen()[0]), + headers(&[("authorization", "Bearer mock-auth-token")]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/edge/google-vertex-anthropic-provider-edge.test.ts › google-vertex-anthropic-provider-edge › should use custom headers in addition to auth token when provided + #[tokio::test] + async fn should_use_custom_headers_in_addition_to_auth_token_when_provided() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + headers: Some(async_headers(&[("Custom-Header", "custom-value")])), + generate_auth_token: Some(Resolvable::Value(Some("mock-auth-token".into()))), + ..settings(&fetch) + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!( + provider_headers(&fetch.seen()[0]), + headers(&[ + ("authorization", "Bearer mock-auth-token"), + ("custom-header", "custom-value"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/edge/google-vertex-anthropic-provider-edge.test.ts › google-vertex-anthropic-provider-edge › should use edge auth token generator + /// + /// Without `generate_auth_token` the default generator makes the token: + /// here aimux's Google auth, from the metadata server. + #[tokio::test] + #[serial_test::serial] + async fn should_use_edge_auth_token_generator() { + let _env = metadata_only_env(); + let fetch = RouteFetch::new(|seen| { + if seen + .url + .contains("/computeMetadata/v1/instance/service-accounts/default/token") + { + metadata_token("mock-auth-token") + } else { + messages_response() + } + }); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + project: Some("test-project".into()), + location: Some("us-east5".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!( + fetch.count( + "GET", + "/computeMetadata/v1/instance/service-accounts/default/token" + ), + 1 + ); + let seen = fetch.seen(); + let model_request = seen + .iter() + .find(|seen| seen.url.contains(":rawPredict")) + .unwrap(); + assert_eq!( + provider_headers(model_request), + headers(&[("authorization", "Bearer mock-auth-token")]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/edge/google-vertex-anthropic-provider-edge.test.ts › google-vertex-anthropic-provider-edge › passes googleCredentials to generateAuthToken + /// + /// The credentials reach aimux's real Google auth (upstream mocks the + /// generator), which signs with the given private key: `test-key` is not + /// a key, so the request fails on it before anything is sent. + #[tokio::test] + async fn passes_google_credentials_to_generate_auth_token() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + google_auth_options: Some(GoogleAuthOptions { + credentials: Some(json!({ + "client_email": "test@example.com", + "private_key": "test-key" + })), + ..Default::default() + }), + ..settings(&fetch) + }) + .unwrap(); + + let error = provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap_err(); + + assert!(error.to_string().contains("private key"), "{error}"); + assert!(fetch.seen().is_empty()); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/edge/google-vertex-anthropic-provider-edge.test.ts › google-vertex-anthropic-provider-edge › uses custom generateAuthToken when provided and skips the default + #[tokio::test] + async fn uses_custom_generate_auth_token_when_provided_and_skips_the_default() { + let fetch = MockFetch::new(vec![messages_response()]); + let (custom_generate, calls) = counting_generator(|_| Some("custom-token".into())); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(custom_generate), + ..settings(&fetch) + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + let seen = fetch.seen(); + assert_eq!( + provider_headers(&seen[0]), + headers(&[("authorization", "Bearer custom-token")]) + ); + assert_eq!(calls.load(Ordering::SeqCst), 1); + // The default generator made no token request: the model request is + // the only one. + assert_eq!(seen.len(), 1); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/edge/google-vertex-anthropic-provider-edge.test.ts › google-vertex-anthropic-provider-edge › merges custom generateAuthToken with user-provided headers + #[tokio::test] + async fn merges_custom_generate_auth_token_with_user_provided_headers() { + let fetch = MockFetch::new(vec![messages_response()]); + let (custom_generate, _) = counting_generator(|_| Some("custom-token".into())); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(custom_generate), + headers: Some(async_headers(&[("Custom-Header", "custom-value")])), + ..settings(&fetch) + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!( + provider_headers(&fetch.seen()[0]), + headers(&[ + ("authorization", "Bearer custom-token"), + ("custom-header", "custom-value"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/edge/google-vertex-anthropic-provider-edge.test.ts › google-vertex-anthropic-provider-edge › invokes custom generateAuthToken on each headers resolution + #[tokio::test] + async fn invokes_custom_generate_auth_token_on_each_headers_resolution() { + let fetch = MockFetch::new(vec![messages_response(), messages_response()]); + let (custom_generate, calls) = counting_generator(|call| Some(format!("token-{call}"))); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(custom_generate), + ..settings(&fetch) + }) + .unwrap(); + let model = provider.call("test-model-id"); + + model.do_generate(&prompt()).await.unwrap(); + model.do_generate(&prompt()).await.unwrap(); + + let seen = fetch.seen(); + assert_eq!( + provider_headers(&seen[0]), + headers(&[("authorization", "Bearer token-1")]) + ); + assert_eq!( + provider_headers(&seen[1]), + headers(&[("authorization", "Bearer token-2")]) + ); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/edge/google-vertex-anthropic-provider-edge.test.ts › google-vertex-anthropic-provider-edge › propagates errors thrown from custom generateAuthToken + #[tokio::test] + async fn propagates_errors_thrown_from_custom_generate_auth_token() { + let fetch = MockFetch::new(vec![messages_response()]); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(Resolvable::from_async_fn(|| async { + Err(AiMuxError::other("token mint failed")) + })), + ..settings(&fetch) + }) + .unwrap(); + + let error = provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap_err(); + + assert!(error.to_string().contains("token mint failed"), "{error}"); + } + + /// upstream: @ai-sdk/google-vertex src/anthropic/edge/google-vertex-anthropic-provider-edge.test.ts › google-vertex-anthropic-provider-edge › user-provided Authorization in headers overrides the generated token + #[tokio::test] + async fn user_provided_authorization_in_headers_overrides_the_generated_token() { + let fetch = MockFetch::new(vec![messages_response()]); + let (custom_generate, _) = counting_generator(|_| Some("custom-token".into())); + let provider = create_google_vertex_anthropic(VertexAnthropicProviderSettings { + generate_auth_token: Some(custom_generate), + headers: Some(async_headers(&[("Authorization", "Bearer user-override")])), + ..settings(&fetch) + }) + .unwrap(); + + provider + .call("test-model-id") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!( + provider_headers(&fetch.seen()[0]), + headers(&[("authorization", "Bearer user-override")]) + ); + } +} diff --git a/aimux-providers/src/vertex/tests/anthropic/edge/mod.rs b/aimux-providers/src/vertex/tests/anthropic/edge/mod.rs new file mode 100644 index 000000000..c3b860845 --- /dev/null +++ b/aimux-providers/src/vertex/tests/anthropic/edge/mod.rs @@ -0,0 +1 @@ +mod anthropic_provider_edge; diff --git a/aimux-providers/src/vertex/tests/anthropic/mod.rs b/aimux-providers/src/vertex/tests/anthropic/mod.rs new file mode 100644 index 000000000..cdc858a8b --- /dev/null +++ b/aimux-providers/src/vertex/tests/anthropic/mod.rs @@ -0,0 +1,3 @@ +mod anthropic_provider; +mod anthropic_provider_node; +mod edge; diff --git a/aimux-providers/src/vertex/tests/auth_google_auth_library.rs b/aimux-providers/src/vertex/tests/auth_google_auth_library.rs new file mode 100644 index 000000000..a2f23c6be --- /dev/null +++ b/aimux-providers/src/vertex/tests/auth_google_auth_library.rs @@ -0,0 +1,196 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 +//! `src/google-vertex-auth-google-auth-library.test.ts`. +//! +//! Upstream's `createAuthTokenGenerator` wraps `google-auth-library`, which the +//! test replaces with a mock `GoogleAuth` whose client returns a fixed token. +//! aimux's counterpart is the provider's Google credential resolver +//! (`GoogleAuth` in `src/vertex/auth.rs`); it cannot be replaced, so the tests +//! give it real service-account credentials (a test RSA key written to a +//! key file) and answer its token request from a mock transport. The scopes +//! and key file upstream passes to the `GoogleAuth` constructor are read +//! from the signed assertion the resolver sends; the mocked client's +//! `{ token: null }` is the `auth_client` option, aimux's replacement for the +//! OAuth client. Tests that set environment variables are serial. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use base64::Engine; +use serde_json::{Value, json}; +use serial_test::serial; + +use crate::test_support::mock_fetch::{Canned, EnvVar, MockFetch, Seen}; +use crate::vertex::GoogleAuthOptions; +use crate::vertex::auth::GoogleAuth; +use aimux_provider_utils::Resolvable; + +const TEST_RSA_KEY: &str = include_str!("fixtures/aimux-test-rsa-key.pem"); +const CLOUD_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; + +/// A directory of key files for one test, removed on drop. +struct KeyDirectory(PathBuf); + +impl KeyDirectory { + fn new(test: &str) -> Self { + let path = std::env::temp_dir().join(format!( + "aimux-vertex-auth-library-{}-{test}", + std::process::id() + )); + std::fs::create_dir_all(&path).unwrap(); + Self(path) + } + + /// A service-account key file named `name` for `client_email`. + fn key_file(&self, name: &str, client_email: &str) -> PathBuf { + let path = self.0.join(name); + let credentials = json!({ + "type": "service_account", + "client_email": client_email, + "private_key": TEST_RSA_KEY, + }); + std::fs::write(&path, credentials.to_string()).unwrap(); + path + } +} + +impl Drop for KeyDirectory { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } +} + +/// `getAccessToken.mockResolvedValue({ token })`: the token endpoint's answer. +fn token_response(token: &str) -> Canned { + Canned::json(&json!({"access_token": token, "expires_in": 3600})) +} + +fn generator(options: GoogleAuthOptions, fetch: &Arc) -> GoogleAuth { + GoogleAuth::new(options, Some(fetch.transport())) +} + +fn key_file_options(path: &Path) -> GoogleAuthOptions { + GoogleAuthOptions { + key_file: Some(path.to_path_buf()), + ..Default::default() + } +} + +/// The claims of the JWT assertion a token request carries. +fn assertion_claims(seen: &Seen) -> Value { + let assertion = url::form_urlencoded::parse(&seen.body) + .find(|(name, _)| name == "assertion") + .map(|(_, value)| value.into_owned()) + .expect("a JWT-bearer token request"); + let claims = assertion.split('.').nth(1).expect("a JWT"); + serde_json::from_slice( + &base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(claims) + .unwrap(), + ) + .unwrap() +} + +mod create_auth_token_generator { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-auth-google-auth-library.test.ts › createAuthTokenGenerator › should generate a valid auth token + #[tokio::test] + #[serial] + async fn should_generate_a_valid_auth_token() { + // `createAuthTokenGenerator()` with no options: the application + // default credentials, here from GOOGLE_APPLICATION_CREDENTIALS. + let keys = KeyDirectory::new("valid-token"); + let path = keys.key_file("adc.json", "adc@test.iam.gserviceaccount.com"); + let _credentials = EnvVar::set( + "GOOGLE_APPLICATION_CREDENTIALS", + Some(path.to_str().unwrap()), + ); + let fetch = MockFetch::new(vec![token_response("mocked-token")]); + let generate_auth_token = generator(GoogleAuthOptions::default(), &fetch); + + let token = generate_auth_token.access_token().await.unwrap(); + + assert_eq!(token, "mocked-token"); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-auth-google-auth-library.test.ts › createAuthTokenGenerator › should return null if no token is received + #[tokio::test] + async fn should_return_null_if_no_token_is_received() { + // The client answers `{ token: null }`. aimux's generator returns the + // token text that goes into `Authorization: Bearer `, where a + // null token is the text `null`, as in upstream's template string. + let fetch = MockFetch::new(vec![]); + let generate_auth_token = generator( + GoogleAuthOptions { + auth_client: Some(Resolvable::Value(None)), + ..Default::default() + }, + &fetch, + ); + + let token = generate_auth_token.access_token().await.unwrap(); + + assert_eq!(token, "null"); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-auth-google-auth-library.test.ts › createAuthTokenGenerator › should create a new auth instance with provided options + #[tokio::test] + async fn should_create_a_new_auth_instance_with_provided_options() { + let keys = KeyDirectory::new("provided-options"); + let path = keys.key_file("test-key.json", "test-key@test.iam.gserviceaccount.com"); + let fetch = MockFetch::new(vec![token_response("mocked-token")]); + + generator(key_file_options(&path), &fetch) + .access_token() + .await + .unwrap(); + + // `{ scopes: [cloud-platform], keyFile: 'test-key.json' }`: the + // credentials of the key file, with the Cloud Platform scope. + let claims = assertion_claims(&fetch.seen()[0]); + assert_eq!(claims["iss"], "test-key@test.iam.gserviceaccount.com"); + assert_eq!(claims["scope"], CLOUD_SCOPE); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-auth-google-auth-library.test.ts › createAuthTokenGenerator › should create only one GoogleAuth instance for repeated calls + #[tokio::test] + async fn should_create_only_one_google_auth_instance_for_repeated_calls() { + let keys = KeyDirectory::new("repeated-calls"); + let path = keys.key_file("key.json", "repeat@test.iam.gserviceaccount.com"); + let fetch = MockFetch::new(vec![token_response("mocked-token")]); + let generate_auth_token = generator(key_file_options(&path), &fetch); + + generate_auth_token.access_token().await.unwrap(); + generate_auth_token.access_token().await.unwrap(); + + // One auth instance: its client is set up (and its token fetched) + // once. + assert_eq!(fetch.seen().len(), 1); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-auth-google-auth-library.test.ts › createAuthTokenGenerator › should create independent generators for separate option sets + #[tokio::test] + async fn should_create_independent_generators_for_separate_option_sets() { + let keys = KeyDirectory::new("separate-options"); + let first = keys.key_file("first-key.json", "first@test.iam.gserviceaccount.com"); + let second = keys.key_file("second-key.json", "second@test.iam.gserviceaccount.com"); + let fetch = MockFetch::new(vec![ + token_response("first-token"), + token_response("second-token"), + ]); + let generate_first_auth_token = generator(key_file_options(&first), &fetch); + let generate_second_auth_token = generator(key_file_options(&second), &fetch); + + generate_first_auth_token.access_token().await.unwrap(); + generate_second_auth_token.access_token().await.unwrap(); + + let seen = fetch.seen(); + assert_eq!(seen.len(), 2); + let first_claims = assertion_claims(&seen[0]); + assert_eq!(first_claims["iss"], "first@test.iam.gserviceaccount.com"); + assert_eq!(first_claims["scope"], CLOUD_SCOPE); + let second_claims = assertion_claims(&seen[1]); + assert_eq!(second_claims["iss"], "second@test.iam.gserviceaccount.com"); + assert_eq!(second_claims["scope"], CLOUD_SCOPE); + } +} diff --git a/aimux-providers/src/vertex/tests/cloud_tts_speech_model.rs b/aimux-providers/src/vertex/tests/cloud_tts_speech_model.rs new file mode 100644 index 000000000..241652dec --- /dev/null +++ b/aimux-providers/src/vertex/tests/cloud_tts_speech_model.rs @@ -0,0 +1,4 @@ +//! `src/google-vertex-cloud-tts-speech-model.test.ts` is not ported: aimux has +//! no Vertex Cloud Text-to-Speech (Chirp) speech model. +//! `VertexProvider::speech` rejects `chirp*` model ids with +//! `AiMuxError::UnsupportedFunctionality`. diff --git a/aimux-providers/src/vertex/tests/edge/auth_edge.rs b/aimux-providers/src/vertex/tests/edge/auth_edge.rs new file mode 100644 index 000000000..c32cd1a4a --- /dev/null +++ b/aimux-providers/src/vertex/tests/edge/auth_edge.rs @@ -0,0 +1,186 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 `src/edge/google-vertex-auth-edge.test.ts`. +//! +//! The edge `generateAuthToken(credentials)` signs a service-account JWT with +//! Web Crypto and exchanges it at the OAuth token endpoint. aimux has one +//! entry point: the provider's Google credential resolver (`GoogleAuth` in +//! `src/vertex/auth.rs`) does the same exchange for service-account +//! credentials. Upstream's `GoogleCredentials` (`clientEmail`, `privateKey`, +//! `privateKeyId`) is the service-account JSON (`client_email`, +//! `private_key`, `private_key_id`) here. Upstream mocks Web Crypto, so its +//! `'mock-private-key'` is never parsed; aimux signs for real, so the harness +//! uses a test RSA key. `global.fetch` is the resolver's transport, and the +//! stubbed `atob` (which returns the same payload for any input) is +//! [`atob`] below. + +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::{Value, json}; + +use crate::test_support::mock_fetch::{Canned, MockFetch}; +use crate::vertex::GoogleAuthOptions; +use crate::vertex::auth::GoogleAuth; +use aimux_provider_utils::{Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse}; + +const TEST_RSA_KEY: &str = include_str!("../fixtures/aimux-test-rsa-key.pem"); + +/// Upstream's `GoogleCredentials`. +#[derive(Clone)] +struct GoogleCredentials { + client_email: &'static str, + private_key: &'static str, + private_key_id: Option<&'static str>, +} + +/// `mockCredentials`, with the test RSA key. +fn mock_credentials() -> GoogleCredentials { + GoogleCredentials { + client_email: "test@test.iam.gserviceaccount.com", + private_key: TEST_RSA_KEY, + private_key_id: Some("test-key-id"), + } +} + +/// The service-account JSON of `credentials`. +fn credential_json(credentials: &GoogleCredentials) -> Value { + let mut value = json!({ + "type": "service_account", + "client_email": credentials.client_email, + "private_key": credentials.private_key, + }); + if let Some(id) = credentials.private_key_id { + value["private_key_id"] = json!(id); + } + value +} + +/// `setupAtobStub(credentials)`: the stubbed `atob` returns this payload for +/// any input. +fn atob(credentials: &GoogleCredentials, _input: &str) -> String { + let mut payload = json!({ + "alg": "RS256", + "typ": "JWT", + "iss": credentials.client_email, + "scope": "https://www.googleapis.com/auth/cloud-platform", + "aud": "https://oauth2.googleapis.com/token", + "iat": 1_616_161_616, + "exp": 1_616_165_216, + }); + if let Some(id) = credentials.private_key_id { + payload["kid"] = json!(id); + } + payload.to_string() +} + +/// `generateAuthToken(credentials)` through `fetch`. +async fn generate_auth_token( + credentials: &GoogleCredentials, + fetch: FetchFunction, +) -> Result { + GoogleAuth::new( + GoogleAuthOptions { + credentials: Some(credential_json(credentials)), + ..Default::default() + }, + Some(fetch), + ) + .access_token() + .await +} + +/// `{ ok: true, json: () => ({ access_token }) }`. +fn token_response(token: &str) -> Canned { + Canned::json(&json!({"access_token": token})) +} + +/// `global.fetch = vi.fn().mockRejectedValue(new Error(message))`. +struct RejectingFetch(&'static str); + +#[async_trait] +impl Fetch for RejectingFetch { + async fn fetch(&self, _request: FetchRequest) -> Result { + Err(FetchError::Connect(self.0.into())) + } +} + +mod google_vertex_edge_auth { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/edge/google-vertex-auth-edge.test.ts › Google Vertex Edge Auth › should generate a valid JWT token + #[tokio::test] + async fn should_generate_a_valid_jwt_token() { + let credentials = mock_credentials(); + let fetch = MockFetch::new(vec![token_response( + "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCIsImtpZCI6InRlc3Qta2V5LWlkIn0.eyJpc3MiOiJ0ZXN0QHRlc3QuaWFtLmdzZXJ2aWNlYWNjb3VudC5jb20iLCJzY29wZSI6Imh0dHBzOi8vd3d3Lmdvb2dsZWFwaXMuY29tL2F1dGgvY2xvdWQtcGxhdGZvcm0iLCJhdWQiOiJodHRwczovL29hdXRoMi5nb29nbGVhcGlzLmNvbS90b2tlbiIsImlhdCI6MTYxNjE2MTYxNiwiZXhwIjoxNjE2MTY1MjE2fQ.signature", + )]); + + let token = generate_auth_token(&credentials, fetch.transport()) + .await + .unwrap(); + + // JWT structure validation + let parts: Vec<&str> = token.split('.').collect(); + assert_eq!(parts.len(), 3); + + // Header validation + let header: Value = serde_json::from_str(&atob(&credentials, parts[0])).unwrap(); + assert_eq!(header["alg"], "RS256"); + assert_eq!(header["typ"], "JWT"); + assert_eq!(header["kid"], "test-key-id"); + assert_eq!(header["iss"], credentials.client_email); + + // Payload validation + let payload: Value = serde_json::from_str(&atob(&credentials, parts[1])).unwrap(); + assert_eq!(payload["iss"], credentials.client_email); + assert_eq!( + payload["scope"], + "https://www.googleapis.com/auth/cloud-platform" + ); + assert_eq!(payload["aud"], "https://oauth2.googleapis.com/token"); + assert!(payload.get("iat").is_some()); + assert!(payload.get("exp").is_some()); + + // Verify exp is ~1 hour after iat + assert_eq!( + payload["exp"].as_i64().unwrap() - payload["iat"].as_i64().unwrap(), + 3600 + ); + } + + /// upstream: @ai-sdk/google-vertex src/edge/google-vertex-auth-edge.test.ts › Google Vertex Edge Auth › should throw error on fetch failure + #[tokio::test] + async fn should_throw_error_on_fetch_failure() { + let error = generate_auth_token( + &mock_credentials(), + Arc::new(RejectingFetch("Network error")), + ) + .await + .unwrap_err(); + + assert!(error.to_string().contains("Network error"), "{error}"); + } + + /// upstream: @ai-sdk/google-vertex src/edge/google-vertex-auth-edge.test.ts › Google Vertex Edge Auth › should work without privateKeyId + #[tokio::test] + async fn should_work_without_private_key_id() { + let credentials_without_key_id = GoogleCredentials { + private_key_id: None, + ..mock_credentials() + }; + let fetch = MockFetch::new(vec![token_response("mock.jwt.token")]); + + let token = generate_auth_token(&credentials_without_key_id, fetch.transport()) + .await + .unwrap(); + assert!(!token.is_empty()); + + // Verify the JWT structure + let parts: Vec<&str> = token.split('.').collect(); + assert_eq!(parts.len(), 3); + + // Verify header doesn't include kid when privateKeyId is not provided + let header: Value = + serde_json::from_str(&atob(&credentials_without_key_id, parts[0])).unwrap(); + assert!(header.get("kid").is_none()); + } +} diff --git a/aimux-providers/src/vertex/tests/edge/mod.rs b/aimux-providers/src/vertex/tests/edge/mod.rs new file mode 100644 index 000000000..779c4845c --- /dev/null +++ b/aimux-providers/src/vertex/tests/edge/mod.rs @@ -0,0 +1,2 @@ +mod auth_edge; +mod provider_edge; diff --git a/aimux-providers/src/vertex/tests/edge/provider_edge.rs b/aimux-providers/src/vertex/tests/edge/provider_edge.rs new file mode 100644 index 000000000..5070e54b7 --- /dev/null +++ b/aimux-providers/src/vertex/tests/edge/provider_edge.rs @@ -0,0 +1,183 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 `src/edge/google-vertex-provider-edge.test.ts`. +//! +//! The edge provider wraps the base provider with the Web-Crypto token +//! generator (`generateAuthToken`). aimux has one provider and one entry +//! point; the same behaviour goes through it: an `access_token` producer +//! stands in for the mocked `generateAuthToken`, `google_auth_options.credentials` +//! (a service-account JSON) carries upstream's `googleCredentials`, and the +//! tests read the provider headers its resolver produces for a request +//! (upstream's `passedOptions.headers`). Every test clears +//! `GOOGLE_VERTEX_API_KEY` (upstream's `beforeEach`) and +//! `GOOGLE_VERTEX_ACCESS_TOKEN` (read before the generator), so all of them +//! are serial. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use futures::FutureExt; +use serde_json::json; +use serial_test::serial; + +use crate::test_support::mock_fetch::EnvVar; +use crate::vertex::{ + GoogleAuthOptions, VertexProvider, VertexProviderSettings, create_google_vertex, +}; +use aimux_provider_utils::{HeaderMapOpt, Resolvable}; + +fn clean_environment() -> Vec { + ["GOOGLE_VERTEX_API_KEY", "GOOGLE_VERTEX_ACCESS_TOKEN"] + .into_iter() + .map(|name| EnvVar::set(name, None)) + .collect() +} + +/// `generateAuthToken: vi.fn().mockResolvedValue('mock-auth-token')`, +/// counting its calls. +fn mock_generate_auth_token(calls: &Arc) -> Resolvable { + let calls = calls.clone(); + Resolvable::AsyncFn(Arc::new(move || { + calls.fetch_add(1, Ordering::SeqCst); + async { Ok("mock-auth-token".to_string()) }.boxed() + })) +} + +fn create(settings: VertexProviderSettings) -> VertexProvider { + create_google_vertex(settings).expect("valid settings") +} + +/// `await resolve(passedOptions.headers)`: the provider headers of one +/// request, in the mode the provider was created in. +async fn provider_headers(provider: &VertexProvider) -> HeaderMapOpt { + let key = provider.resolver.express_key().await.expect("api key"); + provider + .resolver + .headers(key.as_deref()) + .await + .expect("headers") +} + +/// Expected provider headers. aimux combines header maps with lower-cased +/// names (HTTP names are case-insensitive), so the names are lower-cased. +fn headers(entries: &[(&str, &str)]) -> HeaderMapOpt { + entries + .iter() + .map(|(name, value)| (name.to_ascii_lowercase(), Some((*value).to_string()))) + .collect() +} + +mod google_vertex_provider_edge { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/edge/google-vertex-provider-edge.test.ts › google-vertex-provider-edge › default headers function should return auth token + #[tokio::test] + #[serial] + async fn default_headers_function_should_return_auth_token() { + let _env = clean_environment(); + let calls = Arc::new(AtomicUsize::new(0)); + let provider = create(VertexProviderSettings { + project: Some("test-project".into()), + access_token: Some(mock_generate_auth_token(&calls)), + ..Default::default() + }); + + assert_eq!( + provider_headers(&provider).await, + headers(&[("Authorization", "Bearer mock-auth-token")]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/edge/google-vertex-provider-edge.test.ts › google-vertex-provider-edge › should use custom headers in addition to auth token when provided + #[tokio::test] + #[serial] + async fn should_use_custom_headers_in_addition_to_auth_token_when_provided() { + let _env = clean_environment(); + let calls = Arc::new(AtomicUsize::new(0)); + let provider = create(VertexProviderSettings { + project: Some("test-project".into()), + headers: Some(Resolvable::AsyncFn(Arc::new(|| { + async { Ok(headers(&[("Custom-Header", "custom-value")])) }.boxed() + }))), + access_token: Some(mock_generate_auth_token(&calls)), + ..Default::default() + }); + + assert_eq!( + provider_headers(&provider).await, + headers(&[ + ("Authorization", "Bearer mock-auth-token"), + ("Custom-Header", "custom-value"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/edge/google-vertex-provider-edge.test.ts › google-vertex-provider-edge › should use edge auth token generator + #[tokio::test] + #[serial] + async fn should_use_edge_auth_token_generator() { + let _env = clean_environment(); + let calls = Arc::new(AtomicUsize::new(0)); + let provider = create(VertexProviderSettings { + project: Some("test-project".into()), + access_token: Some(mock_generate_auth_token(&calls)), + ..Default::default() + }); + + provider_headers(&provider).await; + assert!(calls.load(Ordering::SeqCst) > 0); + } + + /// upstream: @ai-sdk/google-vertex src/edge/google-vertex-provider-edge.test.ts › google-vertex-provider-edge › passes googleCredentials to generateAuthToken + #[tokio::test] + #[serial] + async fn passes_google_credentials_to_generate_auth_token() { + let _env = clean_environment(); + // `googleCredentials: { clientEmail, privateKey }` is a + // service-account credential JSON in aimux. The credential resolver + // is built with the provider, so no headers call is needed (and + // 'test-key' is no RSA key). + let credentials = json!({ + "client_email": "test@example.com", + "private_key": "test-key", + }); + let provider = create(VertexProviderSettings { + project: Some("test-project".into()), + google_auth_options: Some(GoogleAuthOptions { + credentials: Some(credentials.clone()), + ..Default::default() + }), + ..Default::default() + }); + + assert_eq!( + provider.resolver.auth.options.credentials, + Some(credentials) + ); + } + + /// upstream: @ai-sdk/google-vertex src/edge/google-vertex-provider-edge.test.ts › google-vertex-provider-edge › should pass options through to base provider when apiKey is provided + #[tokio::test] + #[serial] + async fn should_pass_options_through_to_base_provider_when_api_key_is_provided() { + let _env = clean_environment(); + let calls = Arc::new(AtomicUsize::new(0)); + let provider = create(VertexProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".into())), + access_token: Some(mock_generate_auth_token(&calls)), + ..Default::default() + }); + + assert_eq!( + provider.resolver.express_key().await.unwrap().as_deref(), + Some("test-api-key") + ); + // No auth-token headers: only the Express-mode key. + let headers = provider_headers(&provider).await; + assert!( + !headers + .keys() + .any(|name| name.eq_ignore_ascii_case("authorization")), + "{headers:?}" + ); + assert_eq!(calls.load(Ordering::SeqCst), 0); + } +} diff --git a/aimux-providers/src/vertex/tests/embedding_model.rs b/aimux-providers/src/vertex/tests/embedding_model.rs new file mode 100644 index 000000000..744c3aa95 --- /dev/null +++ b/aimux-providers/src/vertex/tests/embedding_model.rs @@ -0,0 +1,540 @@ +use std::collections::BTreeMap; +use std::sync::Arc; + +use serde_json::{Value, json}; + +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::error::AiMuxError; +use aimux_core::shared::SharedProviderOptions; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable}; + +use crate::shared::{Endpoint, EndpointConfig, ResolveEndpoint}; +use crate::test_support::mock_fetch::{Canned, MockFetch, RouteFetch}; +use crate::vertex::{VertexEmbeddingModel, VertexProviderSettings, create_google_vertex}; + +const DEFAULT_URL: &str = "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/textembedding-gecko@001:predict"; +const GEMINI_EMBEDDING_2_URL: &str = "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-embedding-2:embedContent"; +const CUSTOM_URL: &str = "https://custom-endpoint.com/models/textembedding-gecko@001:predict"; +const BASE_URL: &str = "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google"; +const MOCK_MODEL_ID: &str = "textembedding-gecko@001"; +const TEST_VALUES: [&str; 2] = ["test text one", "test text two"]; + +/// The pinned user-agent part of the package (upstream mocks `0.0.0-test`). +const USER_AGENT_PART: &str = "ai-sdk-google-vertex/5.0.100"; + +fn fixture() -> Value { + serde_json::from_str(include_str!("fixtures/google-vertex-embedding.json")).unwrap() +} + +/// A base URL and headers fixed for every request (upstream's `mockConfig`). +struct StaticEndpoint { + base_url: String, + headers: HeaderMapOpt, +} + +#[async_trait::async_trait] +impl ResolveEndpoint for StaticEndpoint { + async fn resolve(&self) -> Result { + Ok(Endpoint { + base_url: self.base_url.clone(), + headers: self.headers.clone(), + }) + } +} + +/// `new GoogleVertexEmbeddingModel(modelId, { provider: 'google-vertex', baseURL, headers: () => ({}), fetch })`. +fn model_with(model_id: &str, base_url: &str, fetch: FetchFunction) -> VertexEmbeddingModel { + VertexEmbeddingModel::from_config( + model_id.to_string(), + EndpointConfig::new( + "google-vertex".to_string(), + StaticEndpoint { + base_url: base_url.to_string(), + headers: HeaderMapOpt::new(), + }, + Some(fetch), + ), + ) +} + +/// A JSON response as the AI SDK test server sends it (compact body, +/// `content-type` and `content-length`). +fn json_response(body: &Value, headers: &[(&str, &str)]) -> Canned { + let body = serde_json::to_vec(body).unwrap(); + let mut all = vec![ + ("content-type".to_string(), "application/json".to_string()), + ("content-length".to_string(), body.len().to_string()), + ]; + all.extend(headers.iter().map(|(k, v)| (k.to_string(), v.to_string()))); + Canned { + status: 200, + headers: all, + body, + } +} + +/// The AI SDK test server: each URL answers its response, anything else 404. +fn test_server(routes: Vec<(&'static str, Canned)>) -> Arc { + RouteFetch::new(move |seen| { + routes + .iter() + .find(|(url, _)| *url == seen.url) + .map(|(_, canned)| canned.clone()) + .unwrap_or(Canned { + status: 404, + headers: Vec::new(), + body: b"Not found".to_vec(), + }) + }) +} + +/// `prepareJsonFixtureResponse('google-vertex-embedding', headers)`. +fn fixture_server(headers: &[(&str, &str)]) -> Arc { + test_server(vec![(DEFAULT_URL, json_response(&fixture(), headers))]) +} + +fn mock_provider_options() -> Value { + json!({ + "outputDimensionality": 768, + "taskType": "SEMANTIC_SIMILARITY", + "title": "test title", + "autoTruncate": false, + }) +} + +fn options(values: &[&str], provider_options: Value) -> EmbeddingCallOptions { + let mut options = EmbeddingCallOptions::new(values[0]); + options.values = values.iter().map(ToString::to_string).collect(); + options.provider_options = + Some(serde_json::from_value::(provider_options).unwrap()); + options +} + +/// aimux embeddings are `f32`: compares them with the expected JSON numbers +/// narrowed to `f32`. +fn assert_embeddings(actual: &[Vec], expected: Value) { + let expected: Vec> = expected + .as_array() + .unwrap() + .iter() + .map(|row| { + row.as_array() + .unwrap() + .iter() + .map(|v| v.as_f64().unwrap() as f32) + .collect() + }) + .collect(); + assert_eq!(actual, expected); +} + +fn fixture_embeddings() -> Value { + json!([ + [ + -0.017999587580561638, + -0.006893285550177097, + -0.036766719073057175, + -0.017558680847287178, + -0.019938766956329346, + ], + [ + -0.06007182598114014, + 0.004907649010419846, + -0.00690646655857563, + -0.007314121350646019, + -0.048464205116033554, + ], + ]) +} + +fn mock_options_body() -> Value { + json!({ + "instances": [ + {"content": "test text one", "task_type": "SEMANTIC_SIMILARITY", "title": "test title"}, + {"content": "test text two", "task_type": "SEMANTIC_SIMILARITY", "title": "test title"}, + ], + "parameters": {"autoTruncate": false, "outputDimensionality": 768}, + }) +} + +mod google_vertex_embedding_model { + use super::*; + + mod embedding { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › embedding › should extract embeddings + #[tokio::test] + async fn should_extract_embeddings() { + let server = fixture_server(&[]); + let model = model_with(MOCK_MODEL_ID, BASE_URL, server.transport()); + let result = model + .do_embed(&options( + &TEST_VALUES, + json!({"google": mock_provider_options()}), + )) + .await + .unwrap(); + assert_embeddings(&result.embeddings, fixture_embeddings()); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › embedding › should return full result snapshot + #[tokio::test] + async fn should_return_full_result_snapshot() { + let server = fixture_server(&[]); + let model = model_with(MOCK_MODEL_ID, BASE_URL, server.transport()); + let result = model + .do_embed(&options( + &TEST_VALUES, + json!({"google": mock_provider_options()}), + )) + .await + .unwrap(); + let mut snapshot: Value = + serde_json::from_str(include_str!("snapshots/embedding_full_result.json")).unwrap(); + assert_embeddings(&result.embeddings, snapshot["embeddings"].take()); + let mut actual = serde_json::to_value(&result).unwrap(); + actual["embeddings"] = Value::Null; + assert_eq!(actual, snapshot); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › embedding › should extract usage + #[tokio::test] + async fn should_extract_usage() { + let server = fixture_server(&[]); + let model = model_with(MOCK_MODEL_ID, BASE_URL, server.transport()); + let result = model + .do_embed(&options( + &TEST_VALUES, + json!({"google": mock_provider_options()}), + )) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(result.usage).unwrap(), + json!({"tokens": 11}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › embedding › should pass the model parameters correctly + #[tokio::test] + async fn should_pass_the_model_parameters_correctly() { + let server = fixture_server(&[]); + let model = model_with(MOCK_MODEL_ID, BASE_URL, server.transport()); + model + .do_embed(&options( + &TEST_VALUES, + json!({"google": mock_provider_options()}), + )) + .await + .unwrap(); + assert_eq!(server.seen()[0].json_body(), mock_options_body()); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › embedding › should accept googleVertex as provider options key + #[tokio::test] + async fn should_accept_google_vertex_as_provider_options_key() { + let server = fixture_server(&[]); + let model = model_with(MOCK_MODEL_ID, BASE_URL, server.transport()); + model + .do_embed(&options( + &TEST_VALUES, + json!({"googleVertex": mock_provider_options()}), + )) + .await + .unwrap(); + assert_eq!(server.seen()[0].json_body(), mock_options_body()); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › embedding › should pass the taskType setting in instances + #[tokio::test] + async fn should_pass_the_task_type_setting_in_instances() { + let server = fixture_server(&[]); + let model = model_with(MOCK_MODEL_ID, BASE_URL, server.transport()); + model + .do_embed(&options( + &TEST_VALUES, + json!({"google": {"taskType": "SEMANTIC_SIMILARITY"}}), + )) + .await + .unwrap(); + assert_eq!( + server.seen()[0].json_body(), + json!({ + "instances": [ + {"content": "test text one", "task_type": "SEMANTIC_SIMILARITY"}, + {"content": "test text two", "task_type": "SEMANTIC_SIMILARITY"}, + ], + "parameters": {}, + }) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › embedding › should pass the title setting in instances + #[tokio::test] + async fn should_pass_the_title_setting_in_instances() { + let server = fixture_server(&[]); + let model = model_with(MOCK_MODEL_ID, BASE_URL, server.transport()); + model + .do_embed(&options( + &TEST_VALUES, + json!({"google": {"title": "test title"}}), + )) + .await + .unwrap(); + assert_eq!( + server.seen()[0].json_body(), + json!({ + "instances": [ + {"content": "test text one", "title": "test title"}, + {"content": "test text two", "title": "test title"}, + ], + "parameters": {}, + }) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › embedding › should use embedContent for gemini-embedding-2 + #[tokio::test] + async fn should_use_embed_content_for_gemini_embedding_2() { + let server = test_server(vec![ + (DEFAULT_URL, json_response(&fixture(), &[])), + ( + GEMINI_EMBEDDING_2_URL, + json_response( + &json!({ + "embedding": {"values": [0.1, 0.2, 0.3]}, + "usageMetadata": {"promptTokenCount": 4}, + }), + &[], + ), + ), + ]); + let model = model_with("gemini-embedding-2", BASE_URL, server.transport()); + let result = model + .do_embed(&options( + &["test text one"], + json!({"google": mock_provider_options()}), + )) + .await + .unwrap(); + let seen = server.seen(); + assert_eq!(seen[0].url, GEMINI_EMBEDDING_2_URL); + assert_eq!( + seen[0].json_body(), + json!({ + "content": {"parts": [{"text": "test text one"}]}, + "embedContentConfig": { + "autoTruncate": false, + "outputDimensionality": 768, + "taskType": "SEMANTIC_SIMILARITY", + "title": "test title", + }, + }) + ); + assert_embeddings(&result.embeddings, json!([[0.1, 0.2, 0.3]])); + assert_eq!( + serde_json::to_value(result.usage).unwrap(), + json!({"tokens": 4}) + ); + } + + mod max_embeddings_per_call { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › embedding › maxEmbeddingsPerCall › should limit predict endpoint models to 250 values per call + #[test] + fn should_limit_predict_endpoint_models_to_250_values_per_call() { + let model = model_with( + "gemini-embedding-001", + BASE_URL, + MockFetch::new(Vec::new()).transport(), + ); + assert_eq!(model.max_embeddings_per_call(), Some(250)); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › embedding › maxEmbeddingsPerCall › should limit %s to one value per call + #[test] + fn should_limit_model_to_one_value_per_call() { + for model_id in ["gemini-embedding-2", "gemini-embedding-2-preview"] { + let model = + model_with(model_id, BASE_URL, MockFetch::new(Vec::new()).transport()); + assert_eq!(model.max_embeddings_per_call(), Some(1), "{model_id}"); + } + } + } + } + + mod response_headers { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › response headers › should expose response headers + #[tokio::test] + async fn should_expose_response_headers() { + let server = fixture_server(&[("test-header", "test-value")]); + let model = model_with(MOCK_MODEL_ID, BASE_URL, server.transport()); + let result = model + .do_embed(&options( + &TEST_VALUES, + json!({"google": mock_provider_options()}), + )) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(result.response.unwrap().headers).unwrap(), + json!({ + "content-length": "429", + "content-type": "application/json", + "test-header": "test-value", + }) + ); + } + } + + mod response_metadata { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › response metadata › should expose the raw response + #[tokio::test] + async fn should_expose_the_raw_response() { + let server = fixture_server(&[("test-header", "test-value")]); + let model = model_with(MOCK_MODEL_ID, BASE_URL, server.transport()); + let result = model + .do_embed(&options( + &TEST_VALUES, + json!({"google": mock_provider_options()}), + )) + .await + .unwrap(); + let snapshot: Value = + serde_json::from_str(include_str!("snapshots/embedding_raw_response.json")) + .unwrap(); + assert_eq!(serde_json::to_value(result.response).unwrap(), snapshot); + } + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › should pass headers correctly + #[tokio::test] + async fn should_pass_headers_correctly() { + let server = fixture_server(&[]); + // aimux has one entry point, and it authenticates: an access token + // stands in for upstream's credential-less base provider, and its + // `authorization` header is left out of the comparison. + let provider = create_google_vertex(VertexProviderSettings { + project: Some("test-project".into()), + location: Some("us-central1".into()), + access_token: Some(Resolvable::Value("test-token".into())), + headers: Some(Resolvable::Value( + [( + "X-Custom-Header".to_string(), + Some("custom-value".to_string()), + )] + .into(), + )), + fetch: Some(server.transport()), + ..Default::default() + }) + .unwrap(); + let mut call = options(&TEST_VALUES, json!({"google": mock_provider_options()})); + call.headers = Some([("X-Request-Header".to_string(), "request-value".to_string())].into()); + provider + .embedding(MOCK_MODEL_ID) + .do_embed(&call) + .await + .unwrap(); + let seen = server.seen(); + let mut headers: BTreeMap = seen[0].headers.clone(); + let user_agent = headers.remove("user-agent").unwrap_or_default(); + headers.remove("authorization"); + assert_eq!( + serde_json::to_value(headers).unwrap(), + json!({ + "content-type": "application/json", + "x-custom-header": "custom-value", + "x-request-header": "request-value", + }) + ); + assert!(user_agent.contains(USER_AGENT_PART), "{user_agent}"); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › should throw TooManyEmbeddingValuesForCallError when too many values provided + #[tokio::test] + async fn should_throw_too_many_embedding_values_for_call_error() { + let server = fixture_server(&[]); + let model = model_with(MOCK_MODEL_ID, BASE_URL, server.transport()); + let too_many_values = vec!["test"; 251]; + let error = model + .do_embed(&options( + &too_many_values, + json!({"google": mock_provider_options()}), + )) + .await + .unwrap_err(); + // aimux reports `TooManyEmbeddingValuesForCallError` as an invalid + // argument with the same meaning. + assert!( + matches!(&error, AiMuxError::InvalidArgument { message, .. } + if message.starts_with("Too many embedding values")), + "{error:?}" + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › should use custom baseURL when provided + #[tokio::test] + async fn should_use_custom_base_url_when_provided() { + let server = test_server(vec![(CUSTOM_URL, json_response(&fixture(), &[]))]); + let model = model_with( + "textembedding-gecko@001", + "https://custom-endpoint.com", + server.transport(), + ); + let result = model + .do_embed(&options( + &TEST_VALUES, + json!({"google": {"outputDimensionality": 768}}), + )) + .await + .unwrap(); + assert_embeddings(&result.embeddings, fixture_embeddings()); + assert_eq!( + server.seen()[0].url, + "https://custom-endpoint.com/models/textembedding-gecko@001:predict" + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-embedding-model.test.ts › GoogleVertexEmbeddingModel › should use custom fetch when provided and include proper request content + #[tokio::test] + async fn should_use_custom_fetch_when_provided_and_include_proper_request_content() { + let custom_fetch = MockFetch::new(vec![Canned { + status: 200, + headers: Vec::new(), + body: serde_json::to_vec(&fixture()).unwrap(), + }]); + let model = model_with( + "textembedding-gecko@001", + "https://custom-endpoint.com", + custom_fetch.transport(), + ); + let result = model + .do_embed(&options( + &TEST_VALUES, + json!({"google": {"outputDimensionality": 768}}), + )) + .await + .unwrap(); + assert_embeddings(&result.embeddings, fixture_embeddings()); + let seen = custom_fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].url, CUSTOM_URL); + assert_eq!( + seen[0].json_body(), + json!({ + "instances": [ + {"content": "test text one"}, + {"content": "test text two"}, + ], + "parameters": {"outputDimensionality": 768}, + }) + ); + } +} diff --git a/aimux-providers/src/vertex/tests/fixtures/aimux-test-rsa-key.pem b/aimux-providers/src/vertex/tests/fixtures/aimux-test-rsa-key.pem new file mode 100644 index 000000000..ec1b9cc75 --- /dev/null +++ b/aimux-providers/src/vertex/tests/fixtures/aimux-test-rsa-key.pem @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQD6rX7hJB3vuNVo +69TL3Gn2dep3657LWz8hqwweedq/7YjuvWpWnKTL70F4f/IHqCbMbHAjR113H5jM +MWmMD2XpNttnEb0RjYNLan3MgNLHZaJbudRrW+v1wsozv2a8E/vBMkSGXNY7Bgaz +JmbCsGD1ov90DFR7psWVRAhByhjAsHfiKosoSMhed6HIVFQMFB/PT1rr1hiHS4SK +YUEf95pnwbgWHeiDwP3atevOqi9ZxlIxQFCh5NowMPirnu8mlJNcL9J48tqNMFwn +qbVHs8HCQKCrhmczOqZDHu8kMQa7PVVRQXEkKCFYmNtpgWFk7yT4kIy//TpuecLa +cPSAZHq9AgMBAAECggEAOk4czwp6wQ19a8ZAjDsj1UynYMWnvVZnmHHjo3oFuHLs +ENVzNJqCFSH82SVvOrswbQdCzN3WPjUV8daDXW2eb1BQV0iPsJO46FNoV6bV/8Vi +rlGKd/8dIqAwLjmUUGYYfh7uQ883j+puqAAkUh+6qZ+EKa6NfuYzGi0E+6JhP/3F +yEYGUgEPOed8+VkcLNIKht+Fyv6FhssrQg07I2I0pIaRIMkapMeL5lL4MkqPtVmV +Jv+Q7KDTRjXNF1qwRf1bphUEHyBUSQui3WQ1Kk5KbiVHg0ewlOp5Q3MyjySu3ix7 +3VkVA1r5GL6HCs3arERyMfR30bq7KFLSP83HMY04MQKBgQD+kc/hTwwcRa5IY6bd +wGuEPmlqgmztBpWYZU+D5M9sNymJYp42hbf6fv06U3cpwk1tiPfqVMn2jQzpK5Yk +/BmzVJyk9EZCyE2qqIJy9epTnbKZpFS6cpga2YcCTh0K9k2dcGIdwT+y2aUD18pe +RJkfbbMY1B3HG3FdO/0rl0O0LQKBgQD8FhXWugoUZxptAfE/bf7D3YvwlFTx7TLa +wTkaFKvTknsXktv+bNuTyIAd+Oz233opP4mERYbPpsRlIbrC23+QBtMl7MH/qkBE +wuvhvURbQ0oA7CBoi3aUB3xkBRGV/V4oNcS/McBOfZ2CoFlDM/dx8FeMD1VNZLnZ +REh7Dysq0QKBgA+9TDEOBh2pYUEwiBylzP/8K+Bd+AYfphgqC2SIPRrfvj20Tzvn +wJxIRbfZeUCntsnS8iEW5PXNesLf4wSeYahoztaUE9QFNmgayxX4hmhtE6HUOzjb +5POkB2jpE3qeB+Z4EFjundWdB0rZiLWTYDRFQBGE6+Jnm/a7wsdzLshxAoGAEtn/ +9hfhixcOxdogVXL/W1xlC6+8Dn+DjxJmF8o181N/+wpTg1pWrV2QHy2vYuLhetYy +jX65jyvUgg+crNfHWVj5T1u+QkVUGxBMnYwO4ZhNn3YJEmhjjLxmlABBCWtyMCWd +OgMZc4U3Uz50jm8DcTTdBy9aXLhMoa/Wj8VZoBECgYBnltH2ntklB8hJ3NEZdNGc +P/8zZBIvzr9mkVOi8oMLQ+GvUw0EeZs9NpvsiIalxj1ZNppqN5v/EVp8wytIV2eV +3zTD7vNnuxrqj9PR67xFpnWbGfSc8r3JBDD7JzIu2cQ0g9Es1Jg3iWygWv4vtdMj +Newwb+6WXO7N9gwIi7Dnbg== +-----END PRIVATE KEY----- diff --git a/aimux-providers/src/vertex/tests/fixtures/google-vertex-embedding.json b/aimux-providers/src/vertex/tests/fixtures/google-vertex-embedding.json new file mode 100644 index 000000000..0d71535cf --- /dev/null +++ b/aimux-providers/src/vertex/tests/fixtures/google-vertex-embedding.json @@ -0,0 +1,31 @@ +{ + "predictions": [ + { + "embeddings": { + "statistics": { + "token_count": 5, + "truncated": false + }, + "values": [ + -0.017999587580561638, -0.006893285550177097, -0.036766719073057175, + -0.017558680847287178, -0.019938766956329346 + ] + } + }, + { + "embeddings": { + "statistics": { + "truncated": false, + "token_count": 6 + }, + "values": [ + -0.06007182598114014, 0.004907649010419846, -0.00690646655857563, + -0.007314121350646019, -0.048464205116033554 + ] + } + } + ], + "metadata": { + "billableCharacterCount": 35 + } +} diff --git a/aimux-providers/src/vertex/tests/gemini_transcription/gemini_transcription_model.rs b/aimux-providers/src/vertex/tests/gemini_transcription/gemini_transcription_model.rs new file mode 100644 index 000000000..9684647d8 --- /dev/null +++ b/aimux-providers/src/vertex/tests/gemini_transcription/gemini_transcription_model.rs @@ -0,0 +1,641 @@ +use std::sync::{Arc, Mutex}; + +use serde_json::{Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::shared::SharedProviderOptions; +use aimux_core::transcription_model::{ + AudioChunk, AudioInput, InputAudioFormat, TranscriptionCallOptions, TranscriptionModel, + TranscriptionStreamOptions, +}; +use aimux_provider_utils::{HeaderMapOpt, Resolvable}; + +use crate::shared::{Endpoint, EndpointConfig, ResolveEndpoint}; +use crate::test_support::mock_fetch::{Canned, RouteFetch}; +use crate::vertex::{VertexGeminiTranscriptionModel, VertexProviderSettings, create_google_vertex}; + +const BASE_URL: &str = "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google"; + +/// A base URL and headers fixed for every request (the upstream model config). +struct StaticEndpoint { + base_url: String, + headers: HeaderMapOpt, +} + +#[async_trait::async_trait] +impl ResolveEndpoint for StaticEndpoint { + async fn resolve(&self) -> Result { + Ok(Endpoint { + base_url: self.base_url.clone(), + headers: self.headers.clone(), + }) + } +} + +/// `createModel(modelId)`: `provider: 'test-provider'`, `baseURL: BASE_URL`, +/// `headers: () => ({ Authorization: 'Bearer test-oauth-token' })`, project +/// `test-project` and location `us-central1` (taken from a provider with those +/// settings), and the mock WebSocket. +fn create_model( + model_id: &str, + fetch: Option<&Arc>, + #[cfg(feature = "realtime")] web_socket: Option>, +) -> VertexGeminiTranscriptionModel { + let provider = create_google_vertex(VertexProviderSettings { + project: Some("test-project".into()), + location: Some("us-central1".into()), + access_token: Some(Resolvable::Value("unused".into())), + ..Default::default() + }) + .unwrap(); + VertexGeminiTranscriptionModel::from_config( + model_id.to_string(), + EndpointConfig::new( + "test-provider".to_string(), + StaticEndpoint { + base_url: BASE_URL.to_string(), + headers: [( + "Authorization".to_string(), + Some("Bearer test-oauth-token".to_string()), + )] + .into(), + }, + fetch.map(RouteFetch::transport), + ), + provider.project_location("transcription"), + #[cfg(feature = "realtime")] + web_socket.map(|socket| socket as Arc), + ) +} + +/// The `doGenerate` test server: the generateContent URL answers `body`, +/// anything else 404. +fn test_server(body: Value) -> Arc { + let url = format!("{BASE_URL}/models/gemini-3.5-transcribe:generateContent"); + let response = Canned::json(&body); + RouteFetch::new(move |seen| { + if seen.url == url { + response.clone() + } else { + Canned { + status: 404, + headers: Vec::new(), + body: b"Not found".to_vec(), + } + } + }) +} + +fn default_server() -> Arc { + test_server(json!({ + "candidates": [{"content": {"parts": [{"text": "Hello "}, {"text": "world."}]}}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 4}, + })) +} + +fn provider_options(value: Value) -> Option { + Some(serde_json::from_value(value).unwrap()) +} + +fn generate_options(audio: Vec, provider_options_value: Value) -> TranscriptionCallOptions { + let mut options = TranscriptionCallOptions::new(AudioInput::Binary(audio), "audio/wav"); + options.provider_options = provider_options(provider_options_value); + options +} + +/// The message of an error, without the variant's display prefix. +fn message(error: &AiMuxError) -> String { + match error { + AiMuxError::InvalidArgument { message, .. } + | AiMuxError::InvalidResponseData { message, .. } + | AiMuxError::UnsupportedFunctionality(message) => message.clone(), + AiMuxError::ApiCall(error) => error.message.clone(), + other => other.to_string(), + } +} + +#[cfg(feature = "realtime")] +use ws_mock::*; + +/// An in-process stand-in for upstream's `MockWebSocket`: the connector pairs +/// the client with a scripted server over an in-memory pipe and records what +/// the client did. +#[cfg(feature = "realtime")] +mod ws_mock { + use super::*; + + use futures::{SinkExt, StreamExt}; + use tokio_tungstenite::WebSocketStream; + use tokio_tungstenite::tungstenite::Message; + use tokio_tungstenite::tungstenite::protocol::Role; + use tokio_tungstenite::tungstenite::protocol::frame::CloseFrame; + use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; + + use aimux_provider_utils::ws::{WebSocketRequest, WsConnection, WsConnector}; + + /// One step of the server side. + pub(super) enum Step { + /// Receive the client's first (setup) message. + ReadSetup, + /// Give the client time to send more before `setupComplete`, then note + /// how many messages it had sent (`expect(ws.send).toHaveBeenCalledTimes`). + CountSentBeforeSetupComplete, + /// Send a JSON message (`ws.message(value)`). + Send(Value), + /// Receive messages until `audioStreamEnd`. + WaitForAudioStreamEnd, + /// Close from the server side (`ws.serverClose(code, reason)`). + ServerClose(u16, &'static str), + } + + #[derive(Default)] + pub(super) struct WsSeen { + pub(super) url: Option, + pub(super) headers: Vec<(String, String)>, + /// Every message the client sent, parsed (`ws.send.mock.calls`). + pub(super) sent: Vec, + pub(super) sent_before_setup_complete: Option, + /// The code of the client's close frame (`ws.close` call argument). + pub(super) close_code: Option, + } + + pub(super) struct MockWebSocket { + pub(super) seen: Arc>, + script: Mutex>>, + } + + impl MockWebSocket { + pub(super) fn new(script: Vec) -> Arc { + Arc::new(Self { + seen: Arc::default(), + script: Mutex::new(Some(script)), + }) + } + + /// Waits (bounded) until the client closed the socket. + pub(super) async fn close_code(&self) -> Option { + for _ in 0..200 { + if let Some(code) = self.seen.lock().unwrap().close_code { + return Some(code); + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + None + } + } + + #[async_trait::async_trait] + impl WsConnector for MockWebSocket { + async fn connect(&self, request: &WebSocketRequest) -> Result { + { + let mut seen = self.seen.lock().unwrap(); + seen.url = Some(request.url.clone()); + seen.headers = request.headers.clone(); + } + let script = self.script.lock().unwrap().take().expect("one connection"); + let (client_io, server_io) = tokio::io::duplex(64 * 1024); + let client = WebSocketStream::from_raw_socket(client_io, Role::Client, None).await; + let server = WebSocketStream::from_raw_socket(server_io, Role::Server, None).await; + tokio::spawn(serve(server, Arc::clone(&self.seen), script)); + Ok(WsConnection::from_transport(Box::new(client), request)) + } + } + + type Server = WebSocketStream; + + /// Receives one frame; records a text frame and returns it. + async fn receive(ws: &mut Server, seen: &Mutex) -> Option> { + match ws.next().await? { + Ok(Message::Text(text)) => { + let value: Value = serde_json::from_str(text.as_str()).expect("client sends JSON"); + seen.lock().unwrap().sent.push(value.clone()); + Some(Some(value)) + } + Ok(Message::Close(frame)) => { + seen.lock().unwrap().close_code = + Some(frame.map_or(1005, |frame| u16::from(frame.code))); + None + } + Ok(_) => Some(None), + Err(_) => None, + } + } + + async fn serve(mut ws: Server, seen: Arc>, script: Vec) { + for step in script { + match step { + Step::ReadSetup => while let Some(None) = receive(&mut ws, &seen).await {}, + Step::CountSentBeforeSetupComplete => { + while let Ok(Some(_)) = tokio::time::timeout( + std::time::Duration::from_millis(100), + receive(&mut ws, &seen), + ) + .await + {} + let count = seen.lock().unwrap().sent.len(); + seen.lock().unwrap().sent_before_setup_complete = Some(count); + } + Step::Send(value) => { + ws.send(Message::Text(value.to_string())).await.unwrap(); + } + Step::WaitForAudioStreamEnd => loop { + match receive(&mut ws, &seen).await { + Some(Some(value)) if value["realtimeInput"]["audioStreamEnd"] == true => { + break; + } + Some(_) => {} + None => return, + } + }, + Step::ServerClose(code, reason) => { + let _ = ws + .send(Message::Close(Some(CloseFrame { + code: CloseCode::from(code), + reason: reason.into(), + }))) + .await; + } + } + } + while receive(&mut ws, &seen).await.is_some() {} + } +} + +#[cfg(feature = "realtime")] +fn stream_options( + audio: futures::stream::BoxStream<'static, AudioChunk>, + rate: u32, + provider_options_value: Value, +) -> TranscriptionStreamOptions { + TranscriptionStreamOptions { + audio, + input_audio_format: InputAudioFormat { + format_type: "audio/pcm".into(), + rate: Some(rate), + }, + provider_options: provider_options(provider_options_value), + abort_signal: None, + headers: None, + include_raw_chunks: false, + timeout: None, + } +} + +#[cfg(feature = "realtime")] +fn chunks(chunks: Vec>) -> futures::stream::BoxStream<'static, AudioChunk> { + use futures::StreamExt; + futures::stream::iter(chunks.into_iter().map(AudioChunk::Binary)).boxed() +} + +/// `convertReadableStreamToArray(result.stream)`, each part as its JSON. +#[cfg(feature = "realtime")] +async fn collect_parts( + stream: std::pin::Pin< + Box< + dyn futures::Stream< + Item = Result< + aimux_core::transcription_model::TranscriptionStreamPart, + AiMuxError, + >, + > + Send, + >, + >, +) -> Result, AiMuxError> { + use futures::StreamExt; + let parts: Vec<_> = tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect()) + .await + .expect("the stream ends"); + parts + .into_iter() + .map(|part| part.map(|part| serde_json::to_value(part).unwrap())) + .collect() +} + +mod do_generate { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/gemini-transcription/google-vertex-gemini-transcription-model.test.ts › doGenerate › transcribes audio via Vertex generateContent with audioTranscriptionConfig + #[tokio::test] + async fn transcribes_audio_via_vertex_generate_content_with_audio_transcription_config() { + let server = default_server(); + let model = create_model( + "gemini-3.5-transcribe", + Some(&server), + #[cfg(feature = "realtime")] + None, + ); + let result = model + .do_generate(&generate_options( + vec![1, 2, 3, 4], + json!({ + "googleVertex": { + "customVocabulary": ["Gemini", "Kubernetes"], + "languageCodes": ["es-ES"], + "mode": "SMART", + }, + }), + )) + .await + .unwrap(); + assert_eq!(result.text, "Hello world."); + let seen = server.seen(); + assert_eq!( + seen[0].json_body(), + json!({ + "contents": [{ + "role": "user", + "parts": [{"inlineData": {"mimeType": "audio/wav", "data": "AQIDBA=="}}], + }], + "generationConfig": { + "audioTranscriptionConfig": { + "languageCodes": ["es-ES"], + "customVocabulary": ["Gemini", "Kubernetes"], + "mode": "SMART", + }, + }, + }) + ); + assert_eq!( + seen[0].headers.get("authorization").map(String::as_str), + Some("Bearer test-oauth-token") + ); + assert_eq!( + serde_json::to_value(result.provider_metadata).unwrap(), + json!({"google": {"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 4}}}) + ); + // `result.usage` is not asserted: aimux's `TranscriptionResult` has no + // `usage` field. + } + + /// upstream: @ai-sdk/google-vertex src/gemini-transcription/google-vertex-gemini-transcription-model.test.ts › doGenerate › accepts options under the google namespace as a fallback + #[tokio::test] + async fn accepts_options_under_the_google_namespace_as_a_fallback() { + let server = default_server(); + let model = create_model( + "gemini-3.5-transcribe", + Some(&server), + #[cfg(feature = "realtime")] + None, + ); + model + .do_generate(&generate_options( + vec![1, 2, 3, 4], + json!({"google": {"mode": "SMART"}}), + )) + .await + .unwrap(); + assert_eq!( + server.seen()[0].json_body()["generationConfig"], + json!({"audioTranscriptionConfig": {"mode": "SMART"}}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/gemini-transcription/google-vertex-gemini-transcription-model.test.ts › doGenerate › extracts text, language, and word segments from the audioTranscription part + #[tokio::test] + async fn extracts_text_language_and_word_segments_from_the_audio_transcription_part() { + let server = test_server(json!({ + "candidates": [{ + "content": { + "parts": [{ + "audioTranscription": { + "text": "The quick brown fox.", + "languageCode": "en-US", + "speakerLabel": "spk:0", + "words": [ + {"word": "The", "startOffset": "0.100s", "endOffset": "0.100s"}, + {"word": "quick", "startOffset": "0.100s", "endOffset": "0.400s"}, + {"word": "brown", "startOffset": "0.400s", "endOffset": "0.700s"}, + {"word": "fox.", "startOffset": "0.700s", "endOffset": "1s"}, + ], + }, + }], + }, + }], + "usageMetadata": {"promptTokenCount": 64}, + })); + let model = create_model( + "gemini-3.5-transcribe", + Some(&server), + #[cfg(feature = "realtime")] + None, + ); + let result = model + .do_generate(&generate_options(vec![1, 2, 3, 4], json!({}))) + .await + .unwrap(); + assert_eq!(result.text, "The quick brown fox."); + assert_eq!(result.language.as_deref(), Some("en-US")); + assert_eq!( + serde_json::to_value(&result.segments).unwrap(), + json!([ + {"text": "The", "startSecond": 0.1, "endSecond": 0.1}, + {"text": "quick", "startSecond": 0.1, "endSecond": 0.4}, + {"text": "brown", "startSecond": 0.4, "endSecond": 0.7}, + {"text": "fox.", "startSecond": 0.7, "endSecond": 1.0}, + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/gemini-transcription/google-vertex-gemini-transcription-model.test.ts › doGenerate › rejects unary transcription on live model ids + #[tokio::test] + async fn rejects_unary_transcription_on_live_model_ids() { + let server = default_server(); + let model = create_model( + "gemini-3.5-transcribe-live", + Some(&server), + #[cfg(feature = "realtime")] + None, + ); + let error = model + .do_generate(&generate_options(vec![1], json!({}))) + .await + .unwrap_err(); + assert!( + message(&error).contains("only supports streaming transcription"), + "{error:?}" + ); + } +} + +#[cfg(feature = "realtime")] +mod do_stream { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/gemini-transcription/google-vertex-gemini-transcription-model.test.ts › doStream › streams transcription over the Vertex Live API WebSocket + #[tokio::test] + async fn streams_transcription_over_the_vertex_live_api_web_socket() { + let ws = MockWebSocket::new(vec![ + Step::ReadSetup, + Step::CountSentBeforeSetupComplete, + Step::Send(json!({"setupComplete": {}})), + Step::WaitForAudioStreamEnd, + Step::Send(json!({"serverContent": {"interimInputTranscription": {"text": "hel"}}})), + Step::Send(json!({ + "serverContent": {"inputTranscription": {"text": "hello world.", "finished": true}}, + })), + Step::Send(json!({ + "usageMetadata": {"promptTokenCount": 7}, + "serverContent": {"turnComplete": true, "interactionStatus": "IDLE"}, + })), + ]); + let model = create_model("gemini-3.5-transcribe-live", None, Some(ws.clone())); + let result = model + .do_stream(stream_options( + chunks(vec![vec![1, 2, 3]]), + 16000, + json!({"googleVertex": {"customVocabulary": ["Gemini"], "mode": "SMART"}}), + )) + .expect("streaming is supported") + .await + .unwrap(); + let parts = collect_parts(result.stream).await.unwrap(); + + { + let seen = ws.seen.lock().unwrap(); + assert_eq!( + seen.url.as_deref(), + Some( + "wss://us-central1-aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" + ) + ); + // Header names compared case-insensitively. + let headers: std::collections::BTreeMap = seen + .headers + .iter() + .map(|(name, value)| (name.to_ascii_lowercase(), value.clone())) + .collect(); + assert_eq!( + serde_json::to_value(headers).unwrap(), + json!({"authorization": "Bearer test-oauth-token"}) + ); + + // audio is gated on setupComplete: only the setup message was sent + assert_eq!(seen.sent_before_setup_complete, Some(1)); + assert_eq!( + seen.sent[0], + json!({ + "setup": { + "model": "projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-transcribe-live", + "inputAudioTranscription": {"customVocabulary": ["Gemini"], "mode": "SMART"}, + }, + }) + ); + + // audio chunk + audioStreamEnd + assert_eq!(seen.sent.len(), 3); + assert_eq!( + seen.sent[1], + json!({"realtimeInput": {"audio": {"data": "AQID", "mimeType": "audio/pcm;rate=16000"}}}) + ); + assert_eq!( + seen.sent[2], + json!({"realtimeInput": {"audioStreamEnd": true}}) + ); + } + + // The finish part's `usage` is left out: aimux's finish part has no + // `usage` field. + assert_eq!( + parts, + vec![ + json!({"type": "stream-start", "warnings": []}), + json!({"type": "transcript-partial", "id": "google-segment-0", "text": "hel"}), + json!({"type": "transcript-delta", "id": "google-segment-0", "delta": "hello world."}), + json!({"type": "transcript-final", "id": "google-segment-0", "text": "hello world."}), + json!({ + "type": "finish", + "text": "hello world.", + "segments": [], + "providerMetadata": {"google": {"usageMetadata": {"promptTokenCount": 7}}}, + }), + ] + ); + assert_eq!(ws.close_code().await, Some(1000)); + } + + /// upstream: @ai-sdk/google-vertex src/gemini-transcription/google-vertex-gemini-transcription-model.test.ts › doStream › falls back to the latest interim partial when no final segment arrives + #[tokio::test] + async fn falls_back_to_the_latest_interim_partial_when_no_final_segment_arrives() { + let ws = MockWebSocket::new(vec![ + Step::ReadSetup, + Step::Send(json!({"setupComplete": {}})), + Step::WaitForAudioStreamEnd, + Step::Send( + json!({"serverContent": {"interimInputTranscription": {"text": "hello world."}}}), + ), + Step::Send(json!({"serverContent": {"interactionStatus": "IDLE"}})), + ]); + let model = create_model("gemini-3.5-transcribe-live", None, Some(ws.clone())); + let result = model + .do_stream(stream_options(chunks(vec![vec![1, 2]]), 16000, json!({}))) + .expect("streaming is supported") + .await + .unwrap(); + let parts = collect_parts(result.stream).await.unwrap(); + // `toMatchObject` + let last = parts.last().expect("a part"); + assert_eq!(last["type"], "finish"); + assert_eq!(last["text"], "hello world."); + } + + /// upstream: @ai-sdk/google-vertex src/gemini-transcription/google-vertex-gemini-transcription-model.test.ts › doStream › errors when the socket closes before the audio ended + #[tokio::test] + async fn errors_when_the_socket_closes_before_the_audio_ended() { + use futures::StreamExt; + + let ws = MockWebSocket::new(vec![ + Step::ReadSetup, + Step::Send(json!({"setupComplete": {}})), + Step::ServerClose(1011, "internal error"), + ]); + let model = create_model("gemini-3.5-transcribe-live", None, Some(ws.clone())); + let result = model + .do_stream(stream_options( + futures::stream::pending().boxed(), + 16000, + json!({}), + )) + .expect("streaming is supported") + .await + .unwrap(); + let error = collect_parts(result.stream).await.unwrap_err(); + assert!( + message(&error).contains( + "closed unexpectedly before finishing (code 1011, reason: internal error)" + ), + "{error:?}" + ); + } + + /// upstream: @ai-sdk/google-vertex src/gemini-transcription/google-vertex-gemini-transcription-model.test.ts › doStream › rejects streaming on unary model ids + #[tokio::test] + async fn rejects_streaming_on_unary_model_ids() { + let ws = MockWebSocket::new(Vec::new()); + let model = create_model("gemini-3.5-transcribe", None, Some(ws)); + let error = model + .do_stream(stream_options(chunks(vec![vec![1]]), 16000, json!({}))) + .expect("the model implements doStream") + .await + .unwrap_err(); + assert!( + message(&error).contains("does not support streaming transcription"), + "{error:?}" + ); + } + + /// upstream: @ai-sdk/google-vertex src/gemini-transcription/google-vertex-gemini-transcription-model.test.ts › doStream › rejects non-16kHz PCM input audio + #[tokio::test] + async fn rejects_non_16khz_pcm_input_audio() { + let ws = MockWebSocket::new(Vec::new()); + let model = create_model("gemini-3.5-transcribe-live", None, Some(ws)); + let error = model + .do_stream(stream_options(chunks(vec![vec![1]]), 24000, json!({}))) + .expect("the model implements doStream") + .await + .unwrap_err(); + assert!( + message(&error).contains("only supports 16kHz 16-bit PCM input audio"), + "{error:?}" + ); + } +} diff --git a/aimux-providers/src/vertex/tests/gemini_transcription/mod.rs b/aimux-providers/src/vertex/tests/gemini_transcription/mod.rs new file mode 100644 index 000000000..03b4825b8 --- /dev/null +++ b/aimux-providers/src/vertex/tests/gemini_transcription/mod.rs @@ -0,0 +1 @@ +mod gemini_transcription_model; diff --git a/aimux-providers/src/vertex/tests/image_model.rs b/aimux-providers/src/vertex/tests/image_model.rs new file mode 100644 index 000000000..c2b9e6969 --- /dev/null +++ b/aimux-providers/src/vertex/tests/image_model.rs @@ -0,0 +1,389 @@ +use std::sync::Arc; + +use serde_json::{Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::image_model::{ + ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, +}; +use aimux_core::shared::{AspectRatio, SharedProviderOptions, Size}; +use aimux_provider_utils::HeaderMapOpt; + +use crate::shared::{Endpoint, EndpointConfig, ResolveEndpoint}; +use crate::test_support::mock_fetch::{Canned, RouteFetch}; +use crate::vertex::VertexImageModel; + +const TEST_URL: &str = "https://api.example.com/models/gemini-2.5-flash-image:generateContent"; + +/// A base URL and headers fixed for every request (the upstream model config). +struct StaticEndpoint { + base_url: String, + headers: HeaderMapOpt, +} + +#[async_trait::async_trait] +impl ResolveEndpoint for StaticEndpoint { + async fn resolve(&self) -> Result { + Ok(Endpoint { + base_url: self.base_url.clone(), + headers: self.headers.clone(), + }) + } +} + +/// `new GoogleVertexImageModel(modelId, { provider: 'google.vertex.image', baseURL: 'https://api.example.com', headers })` +/// against the test server. +fn model_with( + model_id: &str, + headers: &[(&str, &str)], + server: &Arc, +) -> VertexImageModel { + VertexImageModel::from_config( + model_id.to_string(), + EndpointConfig::new( + "google.vertex.image".to_string(), + StaticEndpoint { + base_url: "https://api.example.com".to_string(), + headers: headers + .iter() + .map(|(k, v)| (k.to_string(), Some(v.to_string()))) + .collect(), + }, + Some(server.transport()), + ), + ) +} + +/// The shared `model` of the upstream file. +fn model(server: &Arc) -> VertexImageModel { + model_with("gemini-2.5-flash-image", &[("api-key", "test-key")], server) +} + +fn json_response(body: &Value, headers: &[(&str, &str)]) -> Canned { + let body = serde_json::to_vec(body).unwrap(); + let mut all = vec![ + ("content-type".to_string(), "application/json".to_string()), + ("content-length".to_string(), body.len().to_string()), + ]; + all.extend(headers.iter().map(|(k, v)| (k.to_string(), v.to_string()))); + Canned { + status: 200, + headers: all, + body, + } +} + +/// The AI SDK test server: `TEST_URL` answers `response` (when set), anything +/// else 404. +fn test_server(response: Option) -> Arc { + RouteFetch::new(move |seen| match &response { + Some(canned) if seen.url == TEST_URL => canned.clone(), + _ => Canned { + status: 404, + headers: Vec::new(), + body: b"Not found".to_vec(), + }, + }) +} + +/// `prepareJsonResponse({ usage, headers })` with the default image. +fn prepare_json_response(usage: Option, headers: &[(&str, &str)]) -> Arc { + let usage = usage.unwrap_or_else( + || json!({"promptTokenCount": 10, "candidatesTokenCount": 100, "totalTokenCount": 110}), + ); + test_server(Some(json_response( + &json!({ + "candidates": [{ + "content": { + "parts": [{"inlineData": {"mimeType": "image/png", "data": "base64-generated-image"}}], + "role": "model", + }, + "finishReason": "STOP", + }], + "usageMetadata": usage, + }), + headers, + ))) +} + +fn provider_options(value: Value) -> SharedProviderOptions { + serde_json::from_value(value).unwrap() +} + +/// The message of an error, without the variant's display prefix. +fn message(error: &AiMuxError) -> String { + match error { + AiMuxError::InvalidArgument { message, .. } + | AiMuxError::UnsupportedFunctionality(message) + | AiMuxError::InvalidPrompt { message, .. } + | AiMuxError::InvalidResponseData { message, .. } => message.clone(), + other => other.to_string(), + } +} + +mod google_vertex_image_model { + use super::*; + + mod capabilities { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-image-model.test.ts › GoogleVertexImageModel › capabilities › advertises file=$supportsFileInputs and mask=$supportsMaskInputs for $modelId + #[test] + fn advertises_file_and_mask_support_for_model_id() { + let server = test_server(None); + for (model_id, supports_file_inputs, supports_mask_inputs) in [ + ("gemini-2.5-flash-image", Some(true), Some(false)), + ("gemini-3-pro-image-preview", Some(true), Some(false)), + ("gemini-3.1-flash-image-preview", Some(true), Some(false)), + ("gemini-2.5-pro", None, None), + ("gemini-3-pro-preview", None, None), + ("gemini-custom", None, None), + ("legacy-image-model", None, None), + ] { + let capability_model = model_with(model_id, &[], &server); + assert_eq!( + capability_model.supports_file_inputs(), + supports_file_inputs, + "supportsFileInputs for {model_id}" + ); + assert_eq!( + capability_model.supports_mask_inputs(), + supports_mask_inputs, + "supportsMaskInputs for {model_id}" + ); + } + } + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-image-model.test.ts › GoogleVertexImageModel › should return 1 for maxImagesPerCall + #[test] + fn should_return_1_for_max_images_per_call() { + let server = test_server(None); + assert_eq!(model(&server).max_images_per_call(), Some(1)); + } + + mod do_generate { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-image-model.test.ts › GoogleVertexImageModel › doGenerate › should reject non-Gemini model IDs before sending a request + #[tokio::test] + async fn should_reject_non_gemini_model_ids_before_sending_a_request() { + let server = test_server(None); + let non_gemini_model = + model_with("legacy-image-model", &[("api-key", "test-key")], &server); + let error = non_gemini_model + .do_generate(&ImageCallOptions::new("A beautiful sunset")) + .await + .unwrap_err(); + assert!( + message(&error).contains( + "Google image models other than Gemini are no longer supported. Use a model ID that starts with `gemini-`." + ), + "{error:?}" + ); + assert!(server.seen().is_empty()); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-image-model.test.ts › GoogleVertexImageModel › doGenerate › should use the language model endpoint and extract generated images + #[tokio::test] + async fn should_use_the_language_model_endpoint_and_extract_generated_images() { + let server = prepare_json_response(None, &[]); + let result = model(&server) + .do_generate(&ImageCallOptions::new("A beautiful sunset")) + .await + .unwrap(); + assert_eq!(server.seen()[0].url, TEST_URL); + assert!( + matches!(&result.images, ImageOutputs::Base64(images) if images == &["base64-generated-image"]), + "{:?}", + result.images + ); + assert_eq!( + serde_json::to_value(result.provider_metadata).unwrap(), + json!({ + "googleVertex": {"images": [{}]}, + "vertex": {"images": [{}]}, + }) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-image-model.test.ts › GoogleVertexImageModel › doGenerate › should send response modalities, aspect ratio, seed, and headers + #[tokio::test] + async fn should_send_response_modalities_aspect_ratio_seed_and_headers() { + let server = prepare_json_response(None, &[]); + let model_with_headers = model_with( + "gemini-2.5-flash-image", + &[("Custom-Provider-Header", "provider-header-value")], + &server, + ); + let mut options = ImageCallOptions::new("A beautiful sunset"); + options.aspect_ratio = Some(AspectRatio::new(21, 9)); + options.seed = Some(12345); + options.headers = Some( + [( + "Custom-Request-Header".to_string(), + "request-header-value".to_string(), + )] + .into(), + ); + model_with_headers.do_generate(&options).await.unwrap(); + let seen = server.seen(); + assert_eq!( + seen[0] + .headers + .get("custom-provider-header") + .map(String::as_str), + Some("provider-header-value") + ); + assert_eq!( + seen[0] + .headers + .get("custom-request-header") + .map(String::as_str), + Some("request-header-value") + ); + assert_eq!( + seen[0].json_body(), + json!({ + "contents": [{"parts": [{"text": "A beautiful sunset"}], "role": "user"}], + "generationConfig": { + "imageConfig": {"aspectRatio": "21:9"}, + "responseModalities": ["IMAGE"], + "seed": 12345, + }, + }) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-image-model.test.ts › GoogleVertexImageModel › doGenerate › should accept the legacy vertex options key and prefer googleVertex + #[tokio::test] + async fn should_accept_the_legacy_vertex_options_key_and_prefer_google_vertex() { + let server = prepare_json_response(None, &[]); + let mut options = ImageCallOptions::new("A beautiful sunset"); + options.aspect_ratio = Some(AspectRatio::new(16, 9)); + options.provider_options = provider_options(json!({ + "googleVertex": {"imageConfig": {"imageSize": "4K"}}, + "vertex": {"imageConfig": {"imageSize": "2K"}}, + })); + model(&server).do_generate(&options).await.unwrap(); + assert_eq!( + server.seen()[0].json_body()["generationConfig"]["imageConfig"], + json!({"aspectRatio": "16:9", "imageSize": "4K"}) + ); + + let mut options = ImageCallOptions::new("A beautiful sunset"); + options.provider_options = + provider_options(json!({"vertex": {"imageConfig": {"imageSize": "2K"}}})); + model(&server).do_generate(&options).await.unwrap(); + assert_eq!( + server.seen()[1].json_body()["generationConfig"]["imageConfig"], + json!({"imageSize": "2K"}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-image-model.test.ts › GoogleVertexImageModel › doGenerate › should include usage and response metadata + #[tokio::test] + async fn should_include_usage_and_response_metadata() { + let server = prepare_json_response( + Some( + json!({"promptTokenCount": 20, "candidatesTokenCount": 200, "totalTokenCount": 220}), + ), + &[("request-id", "test-request-id")], + ); + // aimux cannot inject `currentDate`: the timestamp is asserted to + // be set, not to equal a fixed date. + let before = aimux_core::Timestamp::now(); + let result = model(&server) + .do_generate(&ImageCallOptions::new("A beautiful sunset")) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(result.usage).unwrap(), + json!({"inputTokens": 20, "outputTokens": 200, "totalTokens": 220}) + ); + let timestamp = result.response.timestamp.expect("timestamp"); + assert!(timestamp.unix_millis() >= before.unix_millis()); + assert_eq!( + result.response.model_id.as_deref(), + Some("gemini-2.5-flash-image") + ); + assert_eq!( + result + .response + .headers + .as_ref() + .and_then(|headers| headers.get("request-id")) + .map(String::as_str), + Some("test-request-id") + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-image-model.test.ts › GoogleVertexImageModel › doGenerate › should return a warning for the unsupported size option + #[tokio::test] + async fn should_return_a_warning_for_the_unsupported_size_option() { + let server = prepare_json_response(None, &[]); + let mut options = ImageCallOptions::new("A beautiful sunset"); + options.size = Some(Size::new(1024, 1024)); + let result = model(&server).do_generate(&options).await.unwrap(); + assert_eq!( + serde_json::to_value(result.warnings).unwrap(), + json!([{ + "type": "unsupported", + "feature": "size", + "details": "This model does not support the `size` option. Use `aspectRatio` instead.", + }]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-image-model.test.ts › GoogleVertexImageModel › doGenerate › should include input images for editing + #[tokio::test] + async fn should_include_input_images_for_editing() { + let server = prepare_json_response(None, &[]); + let mut options = ImageCallOptions::new("Add a hat to this cat"); + options.files = Some(vec![ImageFile::File { + media_type: "image/png".into(), + data: ImageFileData::Base64("base64-source-image".into()), + }]); + model(&server).do_generate(&options).await.unwrap(); + assert_eq!( + server.seen()[0].json_body()["contents"], + json!([{ + "role": "user", + "parts": [ + {"text": "Add a hat to this cat"}, + {"inlineData": {"mimeType": "image/png", "data": "base64-source-image"}}, + ], + }]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-image-model.test.ts › GoogleVertexImageModel › doGenerate › should reject unsupported URL editing input and masks + #[tokio::test] + async fn should_reject_unsupported_url_editing_input_and_masks() { + let server = prepare_json_response(None, &[]); + let mut options = ImageCallOptions::new("Add a hat to this cat"); + options.files = Some(vec![ImageFile::Url { + url: "https://example.com/cat.png".into(), + }]); + let error = model(&server).do_generate(&options).await.unwrap_err(); + let text = message(&error); + let media_type = text.find("media type \"image/*\""); + assert!( + media_type.is_some_and(|at| text[at..].contains("not passed as inline bytes")), + "{error:?}" + ); + + let mut options = ImageCallOptions::new("Edit this image"); + options.mask = Some(ImageFile::File { + media_type: "image/png".into(), + data: ImageFileData::Base64("base64-mask-image".into()), + }); + let error = model(&server).do_generate(&options).await.unwrap_err(); + assert!( + message(&error) + .contains("Gemini image models do not support mask-based image editing."), + "{error:?}" + ); + } + } +} diff --git a/aimux-providers/src/vertex/tests/language_model.rs b/aimux-providers/src/vertex/tests/language_model.rs new file mode 100644 index 000000000..a37549f30 --- /dev/null +++ b/aimux-providers/src/vertex/tests/language_model.rs @@ -0,0 +1,76 @@ +use serde_json::json; + +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::{CallOptions, Tool}; +use aimux_core::tool::FunctionTool; +use aimux_provider_utils::Resolvable; + +use crate::test_support::mock_fetch::{Canned, RouteFetch}; +use crate::vertex::{VertexProviderSettings, create_google_vertex}; + +const TEST_URL: &str = "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-2.5-flash:generateContent"; + +/// upstream: @ai-sdk/google-vertex src/google-vertex-language-model.test.ts › should preserve local JSON Schema references in tool requests +#[tokio::test] +async fn should_preserve_local_json_schema_references_in_tool_requests() { + let response = Canned::json(&json!({ + "candidates": [{ + "content": {"parts": [{"text": "Done"}], "role": "model"}, + "finishReason": "STOP", + "index": 0, + }], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + })); + // The AI SDK test server: `TEST_URL` answers, anything else 404. + let server = RouteFetch::new(move |seen| { + if seen.url == TEST_URL { + response.clone() + } else { + Canned { + status: 404, + headers: Vec::new(), + body: b"Not found".to_vec(), + } + } + }); + let provider = create_google_vertex(VertexProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".into())), + fetch: Some(server.transport()), + ..Default::default() + }) + .unwrap(); + let mut options = CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]); + options.tools = Some(vec![Tool::Function(FunctionTool { + name: "format-date".into(), + description: Some("Format a date".into()), + input_schema: json!({ + "type": "object", + "properties": { + "locale": {"$ref": "#/$defs/Locale", "description": "Locale for formatting"}, + }, + "required": ["locale"], + "additionalProperties": false, + "$defs": {"Locale": {"type": "string", "enum": ["de", "en"]}}, + }), + strict: None, + provider_options: None, + input_examples: None, + })]); + provider + .call("gemini-2.5-flash") + .do_generate(&options) + .await + .unwrap(); + assert_eq!( + server.seen()[0].json_body()["tools"][0]["functionDeclarations"][0]["parametersJsonSchema"], + json!({ + "type": "object", + "properties": { + "locale": {"$ref": "#/$defs/Locale", "description": "Locale for formatting"}, + }, + "required": ["locale"], + "additionalProperties": false, + "$defs": {"Locale": {"type": "string", "enum": ["de", "en"]}}, + }) + ); +} diff --git a/aimux-providers/src/vertex/tests/location_validation.rs b/aimux-providers/src/vertex/tests/location_validation.rs new file mode 100644 index 000000000..0396da0c8 --- /dev/null +++ b/aimux-providers/src/vertex/tests/location_validation.rs @@ -0,0 +1,154 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 `src/google-vertex-location-validation.test.ts`: +//! the cases the old `tests/vertex_location_validation_test.rs` covers only +//! in part (its other cases carry their markers there). +//! +//! Upstream throws when the model is created; aimux resolves the location per +//! request, so the request fails before anything is fetched. The upstream +//! `describe` is `${name} location validation` for `google`, `anthropic` and +//! `maas`; one test covers all three. Upstream's `argument: 'location'` has no +//! counterpart: aimux's `InvalidArgument` carries a message, which names the +//! location. + +use std::sync::Arc; + +use serial_test::serial; + +use crate::test_support::mock_fetch::{EnvVar, MockFetch}; +use crate::vertex::{ + VertexAnthropicProviderSettings, VertexPartnerProviderSettings, VertexProviderSettings, + create_google_vertex, create_google_vertex_anthropic, create_google_vertex_maas, +}; +use aimux_core::AiMuxError; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_core::transcription_model::{ + AudioInput, InputAudioFormat, TranscriptionCallOptions, TranscriptionStreamOptions, +}; +use aimux_provider_utils::Resolvable; + +/// Whether `error` is the invalid-location error. +fn is_location_error(error: &AiMuxError) -> bool { + matches!( + error, + AiMuxError::InvalidArgument { message, .. } if message.starts_with("Invalid Google Vertex location.") + ) +} + +mod name_location_validation { + use super::*; + + #[derive(Clone, Copy, Debug)] + enum Name { + Google, + Anthropic, + Maas, + } + + /// `create({ baseURL })('test-model')` for each `name`, with a fixed + /// token so that only the settings decide whether a request is sent. + fn create(name: Name, base_url: &str, fetch: &Arc) -> Arc { + let base_url = Some(base_url.to_string()); + let token = Some(Resolvable::Value(Some("test-token".to_string()))); + match name { + Name::Google => create_google_vertex(VertexProviderSettings { + base_url, + access_token: Some("test-token".to_string().into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() + .call("test-model"), + Name::Anthropic => create_google_vertex_anthropic(VertexAnthropicProviderSettings { + base_url, + fetch: Some(fetch.transport()), + generate_auth_token: token, + ..Default::default() + }) + .unwrap() + .call("test-model"), + Name::Maas => Arc::new( + create_google_vertex_maas(VertexPartnerProviderSettings { + base_url, + fetch: Some(fetch.transport()), + generate_auth_token: token, + ..Default::default() + }) + .unwrap() + .chat("test-model"), + ), + } + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-location-validation.test.ts › ${…} › allows a custom URL without a location + #[tokio::test] + #[serial] + async fn allows_a_custom_url_without_a_location() { + let _location = EnvVar::set("GOOGLE_VERTEX_LOCATION", None); + for name in [Name::Google, Name::Anthropic, Name::Maas] { + let fetch = MockFetch::new(vec![]); + let model = create(name, "https://proxy.example/v1", &fetch); + // The mock answers nothing, so the request fails after it is sent. + let error = model + .do_generate(&CallOptions::new(vec![LanguageModelMessage::user_text( + "Hello", + )])) + .await + .expect_err("the mock answers nothing"); + assert!(!is_location_error(&error), "{name:?}: {error:?}"); + assert!( + fetch + .seen() + .first() + .is_some_and(|seen| seen.url.starts_with("https://proxy.example/v1/")), + "{name:?} sent no request: {error:?}" + ); + } + } +} + +/// upstream: @ai-sdk/google-vertex src/google-vertex-location-validation.test.ts › validates location for the separate transcription host of %s +#[tokio::test] +async fn validates_location_for_the_separate_transcription_host_of_s() { + for model_id in ["chirp_3", "gemini-3.5-transcribe-live"] { + let fetch = MockFetch::new(vec![]); + let model = create_google_vertex(VertexProviderSettings { + project: Some("test".into()), + location: Some("user@internal/#".into()), + base_url: Some("https://proxy.example/v1".into()), + access_token: Some("test-token".to_string().into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() + .transcription(model_id); + + // The live model only streams (over a WebSocket); the others are + // unary. + let error = if model_id.contains("-live") { + let options = TranscriptionStreamOptions { + audio: Box::pin(futures::stream::empty()), + input_audio_format: InputAudioFormat { + format_type: "audio/pcm".into(), + rate: Some(16_000), + }, + provider_options: None, + abort_signal: None, + headers: None, + include_raw_chunks: false, + timeout: None, + }; + match model.do_stream(options).expect("a streaming model").await { + Ok(_) => panic!("{model_id}: the stream started"), + Err(error) => error, + } + } else { + let options = + TranscriptionCallOptions::new(AudioInput::Binary(vec![0, 1, 2]), "audio/wav"); + model.do_generate(&options).await.unwrap_err() + }; + + assert!(is_location_error(&error), "{model_id}: {error:?}"); + assert!(fetch.seen().is_empty(), "{model_id}"); + } +} diff --git a/aimux-providers/src/vertex/tests/maas/edge/maas_provider_edge.rs b/aimux-providers/src/vertex/tests/maas/edge/maas_provider_edge.rs new file mode 100644 index 000000000..5e8e5eb33 --- /dev/null +++ b/aimux-providers/src/vertex/tests/maas/edge/maas_provider_edge.rs @@ -0,0 +1,251 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 +//! `src/maas/edge/google-vertex-maas-provider-edge.test.ts`. +//! +//! aimux has one entry point and no separate Web-Crypto edge auth; these +//! cases exercise the same behaviour through it. Upstream mocks the base +//! provider and calls the fetch wrapper the edge entry point hands it with a +//! hand-made `(url, init)`; here each case makes a chat request through a +//! mock transport, so the request URL is the provider's (not +//! `https://example.com/test`) and `init` is the request the +//! OpenAI-compatible model builds. Upstream's mocked `generateAuthToken` +//! (returning `mock-auth-token`) is the `generate_auth_token` setting, except +//! for `googleCredentials` (aimux: `GoogleAuthOptions::credentials`), which +//! runs aimux's Google auth. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_provider_utils::{HeaderMapOpt, Resolvable}; +use serde_json::json; + +use crate::test_support::mock_fetch::{Canned, MockFetch, Seen}; +use crate::vertex::{ + GoogleAuthOptions, VertexMaasProviderSettings, create_google_vertex_maas, google_vertex_maas, +}; + +/// A minimal Chat Completions response. +fn chat_response() -> Canned { + Canned::json(&json!({ + "id": "chatcmpl-1", "object": "chat.completion", "created": 1, "model": "m", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + })) +} + +fn prompt() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +/// A token generator returning `mock-auth-token`, and its call counter. +fn mock_generator() -> (Resolvable>, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let generator = Resolvable::from_async_fn(move || { + counter.fetch_add(1, Ordering::SeqCst); + async { Ok(Some("mock-auth-token".to_string())) } + }); + (generator, calls) +} + +/// Settings of a provider for `test-project` over `fetch`, with the mocked +/// token generator. +fn settings(fetch: &Arc) -> (VertexMaasProviderSettings, Arc) { + let (generator, calls) = mock_generator(); + let settings = VertexMaasProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + fetch: Some(fetch.transport()), + generate_auth_token: Some(generator), + ..Default::default() + }; + (settings, calls) +} + +/// Generate once with `options` and return the request the transport saw. +async fn request( + settings: VertexMaasProviderSettings, + fetch: &Arc, + options: &CallOptions, +) -> Seen { + create_google_vertex_maas(settings) + .unwrap() + .call("test-model") + .do_generate(options) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + seen[0].clone() +} + +/// The request headers except `skip`. +fn sent_headers(seen: &Seen, skip: &[&str]) -> BTreeMap { + seen.headers + .iter() + .filter(|(name, _)| !skip.contains(&name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +fn headers(entries: &[(&str, &str)]) -> BTreeMap { + entries + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() +} + +mod google_vertex_maas_provider_edge { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/maas/edge/google-vertex-maas-provider-edge.test.ts › google-vertex-maas-provider-edge › should generate auth token and add to request headers + /// + /// The `undefined` credentials argument has no counterpart once the + /// generator is replaced; the generator's single call is asserted. + #[tokio::test] + async fn should_generate_auth_token_and_add_to_request_headers() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, calls) = settings(&fetch); + + let seen = request(settings, &fetch, &prompt()).await; + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(seen.method, "POST"); + assert_eq!( + sent_headers(&seen, &["user-agent"]), + headers(&[ + ("content-type", "application/json"), + ("authorization", "Bearer mock-auth-token"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/maas/edge/google-vertex-maas-provider-edge.test.ts › google-vertex-maas-provider-edge › should pass googleCredentials to generateAuthToken + /// + /// The credentials reach aimux's real Google auth (upstream mocks the + /// generator), which signs with the given private key: `test-key` is not + /// a key, so the request fails on it before anything is sent. + #[tokio::test] + async fn should_pass_google_credentials_to_generate_auth_token() { + let fetch = MockFetch::new(vec![chat_response()]); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + fetch: Some(fetch.transport()), + google_auth_options: Some(GoogleAuthOptions { + credentials: Some(json!({ + "client_email": "test@example.com", + "private_key": "test-key" + })), + ..Default::default() + }), + ..Default::default() + }) + .unwrap(); + + let error = provider + .call("test-model") + .do_generate(&prompt()) + .await + .unwrap_err(); + + assert!(error.to_string().contains("private key"), "{error}"); + assert!(fetch.seen().is_empty()); + } + + /// upstream: @ai-sdk/google-vertex src/maas/edge/google-vertex-maas-provider-edge.test.ts › google-vertex-maas-provider-edge › should merge custom headers with auth header + #[tokio::test] + async fn should_merge_custom_headers_with_auth_header() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&fetch); + let settings = VertexMaasProviderSettings { + headers: Some(Resolvable::Value( + [("X-Custom".to_string(), Some("header-value".to_string()))].into(), + )), + ..settings + }; + + let seen = request(settings, &fetch, &prompt()).await; + + assert_eq!( + sent_headers(&seen, &["user-agent"]), + headers(&[ + ("content-type", "application/json"), + ("x-custom", "header-value"), + ("authorization", "Bearer mock-auth-token"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/maas/edge/google-vertex-maas-provider-edge.test.ts › google-vertex-maas-provider-edge › should use custom fetch when provided + #[tokio::test] + async fn should_use_custom_fetch_when_provided() { + let custom_fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&custom_fetch); + + let seen = request(settings, &custom_fetch, &prompt()).await; + + assert_eq!(seen.headers["authorization"], "Bearer mock-auth-token"); + } + + /// upstream: @ai-sdk/google-vertex src/maas/edge/google-vertex-maas-provider-edge.test.ts › google-vertex-maas-provider-edge › should handle async custom headers + #[tokio::test] + async fn should_handle_async_custom_headers() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&fetch); + let async_headers = Resolvable::from_future(async { + Ok::([("X-Async".to_string(), Some("async-value".to_string()))].into()) + }); + let settings = VertexMaasProviderSettings { + headers: Some(async_headers), + ..settings + }; + + let seen = request(settings, &fetch, &prompt()).await; + + assert_eq!(seen.headers["x-async"], "async-value"); + assert_eq!(seen.headers["authorization"], "Bearer mock-auth-token"); + } + + /// upstream: @ai-sdk/google-vertex src/maas/edge/google-vertex-maas-provider-edge.test.ts › google-vertex-maas-provider-edge › should preserve headers from init parameter + /// + /// `init`'s headers are the per-call headers here; the request also + /// carries the model's `content-type`, and `postToApi` appends the + /// pinned package parts to the per-call `User-Agent`. + #[tokio::test] + async fn should_preserve_headers_from_init_parameter() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&fetch); + let mut options = prompt(); + options.headers = Some( + [("User-Agent", "test-agent"), ("X-Request-ID", "test-id")] + .into_iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(), + ); + + let seen = request(settings, &fetch, &options).await; + + assert_eq!( + sent_headers(&seen, &["content-type", "user-agent"]), + headers(&[ + ("x-request-id", "test-id"), + ("authorization", "Bearer mock-auth-token"), + ]) + ); + assert!( + seen.headers["user-agent"].starts_with("test-agent"), + "{}", + seen.headers["user-agent"] + ); + } + + /// upstream: @ai-sdk/google-vertex src/maas/edge/google-vertex-maas-provider-edge.test.ts › google-vertex-maas-provider-edge › should export default googleVertexMaas instance + #[test] + fn should_export_default_google_vertex_maas_instance() { + let model = google_vertex_maas().call("test-model"); + assert_eq!(model.provider(), "vertex.maas.chat"); + } +} diff --git a/aimux-providers/src/vertex/tests/maas/edge/mod.rs b/aimux-providers/src/vertex/tests/maas/edge/mod.rs new file mode 100644 index 000000000..50cdf3cd7 --- /dev/null +++ b/aimux-providers/src/vertex/tests/maas/edge/mod.rs @@ -0,0 +1 @@ +mod maas_provider_edge; diff --git a/aimux-providers/src/vertex/tests/maas/maas_provider.rs b/aimux-providers/src/vertex/tests/maas/maas_provider.rs new file mode 100644 index 000000000..0c7fd7a07 --- /dev/null +++ b/aimux-providers/src/vertex/tests/maas/maas_provider.rs @@ -0,0 +1,252 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 +//! `src/maas/google-vertex-maas-provider.test.ts`. +//! +//! Upstream mocks `createOpenAICompatible` and inspects the settings handed +//! to it. aimux builds the OpenAI-compatible provider itself, so each setting +//! is observed through a request over a mock transport: `baseURL` through the +//! request URL (`{baseURL}/chat/completions`), `name` through +//! `LanguageModel::provider` (`{name}.chat`), `transformRequestBody` through +//! the request body. The transport is a mock (upstream's `fetch` is +//! `undefined`) and the token comes from the `generate_auth_token` setting +//! (upstream's base provider has no auth). + +use std::sync::Arc; + +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_provider_utils::Resolvable; +use serde_json::json; + +use crate::openai_compatible::OpenAICompatibleProvider; +use crate::test_support::mock_fetch::{Canned, EnvVar, MockFetch, Seen}; +use crate::vertex::{VertexMaasProviderSettings, create_google_vertex_maas}; + +/// A minimal Chat Completions response. +fn chat_response() -> Canned { + Canned::json(&json!({ + "id": "chatcmpl-1", "object": "chat.completion", "created": 1, "model": "m", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + })) +} + +fn prompt() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +/// Settings with the mock transport and a fixed token. +fn settings(fetch: &Arc) -> VertexMaasProviderSettings { + VertexMaasProviderSettings { + fetch: Some(fetch.transport()), + generate_auth_token: Some(Resolvable::Value(Some("test-token".into()))), + ..Default::default() + } +} + +/// Generate once with `model_id` and return the request the transport saw. +async fn request( + provider: &OpenAICompatibleProvider, + fetch: &Arc, + model_id: &str, + options: &CallOptions, +) -> Seen { + let model = provider.call(model_id); + assert_eq!(model.provider(), "vertex.maas.chat"); + model.do_generate(options).await.unwrap(); + fetch.seen().last().unwrap().clone() +} + +mod google_vertex_maas_provider { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider.test.ts › google-vertex-maas-provider › should create a provider with correct base URL for global location + #[tokio::test] + async fn should_create_a_provider_with_correct_base_url_for_global_location() { + let fetch = MockFetch::new(vec![chat_response()]); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + ..settings(&fetch) + }) + .unwrap(); + + let seen = request(&provider, &fetch, "test-model", &prompt()).await; + assert_eq!( + seen.url, + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/endpoints/openapi/chat/completions" + ); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider.test.ts › google-vertex-maas-provider › should create a provider with correct base URL for regional location + #[tokio::test] + async fn should_create_a_provider_with_correct_base_url_for_regional_location() { + let fetch = MockFetch::new(vec![chat_response()]); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + location: Some("us-central1".into()), + ..settings(&fetch) + }) + .unwrap(); + + let seen = request(&provider, &fetch, "test-model", &prompt()).await; + assert_eq!( + seen.url, + "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/endpoints/openapi/chat/completions" + ); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider.test.ts › google-vertex-maas-provider › should create a provider with correct base URL for multi-region location + #[tokio::test] + async fn should_create_a_provider_with_correct_base_url_for_multi_region_location() { + let fetch = MockFetch::new(vec![chat_response()]); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + location: Some("eu".into()), + ..settings(&fetch) + }) + .unwrap(); + + let seen = request(&provider, &fetch, "test-model", &prompt()).await; + assert_eq!( + seen.url, + "https://aiplatform.eu.rep.googleapis.com/v1/projects/test-project/locations/eu/endpoints/openapi/chat/completions" + ); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider.test.ts › google-vertex-maas-provider › should default to global location when not specified + /// + /// Upstream's mocked `loadOptionalSetting` ignores the environment; here + /// `GOOGLE_VERTEX_LOCATION` is unset for the test. + #[tokio::test] + #[serial_test::serial] + async fn should_default_to_global_location_when_not_specified() { + let _location = EnvVar::set("GOOGLE_VERTEX_LOCATION", None); + let fetch = MockFetch::new(vec![chat_response()]); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + ..settings(&fetch) + }) + .unwrap(); + + let seen = request(&provider, &fetch, "test-model", &prompt()).await; + assert_eq!( + seen.url, + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/endpoints/openapi/chat/completions" + ); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider.test.ts › google-vertex-maas-provider › should use custom baseURL when provided + #[tokio::test] + async fn should_use_custom_base_url_when_provided() { + let custom_base_url = "https://custom-endpoint.example.com"; + let fetch = MockFetch::new(vec![chat_response()]); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + base_url: Some(custom_base_url.into()), + ..settings(&fetch) + }) + .unwrap(); + + let seen = request(&provider, &fetch, "test-model", &prompt()).await; + assert_eq!(seen.url, format!("{custom_base_url}/chat/completions")); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider.test.ts › google-vertex-maas-provider › should pass custom fetch to openai-compatible provider + #[tokio::test] + async fn should_pass_custom_fetch_to_openai_compatible_provider() { + let custom_fetch = MockFetch::new(vec![chat_response()]); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + ..settings(&custom_fetch) + }) + .unwrap(); + + request(&provider, &custom_fetch, "test-model", &prompt()).await; + assert_eq!(custom_fetch.seen().len(), 1); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider.test.ts › google-vertex-maas-provider › should default Llama 4 max tokens without overriding explicit or other model settings + /// + /// Upstream calls `transformRequestBody` on three bodies; here each body + /// is the one a request with that model and `max_tokens` builds. + #[tokio::test] + async fn should_default_llama_4_max_tokens_without_overriding_explicit_or_other_settings() { + let fetch = MockFetch::new(vec![chat_response(), chat_response(), chat_response()]); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + ..settings(&fetch) + }) + .unwrap(); + + let seen = request( + &provider, + &fetch, + "meta/llama-4-scout-17b-16e-instruct-maas", + &prompt(), + ) + .await; + let body = seen.json_body(); + assert_eq!(body["model"], "meta/llama-4-scout-17b-16e-instruct-maas"); + assert_eq!(body["max_tokens"], 8192); + + let mut explicit = prompt(); + explicit.max_output_tokens = Some(64); + let seen = request( + &provider, + &fetch, + "meta/llama-4-maverick-17b-128e-instruct-maas", + &explicit, + ) + .await; + let body = seen.json_body(); + assert_eq!( + body["model"], + "meta/llama-4-maverick-17b-128e-instruct-maas" + ); + assert_eq!(body["max_tokens"], 64); + + let seen = request(&provider, &fetch, "openai/gpt-oss-20b-maas", &prompt()).await; + let body = seen.json_body(); + assert_eq!(body["model"], "openai/gpt-oss-20b-maas"); + assert!(body.get("max_tokens").is_none(), "{body}"); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider.test.ts › google-vertex-maas-provider › should construct correct URL with trailing slash removed from baseURL + #[tokio::test] + async fn should_construct_correct_url_with_trailing_slash_removed_from_base_url() { + let fetch = MockFetch::new(vec![chat_response()]); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + base_url: Some("https://custom-endpoint.example.com/".into()), + ..settings(&fetch) + }) + .unwrap(); + + let seen = request(&provider, &fetch, "test-model", &prompt()).await; + assert_eq!( + seen.url, + "https://custom-endpoint.example.com/chat/completions" + ); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider.test.ts › google-vertex-maas-provider › should construct correct URL when baseURL is empty string + #[tokio::test] + async fn should_construct_correct_url_when_base_url_is_empty_string() { + let fetch = MockFetch::new(vec![chat_response()]); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + location: Some("us-central1".into()), + base_url: Some(String::new()), + ..settings(&fetch) + }) + .unwrap(); + + let seen = request(&provider, &fetch, "test-model", &prompt()).await; + assert_eq!( + seen.url, + "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/endpoints/openapi/chat/completions" + ); + } +} diff --git a/aimux-providers/src/vertex/tests/maas/maas_provider_node.rs b/aimux-providers/src/vertex/tests/maas/maas_provider_node.rs new file mode 100644 index 000000000..2999bf21d --- /dev/null +++ b/aimux-providers/src/vertex/tests/maas/maas_provider_node.rs @@ -0,0 +1,348 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 +//! `src/maas/google-vertex-maas-provider-node.test.ts`. +//! +//! Upstream mocks the base provider and calls the fetch wrapper the Node +//! entry point hands it with a hand-made `(url, init)`. aimux has one entry +//! point, whose fetch layer adds the same headers: each case makes a chat +//! request through a mock transport, so the request URL is the provider's +//! (not `https://example.com/test`) and `init` is the request the +//! OpenAI-compatible model builds. Upstream's mocked token generator +//! (returning `mock-auth-token`) is the `generate_auth_token` setting, except +//! in the cases about the default generator, which run aimux's Google auth +//! over the mock transport. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_provider_utils::{HeaderMapOpt, Resolvable}; +use serde_json::json; + +use crate::test_support::mock_fetch::{Canned, EnvVar, MockFetch, RouteFetch, Seen}; +use crate::vertex::{ + GoogleAuthOptions, GoogleAuthScopes, VertexMaasProviderSettings, create_google_vertex_maas, + google_vertex_maas, +}; + +/// A minimal Chat Completions response. +fn chat_response() -> Canned { + Canned::json(&json!({ + "id": "chatcmpl-1", "object": "chat.completion", "created": 1, "model": "m", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + })) +} + +fn prompt() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +/// A token generator returning `mock-auth-token`, and its call counter. +fn mock_generator() -> (Resolvable>, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let generator = Resolvable::from_async_fn(move || { + counter.fetch_add(1, Ordering::SeqCst); + async { Ok(Some("mock-auth-token".to_string())) } + }); + (generator, calls) +} + +/// Settings of a provider for `test-project` over `fetch`, with the mocked +/// token generator. +fn settings(fetch: &Arc) -> (VertexMaasProviderSettings, Arc) { + let (generator, calls) = mock_generator(); + let settings = VertexMaasProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + fetch: Some(fetch.transport()), + generate_auth_token: Some(generator), + ..Default::default() + }; + (settings, calls) +} + +/// Generate once with `options` and return the request the transport saw. +async fn request( + settings: VertexMaasProviderSettings, + fetch: &Arc, + options: &CallOptions, +) -> Seen { + create_google_vertex_maas(settings) + .unwrap() + .call("test-model") + .do_generate(options) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + seen[0].clone() +} + +/// The request headers except `skip`. +fn sent_headers(seen: &Seen, skip: &[&str]) -> BTreeMap { + seen.headers + .iter() + .filter(|(name, _)| !skip.contains(&name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +fn headers(entries: &[(&str, &str)]) -> BTreeMap { + entries + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() +} + +/// No credential file or gcloud configuration and the default metadata +/// host: the default Google auth asks the metadata server for a token. +fn metadata_only_env() -> Vec { + vec![ + EnvVar::set("GOOGLE_APPLICATION_CREDENTIALS", None), + EnvVar::set("google_application_credentials", None), + EnvVar::set( + "CLOUDSDK_CONFIG", + Some("/nonexistent/aimux-vertex-test-gcloud"), + ), + EnvVar::set("GCE_METADATA_IP", None), + EnvVar::set("GCE_METADATA_HOST", None), + ] +} + +const METADATA_TOKEN_PATH: &str = "/computeMetadata/v1/instance/service-accounts/default/token"; + +/// The metadata server's token response. +fn metadata_token(token: &str) -> Canned { + let mut canned = Canned::json(&json!({ + "access_token": token, "expires_in": 3600, "token_type": "Bearer" + })); + canned + .headers + .push(("Metadata-Flavor".into(), "Google".into())); + canned +} + +mod google_vertex_maas_provider_node { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider-node.test.ts › google-vertex-maas-provider-node › should generate auth token and add to request headers + /// + /// `createAuthTokenGenerator(undefined)` has no counterpart once the + /// generator is replaced; the generator's single call is asserted. + #[tokio::test] + async fn should_generate_auth_token_and_add_to_request_headers() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, calls) = settings(&fetch); + + let seen = request(settings, &fetch, &prompt()).await; + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(seen.method, "POST"); + assert_eq!( + sent_headers(&seen, &["user-agent"]), + headers(&[ + ("content-type", "application/json"), + ("authorization", "Bearer mock-auth-token"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider-node.test.ts › google-vertex-maas-provider-node › should pass googleAuthOptions to createAuthTokenGenerator + /// + /// The options reach aimux's Google auth (upstream mocks the generator): + /// with no credentials it asks the metadata server for a token of the + /// given scope, once. + #[tokio::test] + #[serial_test::serial] + async fn should_pass_google_auth_options_to_create_auth_token_generator() { + let _env = metadata_only_env(); + let fetch = RouteFetch::new(|seen| { + if seen.url.contains(METADATA_TOKEN_PATH) { + metadata_token("mock-auth-token") + } else { + chat_response() + } + }); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + fetch: Some(fetch.transport()), + google_auth_options: Some(GoogleAuthOptions { + scopes: Some(GoogleAuthScopes::Array(vec!["test-scope".into()])), + ..Default::default() + }), + ..Default::default() + }) + .unwrap(); + + provider + .call("test-model") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!(fetch.count("GET", METADATA_TOKEN_PATH), 1); + let seen = fetch.seen(); + let metadata = seen + .iter() + .find(|seen| seen.url.contains(METADATA_TOKEN_PATH)) + .unwrap(); + assert!( + metadata.url.ends_with("?scopes=test-scope"), + "{}", + metadata.url + ); + let chat = seen + .iter() + .find(|seen| seen.url.ends_with("/chat/completions")) + .unwrap(); + assert_eq!(chat.headers["authorization"], "Bearer mock-auth-token"); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider-node.test.ts › google-vertex-maas-provider-node › should merge custom headers with auth header + #[tokio::test] + async fn should_merge_custom_headers_with_auth_header() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&fetch); + let settings = VertexMaasProviderSettings { + headers: Some(Resolvable::Value( + [("X-Custom".to_string(), Some("header-value".to_string()))].into(), + )), + ..settings + }; + + let seen = request(settings, &fetch, &prompt()).await; + + assert_eq!( + sent_headers(&seen, &["user-agent"]), + headers(&[ + ("content-type", "application/json"), + ("x-custom", "header-value"), + ("authorization", "Bearer mock-auth-token"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider-node.test.ts › google-vertex-maas-provider-node › should use custom fetch when provided + #[tokio::test] + async fn should_use_custom_fetch_when_provided() { + let custom_fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&custom_fetch); + + let seen = request(settings, &custom_fetch, &prompt()).await; + + assert_eq!(seen.headers["authorization"], "Bearer mock-auth-token"); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider-node.test.ts › google-vertex-maas-provider-node › should handle async custom headers + #[tokio::test] + async fn should_handle_async_custom_headers() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&fetch); + let async_headers = Resolvable::from_future(async { + Ok::([("X-Async".to_string(), Some("async-value".to_string()))].into()) + }); + let settings = VertexMaasProviderSettings { + headers: Some(async_headers), + ..settings + }; + + let seen = request(settings, &fetch, &prompt()).await; + + assert_eq!(seen.headers["x-async"], "async-value"); + assert_eq!(seen.headers["authorization"], "Bearer mock-auth-token"); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider-node.test.ts › google-vertex-maas-provider-node › should preserve headers from init parameter + /// + /// `init`'s headers are the per-call headers here; the request also + /// carries the model's `content-type`, and `postToApi` appends the + /// pinned package parts to the per-call `User-Agent`. + #[tokio::test] + async fn should_preserve_headers_from_init_parameter() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&fetch); + let mut options = prompt(); + options.headers = Some( + [("User-Agent", "test-agent"), ("X-Request-ID", "test-id")] + .into_iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(), + ); + + let seen = request(settings, &fetch, &options).await; + + assert_eq!( + sent_headers(&seen, &["content-type", "user-agent"]), + headers(&[ + ("x-request-id", "test-id"), + ("authorization", "Bearer mock-auth-token"), + ]) + ); + assert!( + seen.headers["user-agent"].starts_with("test-agent"), + "{}", + seen.headers["user-agent"] + ); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider-node.test.ts › google-vertex-maas-provider-node › should export default googleVertexMaas instance + #[test] + fn should_export_default_google_vertex_maas_instance() { + let model = google_vertex_maas().call("test-model"); + assert_eq!(model.provider(), "vertex.maas.chat"); + } + + /// upstream: @ai-sdk/google-vertex src/maas/google-vertex-maas-provider-node.test.ts › google-vertex-maas-provider-node › creates the auth token generator once per provider instance + /// + /// The generator is aimux's Google auth of the provider. It is one + /// instance per provider when its token, cached by the instance, serves + /// both requests: one token exchange for two requests. The cache hides + /// how often the instance itself is asked (upstream: twice). + #[tokio::test] + async fn creates_the_auth_token_generator_once_per_provider_instance() { + let fetch = RouteFetch::new(|seen| { + if seen.url.starts_with("https://oauth2.googleapis.com/token") { + Canned::json(&json!({ + "access_token": "mock-auth-token", "expires_in": 3600, "token_type": "Bearer" + })) + } else { + chat_response() + } + }); + let provider = create_google_vertex_maas(VertexMaasProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + fetch: Some(fetch.transport()), + google_auth_options: Some(GoogleAuthOptions { + credentials: Some(json!({ + "type": "authorized_user", + "client_id": "client-id", + "client_secret": "client-secret", + "refresh_token": "refresh-token" + })), + ..Default::default() + }), + ..Default::default() + }) + .unwrap(); + let model = provider.call("test-model"); + + model.do_generate(&prompt()).await.unwrap(); + model.do_generate(&prompt()).await.unwrap(); + + assert_eq!(fetch.count("POST", "/token"), 1); + let chats: Vec<_> = fetch + .seen() + .into_iter() + .filter(|seen| seen.url.ends_with("/chat/completions")) + .collect(); + assert_eq!(chats.len(), 2); + for seen in chats { + assert_eq!(seen.headers["authorization"], "Bearer mock-auth-token"); + } + } +} diff --git a/aimux-providers/src/vertex/tests/maas/mod.rs b/aimux-providers/src/vertex/tests/maas/mod.rs new file mode 100644 index 000000000..3ace3c016 --- /dev/null +++ b/aimux-providers/src/vertex/tests/maas/mod.rs @@ -0,0 +1,3 @@ +mod edge; +mod maas_provider; +mod maas_provider_node; diff --git a/aimux-providers/src/vertex/tests/mod.rs b/aimux-providers/src/vertex/tests/mod.rs new file mode 100644 index 000000000..bc7d32017 --- /dev/null +++ b/aimux-providers/src/vertex/tests/mod.rs @@ -0,0 +1,17 @@ +//! Tests ported from `@ai-sdk/google-vertex` 5.0.100, one file per upstream test file. + +mod anthropic; +mod auth_google_auth_library; +mod cloud_tts_speech_model; +mod edge; +mod embedding_model; +mod gemini_transcription; +mod image_model; +mod language_model; +mod location_validation; +mod maas; +mod provider; +mod provider_base; +mod transcription_model; +mod video_model; +mod xai; diff --git a/aimux-providers/src/vertex/tests/provider.rs b/aimux-providers/src/vertex/tests/provider.rs new file mode 100644 index 000000000..f442c1589 --- /dev/null +++ b/aimux-providers/src/vertex/tests/provider.rs @@ -0,0 +1,254 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 `src/google-vertex-provider.test.ts`. +//! +//! Upstream's Node provider wraps the base provider: it mocks +//! `createAuthTokenGenerator` and the base `createGoogleVertex` and checks the +//! options the wrapper passes on. aimux has one provider, so the tests look at +//! the same values inside it: the provider headers its resolver produces for a +//! request (upstream's `passedOptions.headers`), the options of the Google +//! credential resolver it built (upstream's `createAuthTokenGenerator` +//! argument), and an `access_token` producer standing in for the mocked token +//! generator. Every test clears `GOOGLE_VERTEX_API_KEY` (upstream's +//! `beforeEach`) and `GOOGLE_VERTEX_ACCESS_TOKEN` (read before the generator), +//! so all of them are serial. + +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use futures::FutureExt; +use serde_json::json; +use serial_test::serial; + +use crate::test_support::mock_fetch::{Canned, EnvVar, MockFetch}; +use crate::vertex::{ + GoogleAuthOptions, GoogleAuthScopes, VertexProvider, VertexProviderSettings, + create_google_vertex, +}; +use aimux_provider_utils::{HeaderMapOpt, Resolvable}; + +const TEST_RSA_KEY: &str = include_str!("fixtures/aimux-test-rsa-key.pem"); + +fn clean_environment() -> Vec { + ["GOOGLE_VERTEX_API_KEY", "GOOGLE_VERTEX_ACCESS_TOKEN"] + .into_iter() + .map(|name| EnvVar::set(name, None)) + .collect() +} + +/// `vi.fn().mockResolvedValue('mock-auth-token')`, counting its calls. +fn mock_auth_token(calls: &Arc) -> Resolvable { + let calls = calls.clone(); + Resolvable::AsyncFn(Arc::new(move || { + calls.fetch_add(1, Ordering::SeqCst); + async { Ok("mock-auth-token".to_string()) }.boxed() + })) +} + +fn create(settings: VertexProviderSettings) -> VertexProvider { + create_google_vertex(settings).expect("valid settings") +} + +/// `await resolve(passedOptions.headers)`: the provider headers of one +/// request, in the mode the provider was created in. +async fn provider_headers(provider: &VertexProvider) -> HeaderMapOpt { + let key = provider.resolver.express_key().await.expect("api key"); + provider + .resolver + .headers(key.as_deref()) + .await + .expect("headers") +} + +/// Expected provider headers. aimux combines header maps with lower-cased +/// names (HTTP names are case-insensitive), so the names are lower-cased. +fn headers(entries: &[(&str, &str)]) -> HeaderMapOpt { + entries + .iter() + .map(|(name, value)| (name.to_ascii_lowercase(), Some((*value).to_string()))) + .collect() +} + +/// The options of the provider's Google credential resolver +/// (`createAuthTokenGenerator`'s argument). +fn auth_options(provider: &VertexProvider) -> &GoogleAuthOptions { + &provider.resolver.auth.options +} + +fn scopes(options: &GoogleAuthOptions) -> Option> { + options.scopes.as_ref().map(|scopes| match scopes { + GoogleAuthScopes::String(scope) => vec![scope.clone()], + GoogleAuthScopes::Array(scopes) => scopes.clone(), + }) +} + +mod google_vertex_provider { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider.test.ts › google-vertex-provider › default headers function should return auth token + #[tokio::test] + #[serial] + async fn default_headers_function_should_return_auth_token() { + let _env = clean_environment(); + let calls = Arc::new(AtomicUsize::new(0)); + let provider = create(VertexProviderSettings { + project: Some("test-project".into()), + access_token: Some(mock_auth_token(&calls)), + ..Default::default() + }); + + assert_eq!( + provider_headers(&provider).await, + headers(&[("Authorization", "Bearer mock-auth-token")]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider.test.ts › google-vertex-provider › should use custom headers in addition to auth token when provided + #[tokio::test] + #[serial] + async fn should_use_custom_headers_in_addition_to_auth_token_when_provided() { + let _env = clean_environment(); + let calls = Arc::new(AtomicUsize::new(0)); + let provider = create(VertexProviderSettings { + project: Some("test-project".into()), + headers: Some(Resolvable::AsyncFn(Arc::new(|| { + async { Ok(headers(&[("Custom-Header", "custom-value")])) }.boxed() + }))), + access_token: Some(mock_auth_token(&calls)), + ..Default::default() + }); + + assert_eq!( + provider_headers(&provider).await, + headers(&[ + ("Authorization", "Bearer mock-auth-token"), + ("Custom-Header", "custom-value"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider.test.ts › google-vertex-provider › passes googleAuthOptions to createAuthTokenGenerator + #[tokio::test] + #[serial] + async fn passes_google_auth_options_to_create_auth_token_generator() { + let _env = clean_environment(); + // aimux builds the credential resolver with the provider, so there + // is no headers call to trigger it (the key file does not exist). + let provider = create(VertexProviderSettings { + google_auth_options: Some(GoogleAuthOptions { + scopes: Some(GoogleAuthScopes::Array(vec![ + "https://www.googleapis.com/auth/cloud-platform".into(), + ])), + key_file: Some("path/to/key.json".into()), + ..Default::default() + }), + ..Default::default() + }); + + let options = auth_options(&provider); + assert_eq!( + scopes(options), + Some(vec![ + "https://www.googleapis.com/auth/cloud-platform".to_string() + ]) + ); + assert_eq!(options.key_file, Some(PathBuf::from("path/to/key.json"))); + assert_eq!(options.project_id, None); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider.test.ts › google-vertex-provider › passes project to createAuthTokenGenerator as projectId + #[tokio::test] + #[serial] + async fn passes_project_to_create_auth_token_generator_as_project_id() { + let _env = clean_environment(); + let provider = create(VertexProviderSettings { + project: Some("test-project".into()), + ..Default::default() + }); + + let options = auth_options(&provider); + assert_eq!(options.project_id.as_deref(), Some("test-project")); + assert_eq!(scopes(options), None); + assert_eq!(options.key_file, None); + assert!(options.credentials.is_none()); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider.test.ts › google-vertex-provider › does not override explicit googleAuthOptions projectId + #[tokio::test] + #[serial] + async fn does_not_override_explicit_google_auth_options_project_id() { + let _env = clean_environment(); + let provider = create(VertexProviderSettings { + project: Some("provider-project".into()), + google_auth_options: Some(GoogleAuthOptions { + project_id: Some("auth-project".into()), + key_file: Some("path/to/key.json".into()), + ..Default::default() + }), + ..Default::default() + }); + + let options = auth_options(&provider); + assert_eq!(options.project_id.as_deref(), Some("auth-project")); + assert_eq!(options.key_file, Some(PathBuf::from("path/to/key.json"))); + assert_eq!(scopes(options), None); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider.test.ts › google-vertex-provider › should pass options through to base provider when apiKey is provided + #[tokio::test] + #[serial] + async fn should_pass_options_through_to_base_provider_when_api_key_is_provided() { + let _env = clean_environment(); + let calls = Arc::new(AtomicUsize::new(0)); + let provider = create(VertexProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".into())), + access_token: Some(mock_auth_token(&calls)), + ..Default::default() + }); + + assert_eq!( + provider.resolver.express_key().await.unwrap().as_deref(), + Some("test-api-key") + ); + // No auth-token headers: only the Express-mode key. + let headers = provider_headers(&provider).await; + assert!( + !headers + .keys() + .any(|name| name.eq_ignore_ascii_case("authorization")), + "{headers:?}" + ); + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider.test.ts › google-vertex-provider › creates the auth token generator once per provider instance + #[tokio::test] + #[serial] + async fn creates_the_auth_token_generator_once_per_provider_instance() { + let _env = clean_environment(); + // One generator per provider: its service-account token is fetched + // once and reused by the second headers call. + let fetch = MockFetch::new(vec![Canned::json( + &json!({"access_token": "sa-token", "expires_in": 3600}), + )]); + let provider = create(VertexProviderSettings { + project: Some("test-project".into()), + google_auth_options: Some(GoogleAuthOptions { + credentials: Some(json!({ + "type": "service_account", + "client_email": "test@test.iam.gserviceaccount.com", + "private_key": TEST_RSA_KEY, + })), + ..Default::default() + }), + fetch: Some(fetch.transport()), + ..Default::default() + }); + + let first = provider_headers(&provider).await; + let second = provider_headers(&provider).await; + + assert_eq!(first, headers(&[("Authorization", "Bearer sa-token")])); + assert_eq!(second, first); + assert_eq!(fetch.seen().len(), 1); + } +} diff --git a/aimux-providers/src/vertex/tests/provider_base.rs b/aimux-providers/src/vertex/tests/provider_base.rs new file mode 100644 index 000000000..9d3754ead --- /dev/null +++ b/aimux-providers/src/vertex/tests/provider_base.rs @@ -0,0 +1,646 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 `src/google-vertex-provider-base.test.ts`. +//! +//! Upstream replaces every model class with a `vi.fn()` and checks the +//! constructor arguments (`provider`, `baseURL`, `headers`, `generateId`). +//! aimux builds the model configuration privately, so these tests observe the +//! same values on the model and on the request it sends: `provider()` is the +//! `provider` argument, the request URL starts with `baseURL`, the +//! `Authorization` header of an injected access token shows the `headers` +//! function, and a tool-call id shows a custom `generateId`. Upstream's +//! `beforeEach` deletes `GOOGLE_VERTEX_API_KEY`; every test does the same +//! through an [`EnvVar`] guard, so all of them are serial. + +use std::sync::Arc; + +use serde_json::json; +use serial_test::serial; + +use crate::test_support::mock_fetch::{Canned, EnvVar, MockFetch, Seen}; +use crate::vertex::{VertexProvider, VertexProviderSettings, create_google_vertex}; +use aimux_core::AiMuxError; +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::image_model::{ImageCallOptions, ImageModel}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_core::provider::Provider; +use aimux_core::result::GenerateContent; +use aimux_core::speech_model::{SpeechCallOptions, SpeechModel}; +use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions}; +use aimux_core::video_model::{VideoCallOptions, VideoModel}; +use aimux_provider_utils::Resolvable; + +/// `https://test-location-aiplatform.googleapis.com/…/publishers/google`. +const TEST_LOCATION_BASE_URL: &str = "https://test-location-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/test-location/publishers/google"; +const GLOBAL_BASE_URL: &str = "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/publishers/google"; +const CUSTOM_BASE_URL: &str = "https://custom-endpoint.example.com"; + +/// The upstream `beforeEach`/`afterEach`: no `GOOGLE_VERTEX_API_KEY`. +fn clean_environment() -> EnvVar { + EnvVar::set("GOOGLE_VERTEX_API_KEY", None) +} + +/// `createGoogleVertex({ project, location, ...})` with a fixed access token +/// (upstream's base provider takes the auth headers from its caller) and the +/// mock transport. +fn settings(location: &str, fetch: &Arc) -> VertexProviderSettings { + VertexProviderSettings { + project: Some("test-project".into()), + location: Some(location.into()), + access_token: Some("test-token".to_string().into()), + fetch: Some(fetch.transport()), + ..Default::default() + } +} + +/// `createGoogleVertex({ apiKey: 'test-api-key' })` with the mock transport. +fn express_settings(fetch: &Arc) -> VertexProviderSettings { + VertexProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".into())), + fetch: Some(fetch.transport()), + ..Default::default() + } +} + +fn provider(settings: VertexProviderSettings) -> VertexProvider { + create_google_vertex(settings).expect("valid settings") +} + +fn prompt() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +/// A Gemini response whose only part is a function call without an id, so +/// the model takes the id from `generateId`. +fn function_call_response() -> Canned { + Canned::json(&json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{"functionCall": {"name": "weather", "args": {"city": "Paris"}}}] + }, + "finishReason": "STOP" + }], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2} + })) +} + +/// The first request the transport saw. +fn only_request(fetch: &MockFetch) -> Seen { + let seen = fetch.seen(); + assert!(!seen.is_empty(), "no request was sent"); + seen[0].clone() +} + +/// The request of a `doGenerate` call of a language model. +async fn chat_request(model: &dyn LanguageModel, fetch: &MockFetch) -> Seen { + let _ = model.do_generate(&prompt()).await; + only_request(fetch) +} + +/// `expect.objectContaining({ baseURL })`: the request goes below `base_url`. +fn assert_base_url(seen: &Seen, base_url: &str) { + assert!( + seen.url.starts_with(&format!("{base_url}/")), + "{} is not below {base_url}", + seen.url + ); +} + +/// `headers: expect.any(Function)`: the provider headers reached the request. +fn assert_provider_headers(seen: &Seen) { + assert_eq!( + seen.headers.get("authorization").map(String::as_str), + Some("Bearer test-token"), + "{:?}", + seen.headers + ); +} + +mod google_vertex_provider_base { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should create a language model with default settings + #[tokio::test] + #[serial] + async fn should_create_a_language_model_with_default_settings() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("test-location", &fetch)).call("test-model-id"); + + assert_eq!(model.model_id(), "test-model-id"); + assert_eq!(model.provider(), "google.vertex.chat"); + let seen = chat_request(model.as_ref(), &fetch).await; + assert_eq!( + seen.url, + format!("{TEST_LOCATION_BASE_URL}/models/test-model-id:generateContent") + ); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should create an embedding model with correct settings + #[tokio::test] + #[serial] + async fn should_create_an_embedding_model_with_correct_settings() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("test-location", &fetch)).embedding("test-embedding-model"); + + assert_eq!(model.model_id(), "test-embedding-model"); + assert_eq!(model.provider(), "google.vertex.embedding"); + let _ = model.do_embed(&EmbeddingCallOptions::new("Hello")).await; + let seen = only_request(&fetch); + assert!( + seen.url.starts_with(&format!( + "{TEST_LOCATION_BASE_URL}/models/test-embedding-model:" + )), + "{}", + seen.url + ); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should create a speech model with correct settings + #[tokio::test] + #[serial] + async fn should_create_a_speech_model_with_correct_settings() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("test-location", &fetch)) + .speech("gemini-2.5-flash-tts") + .expect("a Gemini speech model"); + + assert_eq!(model.model_id(), "gemini-2.5-flash-tts"); + assert_eq!(model.provider(), "google.vertex.speech"); + let _ = model.do_generate(&SpeechCallOptions::new("Hello")).await; + let seen = only_request(&fetch); + assert_base_url(&seen, TEST_LOCATION_BASE_URL); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should create a speech model via speechModel() + #[tokio::test] + #[serial] + async fn should_create_a_speech_model_via_speech_model() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("test-location", &fetch)) + .speech_model("gemini-2.5-pro-tts") + .expect("the provider has speech models") + .expect("a Gemini speech model"); + + assert_eq!(model.model_id(), "gemini-2.5-pro-tts"); + assert_eq!(model.provider(), "google.vertex.speech"); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should keep routing gemini speech models to the Gemini speech model + #[tokio::test] + #[serial] + async fn should_keep_routing_gemini_speech_models_to_the_gemini_speech_model() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("test-location", &fetch)) + .speech_model("gemini-2.5-flash-tts") + .expect("the provider has speech models") + .expect("a Gemini speech model"); + + assert_eq!(model.model_id(), "gemini-2.5-flash-tts"); + assert_eq!(model.provider(), "google.vertex.speech"); + // Not the Cloud Text-to-Speech model: the request is a Gemini + // `generateContent` call below the Gemini publisher URL. + let _ = model.do_generate(&SpeechCallOptions::new("Hello")).await; + let seen = only_request(&fetch); + assert_eq!( + seen.url, + format!("{TEST_LOCATION_BASE_URL}/models/gemini-2.5-flash-tts:generateContent") + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should create a transcription model with correct settings + #[tokio::test] + #[serial] + async fn should_create_a_transcription_model_with_correct_settings() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("us-central1", &fetch)).transcription("chirp_2"); + + assert_eq!(model.model_id(), "chirp_2"); + assert_eq!(model.provider(), "google.vertex.transcription"); + let options = TranscriptionCallOptions::new(AudioInput::Binary(vec![0, 1, 2]), "audio/wav"); + let _ = model.do_generate(&options).await; + let seen = only_request(&fetch); + // `project: 'test-project'`, `location: 'us-central1'`. + assert!( + seen.url + .contains("/projects/test-project/locations/us-central1/"), + "{}", + seen.url + ); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should create a transcription model via transcriptionModel() + #[tokio::test] + #[serial] + async fn should_create_a_transcription_model_via_transcription_model() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("us-central1", &fetch)) + .transcription_model("chirp_3") + .expect("the provider has transcription models") + .expect("a transcription model"); + + assert_eq!(model.model_id(), "chirp_3"); + assert_eq!(model.provider(), "google.vertex.transcription"); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should reject Express Mode for transcription models + #[tokio::test] + #[serial] + async fn should_reject_express_mode_for_transcription_models() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(express_settings(&fetch)).transcription("chirp_3"); + + // Upstream throws when the model is created; aimux fails its first + // request, before anything is sent. + let options = TranscriptionCallOptions::new(AudioInput::Binary(vec![0, 1, 2]), "audio/wav"); + let error = model.do_generate(&options).await.unwrap_err(); + assert!( + error.to_string().contains( + "Google Vertex transcription models do not support Express Mode API keys. Use standard Google Cloud credentials instead." + ), + "{error}" + ); + assert!(fetch.seen().is_empty()); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should pass custom headers to the model constructor + #[tokio::test] + #[serial] + async fn should_pass_custom_headers_to_the_model_constructor() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(VertexProviderSettings { + headers: Some(Resolvable::Value( + [( + "Custom-Header".to_string(), + Some("custom-value".to_string()), + )] + .into_iter() + .collect(), + )), + ..settings("test-location", &fetch) + }) + .call("test-model-id"); + + let seen = chat_request(model.as_ref(), &fetch).await; + assert_eq!( + seen.headers.get("custom-header").map(String::as_str), + Some("custom-value") + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should pass custom generateId function to the model constructor + #[tokio::test] + #[serial] + async fn should_pass_custom_generate_id_function_to_the_model_constructor() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![function_call_response()]); + let model = provider(VertexProviderSettings { + generate_id: Some(Arc::new(|| "custom-id".to_string())), + ..settings("test-location", &fetch) + }) + .call("test-model-id"); + + let result = model.do_generate(&prompt()).await.expect("generated"); + let ids: Vec<&str> = result + .content + .iter() + .filter_map(|content| match content { + GenerateContent::ToolCall(call) => Some(call.tool_call_id.as_str()), + _ => None, + }) + .collect(); + assert_eq!(ids, ["custom-id"]); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use languageModel method to create a model + #[tokio::test] + #[serial] + async fn should_use_language_model_method_to_create_a_model() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("test-location", &fetch)) + .language_model("test-model-id") + .expect("a language model"); + + assert_eq!(model.model_id(), "test-model-id"); + assert_eq!(model.provider(), "google.vertex.chat"); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use custom baseURL when provided + #[tokio::test] + #[serial] + async fn should_use_custom_base_url_when_provided() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(VertexProviderSettings { + base_url: Some(CUSTOM_BASE_URL.into()), + ..settings("test-location", &fetch) + }) + .call("test-model-id"); + + let seen = chat_request(model.as_ref(), &fetch).await; + assert_base_url(&seen, CUSTOM_BASE_URL); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use a tuned model endpoints/ id with no publishers/google suffix + #[tokio::test] + #[serial] + async fn should_use_a_tuned_model_endpoints_id_with_no_publishers_google_suffix() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("test-location", &fetch)).call("endpoints/1234567890"); + + assert_eq!(model.model_id(), "endpoints/1234567890"); + assert_eq!(model.provider(), "google.vertex.chat"); + let seen = chat_request(model.as_ref(), &fetch).await; + assert_eq!( + seen.url, + "https://test-location-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/test-location/endpoints/1234567890:generateContent" + ); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should leave base model ids and the publishers/google base URL unchanged + #[tokio::test] + #[serial] + async fn should_leave_base_model_ids_and_the_publishers_google_base_url_unchanged() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("test-location", &fetch)).call("gemini-2.5-pro"); + + assert_eq!(model.model_id(), "gemini-2.5-pro"); + let seen = chat_request(model.as_ref(), &fetch).await; + assert_eq!( + seen.url, + format!("{TEST_LOCATION_BASE_URL}/models/gemini-2.5-pro:generateContent") + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should respect a custom baseURL for tuned model endpoints/ ids + #[tokio::test] + #[serial] + async fn should_respect_a_custom_base_url_for_tuned_model_endpoints_ids() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(VertexProviderSettings { + base_url: Some(CUSTOM_BASE_URL.into()), + ..settings("test-location", &fetch) + }) + .call("endpoints/1234567890"); + + assert_eq!(model.model_id(), "endpoints/1234567890"); + let seen = chat_request(model.as_ref(), &fetch).await; + assert_eq!( + seen.url, + format!("{CUSTOM_BASE_URL}/endpoints/1234567890:generateContent") + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should reject Express Mode for tuned models + #[tokio::test] + #[serial] + async fn should_reject_express_mode_for_tuned_models() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(express_settings(&fetch)).call("endpoints/1234567890"); + + // Upstream throws when the model is created; aimux fails its first + // request, before anything is sent. + let error = model.do_generate(&prompt()).await.unwrap_err(); + assert!( + matches!(&error, AiMuxError::InvalidArgument { .. }), + "{error:?}" + ); + assert!( + error.to_string().contains( + "Google Vertex tuned models do not support Express Mode API keys. Use standard Google Cloud credentials instead." + ), + "{error}" + ); + assert!(fetch.seen().is_empty()); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should create an image model with default settings + #[tokio::test] + #[serial] + async fn should_create_an_image_model_with_default_settings() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("test-location", &fetch)).image("gemini-2.5-flash-image"); + + assert_eq!(model.model_id(), "gemini-2.5-flash-image"); + assert_eq!(model.provider(), "google.vertex.image"); + let _ = model.do_generate(&ImageCallOptions::new("A cat")).await; + let seen = only_request(&fetch); + assert_base_url(&seen, TEST_LOCATION_BASE_URL); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use correct URL for global region + #[tokio::test] + #[serial] + async fn should_use_correct_url_for_global_region() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("global", &fetch)).call("test-model-id"); + + assert_eq!(model.provider(), "google.vertex.chat"); + let seen = chat_request(model.as_ref(), &fetch).await; + assert_eq!( + seen.url, + format!("{GLOBAL_BASE_URL}/models/test-model-id:generateContent") + ); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use correct URL for global region with embedding model + #[tokio::test] + #[serial] + async fn should_use_correct_url_for_global_region_with_embedding_model() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("global", &fetch)).embedding("test-embedding-model"); + + assert_eq!(model.provider(), "google.vertex.embedding"); + let _ = model.do_embed(&EmbeddingCallOptions::new("Hello")).await; + let seen = only_request(&fetch); + assert!( + seen.url + .starts_with(&format!("{GLOBAL_BASE_URL}/models/test-embedding-model:")), + "{}", + seen.url + ); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use correct URL for global region with image model + #[tokio::test] + #[serial] + async fn should_use_correct_url_for_global_region_with_image_model() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("global", &fetch)).image("gemini-2.5-flash-image"); + + assert_eq!(model.provider(), "google.vertex.image"); + let _ = model.do_generate(&ImageCallOptions::new("A cat")).await; + let seen = only_request(&fetch); + assert_base_url(&seen, GLOBAL_BASE_URL); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use region-prefixed URL for non-global regions + #[tokio::test] + #[serial] + async fn should_use_region_prefixed_url_for_non_global_regions() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("us-central1", &fetch)).call("test-model-id"); + + assert_eq!(model.provider(), "google.vertex.chat"); + let seen = chat_request(model.as_ref(), &fetch).await; + assert_base_url( + &seen, + "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google", + ); + assert_provider_headers(&seen); + } + + /// The parts of the REP inline snapshots aimux exposes: `baseURL`, + /// `provider`, `headers` and `supportedUrls`. + async fn assert_multi_region(location: &str, base_url: &str) { + let fetch = MockFetch::new(vec![]); + let model = provider(settings(location, &fetch)).call("test-model-id"); + + assert_eq!(model.provider(), "google.vertex.chat"); + assert!(!model.supported_urls().await.is_empty()); + let seen = chat_request(model.as_ref(), &fetch).await; + assert_base_url(&seen, base_url); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use multi-region REP URL for us location + #[tokio::test] + #[serial] + async fn should_use_multi_region_rep_url_for_us_location() { + let _env = clean_environment(); + assert_multi_region( + "us", + "https://aiplatform.us.rep.googleapis.com/v1beta1/projects/test-project/locations/us/publishers/google", + ) + .await; + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use multi-region REP URL for eu location + #[tokio::test] + #[serial] + async fn should_use_multi_region_rep_url_for_eu_location() { + let _env = clean_environment(); + assert_multi_region( + "eu", + "https://aiplatform.eu.rep.googleapis.com/v1beta1/projects/test-project/locations/eu/publishers/google", + ) + .await; + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use express mode base URL when apiKey is provided + #[tokio::test] + #[serial] + async fn should_use_express_mode_base_url_when_api_key_is_provided() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(express_settings(&fetch)).call("test-model-id"); + + let seen = chat_request(model.as_ref(), &fetch).await; + assert_base_url( + &seen, + "https://aiplatform.googleapis.com/v1/publishers/google", + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should add API key as query parameter via custom fetch + #[tokio::test] + #[serial] + async fn should_add_api_key_as_query_parameter_via_custom_fetch() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + // Upstream calls the provider's fetch wrapper directly with a + // streaming URL; aimux's wrapper is private, so a streaming request + // of the same model goes through it. + let model = provider(express_settings(&fetch)).call("gemini-2.5-pro"); + + let _ = model.do_stream(&prompt()).await; + let seen = only_request(&fetch); + let (url, query) = seen.url.split_once('?').unwrap_or((&seen.url, "")); + assert_eq!( + url, + "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-2.5-pro:streamGenerateContent" + ); + assert!(!query.contains("test-api-key"), "{}", seen.url); + assert_eq!( + seen.headers.get("x-goog-api-key").map(String::as_str), + Some("test-api-key") + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should create a video model with default settings + #[tokio::test] + #[serial] + async fn should_create_a_video_model_with_default_settings() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("test-location", &fetch)).video("veo-2.0-generate-001"); + + assert_eq!(model.model_id(), "veo-2.0-generate-001"); + assert_eq!(model.provider(), "google.vertex.video"); + let _ = model.do_start(&VideoCallOptions::new("A cat")).await; + let seen = only_request(&fetch); + assert_base_url(&seen, TEST_LOCATION_BASE_URL); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use correct URL for global region with video model + #[tokio::test] + #[serial] + async fn should_use_correct_url_for_global_region_with_video_model() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(settings("global", &fetch)).video("veo-3.0-generate-001"); + + assert_eq!(model.model_id(), "veo-3.0-generate-001"); + assert_eq!(model.provider(), "google.vertex.video"); + let _ = model.do_start(&VideoCallOptions::new("A cat")).await; + let seen = only_request(&fetch); + assert_base_url(&seen, GLOBAL_BASE_URL); + assert_provider_headers(&seen); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-provider-base.test.ts › google-vertex-provider-base › should use custom baseURL for video model when provided + #[tokio::test] + #[serial] + async fn should_use_custom_base_url_for_video_model_when_provided() { + let _env = clean_environment(); + let fetch = MockFetch::new(vec![]); + let model = provider(VertexProviderSettings { + base_url: Some(CUSTOM_BASE_URL.into()), + ..settings("test-location", &fetch) + }) + .video("veo-2.0-generate-001"); + + let _ = model.do_start(&VideoCallOptions::new("A cat")).await; + let seen = only_request(&fetch); + assert_base_url(&seen, CUSTOM_BASE_URL); + } +} diff --git a/aimux-providers/src/vertex/tests/snapshots/embedding_full_result.json b/aimux-providers/src/vertex/tests/snapshots/embedding_full_result.json new file mode 100644 index 000000000..ce5c54c2e --- /dev/null +++ b/aimux-providers/src/vertex/tests/snapshots/embedding_full_result.json @@ -0,0 +1,65 @@ +{ + "embeddings": [ + [ + -0.017999587580561638, + -0.006893285550177097, + -0.036766719073057175, + -0.017558680847287178, + -0.019938766956329346 + ], + [ + -0.06007182598114014, + 0.004907649010419846, + -0.00690646655857563, + -0.007314121350646019, + -0.048464205116033554 + ] + ], + "response": { + "body": { + "metadata": { + "billableCharacterCount": 35 + }, + "predictions": [ + { + "embeddings": { + "statistics": { + "token_count": 5, + "truncated": false + }, + "values": [ + -0.017999587580561638, + -0.006893285550177097, + -0.036766719073057175, + -0.017558680847287178, + -0.019938766956329346 + ] + } + }, + { + "embeddings": { + "statistics": { + "truncated": false, + "token_count": 6 + }, + "values": [ + -0.06007182598114014, + 0.004907649010419846, + -0.00690646655857563, + -0.007314121350646019, + -0.048464205116033554 + ] + } + } + ] + }, + "headers": { + "content-length": "429", + "content-type": "application/json" + } + }, + "usage": { + "tokens": 11 + }, + "warnings": [] +} diff --git a/aimux-providers/src/vertex/tests/snapshots/embedding_raw_response.json b/aimux-providers/src/vertex/tests/snapshots/embedding_raw_response.json new file mode 100644 index 000000000..64310b8e2 --- /dev/null +++ b/aimux-providers/src/vertex/tests/snapshots/embedding_raw_response.json @@ -0,0 +1,44 @@ +{ + "body": { + "metadata": { + "billableCharacterCount": 35 + }, + "predictions": [ + { + "embeddings": { + "statistics": { + "token_count": 5, + "truncated": false + }, + "values": [ + -0.017999587580561638, + -0.006893285550177097, + -0.036766719073057175, + -0.017558680847287178, + -0.019938766956329346 + ] + } + }, + { + "embeddings": { + "statistics": { + "truncated": false, + "token_count": 6 + }, + "values": [ + -0.06007182598114014, + 0.004907649010419846, + -0.00690646655857563, + -0.007314121350646019, + -0.048464205116033554 + ] + } + } + ] + }, + "headers": { + "content-length": "429", + "content-type": "application/json", + "test-header": "test-value" + } +} diff --git a/aimux-providers/src/vertex/tests/transcription_model.rs b/aimux-providers/src/vertex/tests/transcription_model.rs new file mode 100644 index 000000000..20ca6e778 --- /dev/null +++ b/aimux-providers/src/vertex/tests/transcription_model.rs @@ -0,0 +1,514 @@ +use std::sync::Arc; + +use serde_json::{Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::shared::SharedProviderOptions; +use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; +use aimux_provider_utils::{HeaderMapOpt, Resolvable}; + +use crate::shared::{Endpoint, EndpointConfig, ResolveEndpoint}; +use crate::test_support::mock_fetch::{Canned, MockFetch, RouteFetch}; +use crate::vertex::{ + VertexProvider, VertexProviderSettings, VertexTranscriptionModel, create_google_vertex, +}; + +/// 8 bytes of audio; base64 -> 'AQIDBAUGBwg='. +const AUDIO_DATA: [u8; 8] = [1, 2, 3, 4, 5, 6, 7, 8]; +const AUDIO_BASE64: &str = "AQIDBAUGBwg="; + +const URL: &str = "https://us-central1-speech.googleapis.com/v2/projects/test-project/locations/us-central1/recognizers/_:recognize"; +const US_URL: &str = "https://us-speech.googleapis.com/v2/projects/test-project/locations/us/recognizers/_:recognize"; +const GLOBAL_URL: &str = "https://speech.googleapis.com/v2/projects/test-project/locations/global/recognizers/_:recognize"; + +/// The pinned user-agent part of the package (upstream mocks `0.0.0-test`). +const USER_AGENT_PART: &str = "ai-sdk-google-vertex/5.0.100"; + +fn default_body() -> Value { + json!({ + "results": [{ + "alternatives": [{ + "transcript": "hello world", + "words": [ + {"word": "hello", "startOffset": "0s", "endOffset": "0.500s"}, + {"word": "world", "startOffset": "0.500s", "endOffset": "1s"}, + ], + }], + "languageCode": "en-US", + }], + "metadata": {"totalBilledDuration": "1s"}, + }) +} + +fn json_response(body: &Value, headers: &[(&str, &str)]) -> Canned { + let body = serde_json::to_vec(body).unwrap(); + let mut all = vec![ + ("content-type".to_string(), "application/json".to_string()), + ("content-length".to_string(), body.len().to_string()), + ]; + all.extend(headers.iter().map(|(k, v)| (k.to_string(), v.to_string()))); + Canned { + status: 200, + headers: all, + body, + } +} + +/// The AI SDK test server: `url` answers `response`, anything else 404. +fn test_server(url: &'static str, response: Canned) -> Arc { + RouteFetch::new(move |seen| { + if seen.url == url { + response.clone() + } else { + Canned { + status: 404, + headers: Vec::new(), + body: b"Not found".to_vec(), + } + } + }) +} + +/// `prepareJsonResponse({ headers })`. +fn prepare_json_response(headers: &[(&str, &str)]) -> Arc { + test_server(URL, json_response(&default_body(), headers)) +} + +/// `createGoogleVertex({ project: 'test-project', location: 'us-central1' })`. +/// aimux has one entry point, and it authenticates: an access token stands in +/// for upstream's credential-less base provider. +fn provider_with(fetch: &Arc, headers: Option) -> VertexProvider { + create_google_vertex(VertexProviderSettings { + project: Some("test-project".into()), + location: Some("us-central1".into()), + access_token: Some(Resolvable::Value("test-token".into())), + headers: headers.map(Resolvable::Value), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} + +/// `provider.transcription('chirp_2')`. +fn model(fetch: &Arc) -> Arc { + provider_with(fetch, None).transcription("chirp_2") +} + +fn options(audio: AudioInput, provider_options: Option) -> TranscriptionCallOptions { + let mut options = TranscriptionCallOptions::new(audio, "audio/wav"); + options.provider_options = provider_options + .map(|value| serde_json::from_value::(value).unwrap()); + options +} + +fn audio() -> AudioInput { + AudioInput::Binary(AUDIO_DATA.to_vec()) +} + +/// A base URL and headers fixed for every request (the upstream model config). +struct StaticEndpoint { + base_url: String, + headers: HeaderMapOpt, +} + +#[async_trait::async_trait] +impl ResolveEndpoint for StaticEndpoint { + async fn resolve(&self) -> Result { + Ok(Endpoint { + base_url: self.base_url.clone(), + headers: self.headers.clone(), + }) + } +} + +mod do_generate { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should send the model, languageCodes, features and base64 content + #[tokio::test] + async fn should_send_the_model_language_codes_features_and_base64_content() { + let server = prepare_json_response(&[]); + model(&server) + .do_generate(&options(audio(), None)) + .await + .unwrap(); + assert_eq!( + server.seen()[0].json_body(), + json!({ + "config": { + "model": "chirp_2", + "languageCodes": ["auto"], + "autoDecodingConfig": {}, + "features": { + "enableWordTimeOffsets": true, + "enableAutomaticPunctuation": true, + }, + }, + "content": AUDIO_BASE64, + }) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should target the regional Speech-to-Text endpoint + #[tokio::test] + async fn should_target_the_regional_speech_to_text_endpoint() { + let server = prepare_json_response(&[]); + model(&server) + .do_generate(&options(audio(), None)) + .await + .unwrap(); + assert_eq!(server.seen()[0].url, URL); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should extract text, segments, language and duration + #[tokio::test] + async fn should_extract_text_segments_language_and_duration() { + let server = prepare_json_response(&[]); + let result = model(&server) + .do_generate(&options(audio(), None)) + .await + .unwrap(); + assert_eq!(result.text, "hello world"); + assert_eq!( + serde_json::to_value(&result.segments).unwrap(), + json!([ + {"text": "hello", "startSecond": 0.0, "endSecond": 0.5}, + {"text": "world", "startSecond": 0.5, "endSecond": 1.0}, + ]) + ); + // BCP-47 `en-US` is reduced to the ISO-639-1 `en`. + assert_eq!(result.language.as_deref(), Some("en")); + assert_eq!(result.duration_in_seconds, Some(1.0)); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should accept a base64 string audio input directly + #[tokio::test] + async fn should_accept_a_base64_string_audio_input_directly() { + let server = prepare_json_response(&[]); + model(&server) + .do_generate(&options(AudioInput::Base64(AUDIO_BASE64.into()), None)) + .await + .unwrap(); + assert_eq!(server.seen()[0].json_body()["content"], AUDIO_BASE64); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should pass languageCodes and features from provider options + #[tokio::test] + async fn should_pass_language_codes_and_features_from_provider_options() { + let server = prepare_json_response(&[]); + model(&server) + .do_generate(&options( + audio(), + Some(json!({ + "googleVertex": { + "languageCodes": ["en-US", "fr-FR"], + "enableAutomaticPunctuation": false, + "enableWordTimeOffsets": false, + }, + })), + )) + .await + .unwrap(); + // `toMatchObject` + let body = server.seen()[0].json_body(); + assert_eq!(body["config"]["languageCodes"], json!(["en-US", "fr-FR"])); + assert_eq!( + body["config"]["features"]["enableAutomaticPunctuation"], + json!(false) + ); + assert_eq!( + body["config"]["features"]["enableWordTimeOffsets"], + json!(false) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should accept the legacy vertex provider options key + #[tokio::test] + async fn should_accept_the_legacy_vertex_provider_options_key() { + let server = test_server(US_URL, json_response(&default_body(), &[])); + model(&server) + .do_generate(&options(audio(), Some(json!({"vertex": {"region": "us"}})))) + .await + .unwrap(); + assert_eq!(server.seen()[0].url, US_URL); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should accept the google provider options key + #[tokio::test] + async fn should_accept_the_google_provider_options_key() { + let server = test_server(US_URL, json_response(&default_body(), &[])); + model(&server) + .do_generate(&options(audio(), Some(json!({"google": {"region": "us"}})))) + .await + .unwrap(); + assert_eq!(server.seen()[0].url, US_URL); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should prefer googleVertex provider options + #[tokio::test] + async fn should_prefer_google_vertex_provider_options() { + let server = test_server(US_URL, json_response(&default_body(), &[])); + model(&server) + .do_generate(&options( + audio(), + Some(json!({ + "googleVertex": {"region": "us"}, + "vertex": {"region": "eu"}, + "google": {"region": "asia-southeast1"}, + })), + )) + .await + .unwrap(); + assert_eq!(server.seen()[0].url, US_URL); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should route to a region override from provider options + #[tokio::test] + async fn should_route_to_a_region_override_from_provider_options() { + let server = test_server(US_URL, json_response(&default_body(), &[])); + model(&server) + .do_generate(&options( + audio(), + Some(json!({"googleVertex": {"region": "us"}})), + )) + .await + .unwrap(); + assert_eq!(server.seen()[0].url, US_URL); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should use the unprefixed host for the global location + #[tokio::test] + async fn should_use_the_unprefixed_host_for_the_global_location() { + let server = test_server(GLOBAL_URL, json_response(&default_body(), &[])); + provider_with(&server, None) + .transcription("latest_long") + .do_generate(&options( + audio(), + Some(json!({"googleVertex": {"region": "global"}})), + )) + .await + .unwrap(); + assert_eq!(server.seen()[0].url, GLOBAL_URL); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should pass headers and user agent + #[tokio::test] + async fn should_pass_headers_and_user_agent() { + let server = prepare_json_response(&[]); + let provider_with_headers = provider_with( + &server, + Some( + [( + "Custom-Provider-Header".to_string(), + Some("provider-header-value".to_string()), + )] + .into(), + ), + ); + let mut call = options(audio(), None); + call.headers = Some( + [( + "Custom-Request-Header".to_string(), + "request-header-value".to_string(), + )] + .into(), + ); + provider_with_headers + .transcription("chirp_2") + .do_generate(&call) + .await + .unwrap(); + let seen = server.seen(); + let headers = &seen[0].headers; + assert_eq!( + headers.get("custom-provider-header").map(String::as_str), + Some("provider-header-value") + ); + assert_eq!( + headers.get("custom-request-header").map(String::as_str), + Some("request-header-value") + ); + let user_agent = headers.get("user-agent").cloned().unwrap_or_default(); + assert!(user_agent.contains(USER_AGENT_PART), "{user_agent}"); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should return empty text and segments when there are no results + #[tokio::test] + async fn should_return_empty_text_and_segments_when_there_are_no_results() { + let server = test_server(URL, json_response(&json!({}), &[])); + let result = model(&server) + .do_generate(&options(audio(), None)) + .await + .unwrap(); + assert_eq!(result.text, ""); + assert!(result.segments.is_empty()); + assert_eq!(result.language, None); + assert_eq!(result.duration_in_seconds, None); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should omit segments without word time offsets + #[tokio::test] + async fn should_omit_segments_without_word_time_offsets() { + let server = test_server( + URL, + json_response( + &json!({ + "results": [{ + "alternatives": [{ + "transcript": "hello world", + "words": [{"word": "hello"}, {"word": "world"}], + }], + "languageCode": "en-US", + }], + }), + &[], + ), + ); + let result = model(&server) + .do_generate(&options(audio(), None)) + .await + .unwrap(); + assert!(result.segments.is_empty()); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should canonicalize detected languages to ISO-639-1 + #[tokio::test] + async fn should_canonicalize_detected_languages_to_iso_639_1() { + let server = test_server( + URL, + json_response( + &json!({ + "results": [{ + "alternatives": [{"transcript": "你好"}], + "languageCode": "cmn-Hans-CN", + }], + }), + &[], + ), + ); + let result = model(&server) + .do_generate(&options(audio(), None)) + .await + .unwrap(); + assert_eq!(result.language.as_deref(), Some("zh")); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should omit detected languages without an ISO-639-1 code + #[tokio::test] + async fn should_omit_detected_languages_without_an_iso_639_1_code() { + let server = test_server( + URL, + json_response( + &json!({ + "results": [{ + "alternatives": [{"transcript": "hello world"}], + "languageCode": "ast-ES", + }], + }), + &[], + ), + ); + let result = model(&server) + .do_generate(&options(audio(), None)) + .await + .unwrap(); + assert_eq!(result.language, None); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should parse billed duration + #[tokio::test] + async fn should_parse_billed_duration() { + let server = test_server( + URL, + json_response( + &json!({ + "results": [{ + "alternatives": [{"transcript": "hello world"}], + "languageCode": "en-US", + }], + "metadata": {"totalBilledDuration": "1.250s"}, + }), + &[], + ), + ); + let result = model(&server) + .do_generate(&options(audio(), None)) + .await + .unwrap(); + assert_eq!(result.duration_in_seconds, Some(1.25)); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › doGenerate › should include response data with timestamp, modelId and headers + #[tokio::test] + async fn should_include_response_data_with_timestamp_model_id_and_headers() { + let server = prepare_json_response(&[("x-request-id", "test-request-id")]); + // `new GoogleVertexTranscriptionModel('chirp_2', { provider, project, + // location, headers: () => ({}) })`: the project and location come + // from a provider with those settings. + let custom_model = VertexTranscriptionModel::from_config( + "chirp_2".to_string(), + EndpointConfig::new( + "google.vertex.transcription".to_string(), + StaticEndpoint { + base_url: "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google".to_string(), + headers: HeaderMapOpt::new(), + }, + Some(server.transport()), + ), + provider_with(&server, None).project_location("transcription"), + ); + // aimux cannot inject `currentDate`: the timestamp is asserted to be + // set, not to equal a fixed date. + let before = aimux_core::Timestamp::now(); + let result = custom_model + .do_generate(&options(audio(), None)) + .await + .unwrap(); + let timestamp = result.response.timestamp.expect("timestamp"); + assert!(timestamp.unix_millis() >= before.unix_millis()); + assert_eq!(result.response.model_id.as_deref(), Some("chirp_2")); + assert_eq!( + result + .response + .headers + .as_ref() + .and_then(|headers| headers.get("x-request-id")) + .map(String::as_str), + Some("test-request-id") + ); + } +} + +mod region_validation { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-transcription-model.test.ts › region validation › rejects an injected per-request region under %s before fetching + #[tokio::test] + async fn rejects_an_injected_per_request_region_before_fetching() { + for namespace in ["googleVertex", "vertex", "google"] { + let fetch = MockFetch::new(Vec::new()); + let model = create_google_vertex(VertexProviderSettings { + project: Some("test".into()), + location: Some("us-central1".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() + .transcription("chirp_3"); + let error = model + .do_generate(&options( + audio(), + Some(json!({namespace: {"region": "user@internal:8080/#"}})), + )) + .await + .unwrap_err(); + // `AI_InvalidArgumentError` with `argument: 'region'`: aimux's + // invalid-argument error names the argument in its message. + assert!( + matches!(&error, AiMuxError::InvalidArgument { message, .. } if message.contains("region")), + "{namespace}: {error:?}" + ); + assert!(fetch.seen().is_empty(), "{namespace}: fetch was called"); + } + } +} diff --git a/aimux-providers/src/vertex/tests/video_model.rs b/aimux-providers/src/vertex/tests/video_model.rs new file mode 100644 index 000000000..82c99d550 --- /dev/null +++ b/aimux-providers/src/vertex/tests/video_model.rs @@ -0,0 +1,527 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use serde_json::{Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::shared::{AspectRatio, SharedProviderOptions, Size}; +use aimux_core::video_model::{ + VideoCallOptions, VideoFile, VideoFileData, VideoModel, VideoOperationStatus, +}; +use aimux_provider_utils::HeaderMapOpt; + +use crate::shared::{Endpoint, EndpointConfig, ResolveEndpoint}; +use crate::test_support::mock_fetch::{Canned, RouteFetch, Seen}; +use crate::vertex::VertexVideoModel; + +const PROMPT: &str = "A futuristic city with flying cars"; + +/// A base URL and headers fixed for every request (the upstream model config). +struct StaticEndpoint { + base_url: String, + headers: HeaderMapOpt, +} + +#[async_trait::async_trait] +impl ResolveEndpoint for StaticEndpoint { + async fn resolve(&self) -> Result { + Ok(Endpoint { + base_url: self.base_url.clone(), + headers: self.headers.clone(), + }) + } +} + +fn json_response(body: &Value) -> Canned { + Canned::json(body) +} + +fn not_found() -> Canned { + Canned { + status: 404, + headers: Vec::new(), + body: b"Not found".to_vec(), + } +} + +/// `new GoogleVertexVideoModel(modelId, { provider: 'google-vertex', baseURL: 'https://api.example.com', headers: () => ({ 'api-key': 'test-key' }), fetch })`. +fn model_with_fetch(model_id: &str, fetch: &Arc) -> VertexVideoModel { + VertexVideoModel::from_config( + model_id.to_string(), + EndpointConfig::new( + "google-vertex".to_string(), + StaticEndpoint { + base_url: "https://api.example.com".to_string(), + headers: [("api-key".to_string(), Some("test-key".to_string()))].into(), + }, + Some(fetch.transport()), + ), + ) +} + +/// The parameters of upstream's `createMockModel`. +struct MockModel { + model_id: &'static str, + operation_name: &'static str, + operation_done: bool, + operation_error: Option, + videos: Vec, + polls_until_done: usize, +} + +impl Default for MockModel { + fn default() -> Self { + Self { + model_id: "veo-2.0-generate-001", + operation_name: "operations/test-operation-id", + operation_done: true, + operation_error: None, + videos: vec![ + json!({"bytesBase64Encoded": "base64-video-data", "mimeType": "video/mp4"}), + ], + polls_until_done: 1, + } + } +} + +/// `createMockModel(...)`: the model and its recording fetch. +fn create_mock_model(mock: MockModel) -> (VertexVideoModel, Arc) { + let poll_count = AtomicUsize::new(0); + let MockModel { + model_id, + operation_name, + operation_done, + operation_error, + videos, + polls_until_done, + } = mock; + let fetch = RouteFetch::new(move |seen: &Seen| { + if seen.url.contains(":predictLongRunning") { + return json_response(&json!({"name": operation_name, "done": false})); + } + if seen.url.contains(":fetchPredictOperation") { + let count = poll_count.fetch_add(1, Ordering::SeqCst) + 1; + if count < polls_until_done { + return json_response(&json!({"name": operation_name, "done": false})); + } + let mut body = json!({"name": operation_name, "done": operation_done}); + match &operation_error { + Some(error) => body["error"] = error.clone(), + None => body["response"] = json!({"videos": videos}), + } + return json_response(&body); + } + not_found() + }); + (model_with_fetch(model_id, &fetch), fetch) +} + +/// `defaultOptions`. +fn default_options() -> VideoCallOptions { + VideoCallOptions::new(PROMPT) +} + +fn provider_options(value: Value) -> SharedProviderOptions { + serde_json::from_value(value).unwrap() +} + +/// The `:predictLongRunning` body the fetch saw. +fn captured_body(fetch: &RouteFetch) -> Value { + fetch + .seen() + .iter() + .find(|seen| seen.url.contains(":predictLongRunning")) + .expect("a predictLongRunning request") + .json_body() +} + +async fn start_with(options: VideoCallOptions) -> Value { + let (model, fetch) = create_mock_model(MockModel::default()); + model.do_start(&options).await.unwrap(); + captured_body(&fetch) +} + +/// The message of an error, without the variant's display prefix. +fn message(error: &AiMuxError) -> String { + match error { + AiMuxError::InvalidArgument { message, .. } + | AiMuxError::InvalidResponseData { message, .. } => message.clone(), + other => other.to_string(), + } +} + +mod google_vertex_video_model { + use super::*; + + mod constructor { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › constructor › should expose correct provider and model information + #[test] + fn should_expose_correct_provider_and_model_information() { + let (model, _) = create_mock_model(MockModel::default()); + assert_eq!(model.provider(), "google-vertex"); + assert_eq!(model.model_id(), "veo-2.0-generate-001"); + assert_eq!(model.specification_version(), "v4"); + assert_eq!(model.max_videos_per_call(), Some(4)); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › constructor › should support different model IDs + #[test] + fn should_support_different_model_ids() { + let (model, _) = create_mock_model(MockModel { + model_id: "veo-3.0-generate-001", + ..Default::default() + }); + assert_eq!(model.model_id(), "veo-3.0-generate-001"); + } + } + + mod do_start { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › should return operation with operationName + #[tokio::test] + async fn should_return_operation_with_operation_name() { + let (model, _) = create_mock_model(MockModel { + operation_name: "operations/my-op-123", + ..Default::default() + }); + let result = model.do_start(&default_options()).await.unwrap(); + assert_eq!( + result.operation, + json!({"operationName": "operations/my-op-123"}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › should pass correct request body + #[tokio::test] + async fn should_pass_correct_request_body() { + assert_eq!( + start_with(default_options()).await, + json!({"instances": [{"prompt": PROMPT}], "parameters": {"sampleCount": 1}}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › should pass seed when provided + #[tokio::test] + async fn should_pass_seed_when_provided() { + let mut options = default_options(); + options.seed = Some(42); + assert_eq!( + start_with(options).await, + json!({"instances": [{"prompt": PROMPT}], "parameters": {"sampleCount": 1, "seed": 42}}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › should pass aspect ratio when provided + #[tokio::test] + async fn should_pass_aspect_ratio_when_provided() { + let mut options = default_options(); + options.aspect_ratio = Some(AspectRatio::new(16, 9)); + assert_eq!( + start_with(options).await, + json!({"instances": [{"prompt": PROMPT}], "parameters": {"sampleCount": 1, "aspectRatio": "16:9"}}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › should convert resolution to Vertex format + #[tokio::test] + async fn should_convert_resolution_to_vertex_format() { + let mut options = default_options(); + options.resolution = Some(Size::new(1920, 1080)); + assert_eq!( + start_with(options).await, + json!({"instances": [{"prompt": PROMPT}], "parameters": {"sampleCount": 1, "resolution": "1080p"}}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › should pass duration as durationSeconds + #[tokio::test] + async fn should_pass_duration_as_duration_seconds() { + let mut options = default_options(); + options.duration = Some(5); + assert_eq!( + start_with(options).await, + json!({"instances": [{"prompt": PROMPT}], "parameters": {"sampleCount": 1, "durationSeconds": 5}}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › should pass n as sampleCount + #[tokio::test] + async fn should_pass_n_as_sample_count() { + let mut options = default_options(); + options.n = 2; + assert_eq!( + start_with(options).await, + json!({"instances": [{"prompt": PROMPT}], "parameters": {"sampleCount": 2}}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › should return warnings array + #[tokio::test] + async fn should_return_warnings_array() { + let (model, _) = create_mock_model(MockModel::default()); + let result = model.do_start(&default_options()).await.unwrap(); + assert!(result.warnings.is_empty()); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › should throw when no operation name returned + #[tokio::test] + async fn should_throw_when_no_operation_name_returned() { + let fetch = RouteFetch::new(|_| json_response(&json!({"done": false}))); + let model = model_with_fetch("veo-2.0-generate-001", &fetch); + let error = model.do_start(&default_options()).await.unwrap_err(); + assert_eq!(message(&error), "No operation name returned from API"); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › should include warnings and response metadata + #[tokio::test] + async fn should_include_warnings_and_response_metadata() { + // aimux cannot inject `currentDate`: the timestamp is asserted to + // be set, not to equal a fixed date. + let before = aimux_core::Timestamp::now(); + let (model, _) = create_mock_model(MockModel::default()); + let result = model.do_start(&default_options()).await.unwrap(); + assert!(result.warnings.is_empty()); + let timestamp = result.response.timestamp.expect("timestamp"); + assert!(timestamp.unix_millis() >= before.unix_millis()); + assert_eq!( + result.response.model_id.as_deref(), + Some("veo-2.0-generate-001") + ); + } + + mod image_to_video { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › Image-to-Video › should send image as bytesBase64Encoded + #[tokio::test] + async fn should_send_image_as_bytes_base64_encoded() { + let mut options = default_options(); + options.image = Some(VideoFile::File { + media_type: "image/png".into(), + data: VideoFileData::Base64("base64-image-data".into()), + }); + let body = start_with(options).await; + // `toMatchObject` + assert_eq!( + body["instances"], + json!([{ + "prompt": PROMPT, + "image": {"bytesBase64Encoded": "base64-image-data", "mimeType": "image/png"}, + }]) + ); + assert_eq!(body["parameters"]["sampleCount"], json!(1)); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › Image-to-Video › should warn when URL-based image is provided + #[tokio::test] + async fn should_warn_when_url_based_image_is_provided() { + let (model, _) = create_mock_model(MockModel::default()); + let mut options = default_options(); + options.image = Some(VideoFile::Url { + url: "https://example.com/image.png".into(), + media_type: None, + }); + let result = model.do_start(&options).await.unwrap(); + assert_eq!(result.warnings.len(), 1); + let warning = serde_json::to_value(&result.warnings[0]).unwrap(); + assert_eq!(warning["type"], "unsupported"); + assert_eq!(warning["feature"], "URL-based image input"); + } + } + + mod provider_options { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › Provider Options › should pass personGeneration option + #[tokio::test] + async fn should_pass_person_generation_option() { + let mut options = default_options(); + options.provider_options = + provider_options(json!({"vertex": {"personGeneration": "allow_adult"}})); + assert_eq!( + start_with(options).await, + json!({ + "instances": [{"prompt": PROMPT}], + "parameters": {"sampleCount": 1, "personGeneration": "allow_adult"}, + }) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › Provider Options › should pass negativePrompt option + #[tokio::test] + async fn should_pass_negative_prompt_option() { + let mut options = default_options(); + options.provider_options = + provider_options(json!({"vertex": {"negativePrompt": "blurry, low quality"}})); + assert_eq!( + start_with(options).await, + json!({ + "instances": [{"prompt": PROMPT}], + "parameters": {"sampleCount": 1, "negativePrompt": "blurry, low quality"}, + }) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › Provider Options › should pass generateAudio option + #[tokio::test] + async fn should_pass_generate_audio_option() { + let mut options = default_options(); + options.provider_options = + provider_options(json!({"vertex": {"generateAudio": true}})); + assert_eq!( + start_with(options).await, + json!({ + "instances": [{"prompt": PROMPT}], + "parameters": {"sampleCount": 1, "generateAudio": true}, + }) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › Provider Options › should pass gcsOutputDirectory option + #[tokio::test] + async fn should_pass_gcs_output_directory_option() { + let mut options = default_options(); + options.provider_options = provider_options( + json!({"vertex": {"gcsOutputDirectory": "gs://bucket/output/"}}), + ); + assert_eq!( + start_with(options).await, + json!({ + "instances": [{"prompt": PROMPT}], + "parameters": {"sampleCount": 1, "gcsOutputDirectory": "gs://bucket/output/"}, + }) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStart › Provider Options › should pass referenceImages option + #[tokio::test] + async fn should_pass_reference_images_option() { + let mut options = default_options(); + options.provider_options = provider_options(json!({ + "vertex": { + "referenceImages": [ + {"bytesBase64Encoded": "reference-image-data"}, + {"gcsUri": "gs://bucket/reference.png"}, + ], + }, + })); + let body = start_with(options).await; + assert_eq!( + body["instances"][0]["referenceImages"], + json!([ + {"bytesBase64Encoded": "reference-image-data"}, + {"gcsUri": "gs://bucket/reference.png"}, + ]) + ); + } + } + } + + mod do_status { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStatus › should return completed with video data when done + #[tokio::test] + async fn should_return_completed_with_video_data_when_done() { + let (model, _) = create_mock_model(MockModel { + operation_name: "operations/my-op-123", + ..Default::default() + }); + let result = model + .do_status( + &json!({"operationName": "operations/my-op-123"}), + &default_options(), + ) + .await + .unwrap(); + let VideoOperationStatus::Completed(result) = result else { + panic!("expected status completed, got pending"); + }; + assert_eq!( + serde_json::to_value(result.videos).unwrap(), + json!([{"type": "base64", "data": "base64-video-data", "mediaType": "video/mp4"}]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStatus › should return pending when not done + #[tokio::test] + async fn should_return_pending_when_not_done() { + let (model, _) = create_mock_model(MockModel { + operation_name: "operations/my-op-123", + polls_until_done: 3, + ..Default::default() + }); + let result = model + .do_status( + &json!({"operationName": "operations/my-op-123"}), + &default_options(), + ) + .await + .unwrap(); + assert!(matches!(result, VideoOperationStatus::Pending)); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStatus › should use POST with operationName in body + #[tokio::test] + async fn should_use_post_with_operation_name_in_body() { + let fetch = RouteFetch::new(|seen: &Seen| { + if seen.url.contains(":predictLongRunning") { + return json_response(&json!({"name": "operations/check-post", "done": false})); + } + if seen.url.contains(":fetchPredictOperation") { + return json_response(&json!({ + "name": "operations/check-post", + "done": true, + "response": {"videos": [{"bytesBase64Encoded": "video-data", "mimeType": "video/mp4"}]}, + })); + } + not_found() + }); + let model = model_with_fetch("veo-2.0-generate-001", &fetch); + model + .do_status( + &json!({"operationName": "operations/check-post"}), + &default_options(), + ) + .await + .unwrap(); + let seen = fetch.seen(); + let poll = seen + .iter() + .find(|seen| seen.url.contains(":fetchPredictOperation")) + .expect("captured url contains :fetchPredictOperation"); + assert_eq!(poll.method, "POST"); + assert_eq!( + poll.json_body(), + json!({"operationName": "operations/check-post"}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/google-vertex-video-model.test.ts › GoogleVertexVideoModel › doStatus › should return completed with GCS URI video + #[tokio::test] + async fn should_return_completed_with_gcs_uri_video() { + let (model, _) = create_mock_model(MockModel { + operation_name: "operations/gcs-op", + videos: vec![json!({"gcsUri": "gs://bucket/video.mp4", "mimeType": "video/mp4"})], + ..Default::default() + }); + let result = model + .do_status( + &json!({"operationName": "operations/gcs-op"}), + &default_options(), + ) + .await + .unwrap(); + let VideoOperationStatus::Completed(result) = result else { + panic!("expected status completed, got pending"); + }; + assert_eq!( + serde_json::to_value(result.videos).unwrap(), + json!([{"type": "url", "url": "gs://bucket/video.mp4", "mediaType": "video/mp4"}]) + ); + } + } +} diff --git a/aimux-providers/src/vertex/tests/xai/edge/mod.rs b/aimux-providers/src/vertex/tests/xai/edge/mod.rs new file mode 100644 index 000000000..f8a39f7a7 --- /dev/null +++ b/aimux-providers/src/vertex/tests/xai/edge/mod.rs @@ -0,0 +1 @@ +mod xai_provider_edge; diff --git a/aimux-providers/src/vertex/tests/xai/edge/xai_provider_edge.rs b/aimux-providers/src/vertex/tests/xai/edge/xai_provider_edge.rs new file mode 100644 index 000000000..c6ec55df5 --- /dev/null +++ b/aimux-providers/src/vertex/tests/xai/edge/xai_provider_edge.rs @@ -0,0 +1,202 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 +//! `src/xai/edge/google-vertex-xai-provider-edge.test.ts`. +//! +//! aimux has one entry point and no separate Web-Crypto edge auth; these +//! cases exercise the same behaviour through it. Upstream mocks the base +//! provider and calls the fetch wrapper the edge entry point hands it with a +//! hand-made `(url, init)`; here each case makes a chat request through a +//! mock transport, so the request URL is the provider's (not +//! `https://example.com/test`) and `init` is the request the +//! OpenAI-compatible model builds. Upstream's mocked `generateAuthToken` +//! (returning `mock-auth-token`) is the `generate_auth_token` setting, except +//! for `googleCredentials` (aimux: `GoogleAuthOptions::credentials`), which +//! runs aimux's Google auth. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_provider_utils::Resolvable; +use serde_json::json; + +use crate::test_support::mock_fetch::{Canned, MockFetch, Seen}; +use crate::vertex::{ + GoogleAuthOptions, VertexXaiProviderSettings, create_google_vertex_xai, google_vertex_xai, +}; + +/// A minimal Chat Completions response. +fn chat_response() -> Canned { + Canned::json(&json!({ + "id": "chatcmpl-1", "object": "chat.completion", "created": 1, "model": "m", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + })) +} + +fn prompt() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +/// A token generator returning `mock-auth-token`, and its call counter. +fn mock_generator() -> (Resolvable>, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let generator = Resolvable::from_async_fn(move || { + counter.fetch_add(1, Ordering::SeqCst); + async { Ok(Some("mock-auth-token".to_string())) } + }); + (generator, calls) +} + +/// Settings of a provider for `test-project` over `fetch`, with the mocked +/// token generator. +fn settings(fetch: &Arc) -> (VertexXaiProviderSettings, Arc) { + let (generator, calls) = mock_generator(); + let settings = VertexXaiProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + fetch: Some(fetch.transport()), + generate_auth_token: Some(generator), + ..Default::default() + }; + (settings, calls) +} + +/// Generate once with `options` and return the request the transport saw. +async fn request( + settings: VertexXaiProviderSettings, + fetch: &Arc, + options: &CallOptions, +) -> Seen { + create_google_vertex_xai(settings) + .unwrap() + .call("xai/grok-4.1-fast-reasoning") + .do_generate(options) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + seen[0].clone() +} + +/// The request headers except `skip`. +fn sent_headers(seen: &Seen, skip: &[&str]) -> BTreeMap { + seen.headers + .iter() + .filter(|(name, _)| !skip.contains(&name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +fn headers(entries: &[(&str, &str)]) -> BTreeMap { + entries + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() +} + +mod google_vertex_xai_provider_edge { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/xai/edge/google-vertex-xai-provider-edge.test.ts › google-vertex-xai-provider-edge › should generate auth token and add to request headers + /// + /// The `undefined` credentials argument has no counterpart once the + /// generator is replaced; the generator's single call is asserted. + #[tokio::test] + async fn should_generate_auth_token_and_add_to_request_headers() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, calls) = settings(&fetch); + + let seen = request(settings, &fetch, &prompt()).await; + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(seen.method, "POST"); + assert_eq!( + sent_headers(&seen, &["user-agent"]), + headers(&[ + ("content-type", "application/json"), + ("authorization", "Bearer mock-auth-token"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/xai/edge/google-vertex-xai-provider-edge.test.ts › google-vertex-xai-provider-edge › should pass googleCredentials to generateAuthToken + /// + /// The credentials reach aimux's real Google auth (upstream mocks the + /// generator), which signs with the given private key: `test-key` is not + /// a key, so the request fails on it before anything is sent. + #[tokio::test] + async fn should_pass_google_credentials_to_generate_auth_token() { + let fetch = MockFetch::new(vec![chat_response()]); + let provider = create_google_vertex_xai(VertexXaiProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + fetch: Some(fetch.transport()), + google_auth_options: Some(GoogleAuthOptions { + credentials: Some(json!({ + "client_email": "test@example.com", + "private_key": "test-key" + })), + ..Default::default() + }), + ..Default::default() + }) + .unwrap(); + + let error = provider + .call("xai/grok-4.1-fast-reasoning") + .do_generate(&prompt()) + .await + .unwrap_err(); + + assert!(error.to_string().contains("private key"), "{error}"); + assert!(fetch.seen().is_empty()); + } + + /// upstream: @ai-sdk/google-vertex src/xai/edge/google-vertex-xai-provider-edge.test.ts › google-vertex-xai-provider-edge › should merge custom headers with auth header + #[tokio::test] + async fn should_merge_custom_headers_with_auth_header() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&fetch); + let settings = VertexXaiProviderSettings { + headers: Some(Resolvable::Value( + [("X-Custom".to_string(), Some("header-value".to_string()))].into(), + )), + ..settings + }; + + let seen = request(settings, &fetch, &prompt()).await; + + assert_eq!( + sent_headers(&seen, &["user-agent"]), + headers(&[ + ("content-type", "application/json"), + ("x-custom", "header-value"), + ("authorization", "Bearer mock-auth-token"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/xai/edge/google-vertex-xai-provider-edge.test.ts › google-vertex-xai-provider-edge › should use custom fetch when provided + #[tokio::test] + async fn should_use_custom_fetch_when_provided() { + let custom_fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&custom_fetch); + + let seen = request(settings, &custom_fetch, &prompt()).await; + + assert_eq!( + sent_headers(&seen, &["content-type", "user-agent"]), + headers(&[("authorization", "Bearer mock-auth-token")]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/xai/edge/google-vertex-xai-provider-edge.test.ts › google-vertex-xai-provider-edge › should export default googleVertexXai instance + #[test] + fn should_export_default_google_vertex_xai_instance() { + let model = google_vertex_xai().call("xai/grok-4.1-fast-reasoning"); + assert_eq!(model.provider(), "googleVertex.xai.chat"); + } +} diff --git a/aimux-providers/src/vertex/tests/xai/mod.rs b/aimux-providers/src/vertex/tests/xai/mod.rs new file mode 100644 index 000000000..98b552ae7 --- /dev/null +++ b/aimux-providers/src/vertex/tests/xai/mod.rs @@ -0,0 +1,3 @@ +mod edge; +mod xai_provider; +mod xai_provider_node; diff --git a/aimux-providers/src/vertex/tests/xai/xai_provider.rs b/aimux-providers/src/vertex/tests/xai/xai_provider.rs new file mode 100644 index 000000000..e9e885efa --- /dev/null +++ b/aimux-providers/src/vertex/tests/xai/xai_provider.rs @@ -0,0 +1,282 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 +//! `src/xai/google-vertex-xai-provider.test.ts`. +//! +//! Upstream mocks `createOpenAICompatible` and inspects the settings handed +//! to it, calling its `transformRequestBody`, `convertUsage` and +//! `supportedUrls` directly. aimux builds the OpenAI-compatible provider +//! itself, so each setting is observed through the model: `baseURL` through +//! the request URL (`{baseURL}/chat/completions`), `name` through +//! `LanguageModel::provider` (`{name}.chat`), `includeUsage` and +//! `supportsStructuredOutputs` through the request bodies, +//! `transformRequestBody` through the request body, `convertUsage` through +//! the result's usage, `supportedUrls` through `LanguageModel::supported_urls`. +//! The transport is a mock (upstream's `fetch` is `undefined`) and the token +//! comes from the `generate_auth_token` setting (upstream's base provider has +//! no auth). + +use std::sync::Arc; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::{CallOptions, ResponseFormat}; +use aimux_core::provider::Provider; +use aimux_core::types::ReasoningEffort; +use aimux_provider_utils::Resolvable; +use serde_json::{Value, json}; + +use crate::test_support::mock_fetch::{Canned, MockFetch, Seen}; +use crate::vertex::{VertexXaiProvider, VertexXaiProviderSettings, create_google_vertex_xai}; + +/// A Chat Completions response with `usage`. +fn chat_response(usage: Value) -> Canned { + Canned::json(&json!({ + "id": "chatcmpl-1", "object": "chat.completion", "created": 1, "model": "m", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": usage + })) +} + +fn ok_response() -> Canned { + chat_response(json!({"prompt_tokens": 1, "completion_tokens": 1})) +} + +/// A Chat Completions stream with one empty chunk. +fn stream_response() -> Canned { + let chunk = json!({ + "id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}] + }); + Canned { + status: 200, + headers: vec![("content-type".into(), "text/event-stream".into())], + body: format!("data: {chunk}\n\ndata: [DONE]\n\n").into_bytes(), + } +} + +fn prompt() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +/// Settings of `test-project` with the mock transport and a fixed token. +fn settings(fetch: &Arc) -> VertexXaiProviderSettings { + VertexXaiProviderSettings { + project: Some("test-project".into()), + fetch: Some(fetch.transport()), + generate_auth_token: Some(Resolvable::Value(Some("test-token".into()))), + ..Default::default() + } +} + +/// The provider of `settings` in `location` (with `base_url`). +fn provider( + fetch: &Arc, + location: Option<&str>, + base_url: Option<&str>, +) -> VertexXaiProvider { + create_google_vertex_xai(VertexXaiProviderSettings { + location: location.map(String::from), + base_url: base_url.map(String::from), + ..settings(fetch) + }) + .unwrap() +} + +/// Generate once with `options` and return the request the transport saw. +async fn request( + provider: &VertexXaiProvider, + fetch: &Arc, + options: &CallOptions, +) -> Seen { + let model = provider.call("xai/grok-4.1-fast-reasoning"); + assert_eq!(model.provider(), "googleVertex.xai.chat"); + model.do_generate(options).await.unwrap(); + fetch.seen().last().unwrap().clone() +} + +mod google_vertex_xai_provider { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider.test.ts › google-vertex-xai-provider › should create a provider with correct base URL for global location + #[tokio::test] + async fn should_create_a_provider_with_correct_base_url_for_global_location() { + let fetch = MockFetch::new(vec![ok_response()]); + let provider = provider(&fetch, Some("global"), None); + + let seen = request(&provider, &fetch, &prompt()).await; + assert_eq!( + seen.url, + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/endpoints/openapi/chat/completions" + ); + } + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider.test.ts › google-vertex-xai-provider › should construct the default base URL when baseURL is an empty string + #[tokio::test] + async fn should_construct_the_default_base_url_when_base_url_is_an_empty_string() { + let fetch = MockFetch::new(vec![ok_response()]); + let provider = provider(&fetch, Some("global"), Some("")); + + let seen = request(&provider, &fetch, &prompt()).await; + assert_eq!( + seen.url, + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/endpoints/openapi/chat/completions" + ); + } + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider.test.ts › google-vertex-xai-provider › should create a provider with correct base URL for regional location + #[tokio::test] + async fn should_create_a_provider_with_correct_base_url_for_regional_location() { + let fetch = MockFetch::new(vec![ok_response()]); + let provider = provider(&fetch, Some("us-central1"), None); + + let seen = request(&provider, &fetch, &prompt()).await; + assert_eq!( + seen.url, + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/endpoints/openapi/chat/completions" + ); + } + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider.test.ts › google-vertex-xai-provider › should configure OpenAI-compatible behavior for Vertex Grok + #[tokio::test] + async fn should_configure_openai_compatible_behavior_for_vertex_grok() { + let fetch = MockFetch::new(vec![ok_response(), stream_response()]); + let provider = provider(&fetch, Some("global"), None); + let model = provider.call("xai/grok-4.1-fast-reasoning"); + assert_eq!(model.provider(), "googleVertex.xai.chat"); + + // supportsStructuredOutputs: a JSON schema goes out as `json_schema`. + let mut options = prompt(); + options.response_format = Some(ResponseFormat::Json { + schema: Some(json!({"type": "object", "properties": {"a": {"type": "string"}}})), + name: None, + description: None, + }); + model.do_generate(&options).await.unwrap(); + // includeUsage: streams ask for usage. + let _stream = model.do_stream(&prompt()).await.unwrap(); + + let seen = fetch.seen(); + assert_eq!(seen.len(), 2); + let base_url = "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/endpoints/openapi"; + for seen in &seen { + assert_eq!(seen.url, format!("{base_url}/chat/completions")); + } + assert_eq!( + seen[0].json_body()["response_format"]["type"], + "json_schema" + ); + assert_eq!( + seen[1].json_body()["stream_options"], + json!({"include_usage": true}) + ); + } + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider.test.ts › google-vertex-xai-provider › should strip reasoning_effort from request bodies + /// + /// Upstream calls `transformRequestBody` on a body with + /// `reasoning_effort: 'high'`; here the body is the one a request with + /// `reasoning: high` builds. + #[tokio::test] + async fn should_strip_reasoning_effort_from_request_bodies() { + let fetch = MockFetch::new(vec![ok_response()]); + let provider = provider(&fetch, Some("global"), None); + let mut options = prompt(); + options.reasoning = Some(ReasoningEffort::High); + + let body = request(&provider, &fetch, &options).await.json_body(); + + assert!(body.get("reasoning_effort").is_none(), "{body}"); + assert_eq!(body["model"], "xai/grok-4.1-fast-reasoning"); + assert!(body["messages"].is_array(), "{body}"); + } + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider.test.ts › google-vertex-xai-provider › should count Grok reasoning tokens separately from completion tokens + /// + /// Upstream calls `convertUsage` on the usage object; here it is the + /// usage of a response. + #[tokio::test] + async fn should_count_grok_reasoning_tokens_separately_from_completion_tokens() { + let usage = json!({ + "prompt_tokens": 663, + "prompt_tokens_details": {"cached_tokens": 654}, + "completion_tokens": 50, + "completion_tokens_details": {"reasoning_tokens": 124} + }); + let fetch = MockFetch::new(vec![chat_response(usage)]); + let provider = provider(&fetch, Some("global"), None); + + let result = provider + .call("xai/grok-4.1-fast-reasoning") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!( + serde_json::to_value(&result.usage).unwrap(), + json!({ + "inputTokens": { + "cacheRead": 654, + "noCache": 9, + "total": 663 + }, + "outputTokens": { + "reasoning": 124, + "text": 50, + "total": 174 + }, + "raw": { + "completion_tokens": 50, + "completion_tokens_details": { + "reasoning_tokens": 124 + }, + "prompt_tokens": 663, + "prompt_tokens_details": { + "cached_tokens": 654 + } + } + }) + ); + } + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider.test.ts › google-vertex-xai-provider › should support HTTP image URLs + #[tokio::test] + async fn should_support_http_image_urls() { + let fetch = MockFetch::new(vec![]); + let provider = provider(&fetch, Some("global"), None); + + let supported_urls = provider + .call("xai/grok-4.1-fast-reasoning") + .supported_urls() + .await; + + let patterns: Vec<(&str, Vec<&str>)> = supported_urls + .0 + .iter() + .map(|(media_type, patterns)| { + ( + media_type.as_str(), + patterns.iter().map(regex::Regex::as_str).collect(), + ) + }) + .collect(); + assert_eq!(patterns, vec![("image/*", vec![r"^https?://.*$"])]); + } + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider.test.ts › google-vertex-xai-provider › should throw NoSuchModelError for embedding and image models + #[test] + fn should_throw_no_such_model_error_for_embedding_and_image_models() { + let provider = create_google_vertex_xai(VertexXaiProviderSettings { + project: Some("test-project".into()), + ..Default::default() + }) + .unwrap(); + + assert!(matches!( + provider.embedding_model("invalid-model-id"), + Err(AiMuxError::NoSuchModel { .. }) + )); + assert!(matches!( + provider.image_model("invalid-model-id"), + Err(AiMuxError::NoSuchModel { .. }) + )); + } +} diff --git a/aimux-providers/src/vertex/tests/xai/xai_provider_node.rs b/aimux-providers/src/vertex/tests/xai/xai_provider_node.rs new file mode 100644 index 000000000..9ac051f54 --- /dev/null +++ b/aimux-providers/src/vertex/tests/xai/xai_provider_node.rs @@ -0,0 +1,249 @@ +//! Port of `@ai-sdk/google-vertex` 5.0.100 +//! `src/xai/google-vertex-xai-provider-node.test.ts`. +//! +//! Upstream mocks the base provider and calls the fetch wrapper the Node +//! entry point hands it with a hand-made `(url, init)`. aimux has one entry +//! point, whose fetch layer adds the same headers: each case makes a chat +//! request through a mock transport, so the request URL is the provider's +//! (not `https://example.com/test`) and `init` is the request the +//! OpenAI-compatible model builds. Upstream's mocked token generator +//! (returning `mock-auth-token`) is the `generate_auth_token` setting, except +//! in the cases about the default generator, which run aimux's Google auth +//! over the mock transport. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_provider_utils::Resolvable; +use serde_json::json; + +use crate::test_support::mock_fetch::{Canned, EnvVar, MockFetch, RouteFetch, Seen}; +use crate::vertex::{ + GoogleAuthOptions, GoogleAuthScopes, VertexXaiProviderSettings, create_google_vertex_xai, + google_vertex_xai, +}; + +/// A minimal Chat Completions response. +fn chat_response() -> Canned { + Canned::json(&json!({ + "id": "chatcmpl-1", "object": "chat.completion", "created": 1, "model": "m", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + })) +} + +fn prompt() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +/// A token generator returning `mock-auth-token`, and its call counter. +fn mock_generator() -> (Resolvable>, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let generator = Resolvable::from_async_fn(move || { + counter.fetch_add(1, Ordering::SeqCst); + async { Ok(Some("mock-auth-token".to_string())) } + }); + (generator, calls) +} + +/// Settings of a provider for `test-project` over `fetch`, with the mocked +/// token generator. +fn settings(fetch: &Arc) -> (VertexXaiProviderSettings, Arc) { + let (generator, calls) = mock_generator(); + let settings = VertexXaiProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + fetch: Some(fetch.transport()), + generate_auth_token: Some(generator), + ..Default::default() + }; + (settings, calls) +} + +/// Generate once with `options` and return the request the transport saw. +async fn request( + settings: VertexXaiProviderSettings, + fetch: &Arc, + options: &CallOptions, +) -> Seen { + create_google_vertex_xai(settings) + .unwrap() + .call("xai/grok-4.1-fast-reasoning") + .do_generate(options) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + seen[0].clone() +} + +/// The request headers except `skip`. +fn sent_headers(seen: &Seen, skip: &[&str]) -> BTreeMap { + seen.headers + .iter() + .filter(|(name, _)| !skip.contains(&name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +fn headers(entries: &[(&str, &str)]) -> BTreeMap { + entries + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() +} + +/// No credential file or gcloud configuration and the default metadata +/// host: the default Google auth asks the metadata server for a token. +fn metadata_only_env() -> Vec { + vec![ + EnvVar::set("GOOGLE_APPLICATION_CREDENTIALS", None), + EnvVar::set("google_application_credentials", None), + EnvVar::set( + "CLOUDSDK_CONFIG", + Some("/nonexistent/aimux-vertex-test-gcloud"), + ), + EnvVar::set("GCE_METADATA_IP", None), + EnvVar::set("GCE_METADATA_HOST", None), + ] +} + +const METADATA_TOKEN_PATH: &str = "/computeMetadata/v1/instance/service-accounts/default/token"; + +/// The metadata server's token response. +fn metadata_token(token: &str) -> Canned { + let mut canned = Canned::json(&json!({ + "access_token": token, "expires_in": 3600, "token_type": "Bearer" + })); + canned + .headers + .push(("Metadata-Flavor".into(), "Google".into())); + canned +} + +mod google_vertex_xai_provider_node { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider-node.test.ts › google-vertex-xai-provider-node › should generate auth token and add to request headers + /// + /// `createAuthTokenGenerator(undefined)` has no counterpart once the + /// generator is replaced; the generator's single call is asserted. + #[tokio::test] + async fn should_generate_auth_token_and_add_to_request_headers() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, calls) = settings(&fetch); + + let seen = request(settings, &fetch, &prompt()).await; + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(seen.method, "POST"); + assert_eq!( + sent_headers(&seen, &["user-agent"]), + headers(&[ + ("content-type", "application/json"), + ("authorization", "Bearer mock-auth-token"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider-node.test.ts › google-vertex-xai-provider-node › should pass googleAuthOptions to createAuthTokenGenerator + /// + /// The options reach aimux's Google auth (upstream mocks the generator): + /// with no credentials it asks the metadata server for a token of the + /// given scope, once. + #[tokio::test] + #[serial_test::serial] + async fn should_pass_google_auth_options_to_create_auth_token_generator() { + let _env = metadata_only_env(); + let fetch = RouteFetch::new(|seen| { + if seen.url.contains(METADATA_TOKEN_PATH) { + metadata_token("mock-auth-token") + } else { + chat_response() + } + }); + let provider = create_google_vertex_xai(VertexXaiProviderSettings { + project: Some("test-project".into()), + location: Some("global".into()), + fetch: Some(fetch.transport()), + google_auth_options: Some(GoogleAuthOptions { + scopes: Some(GoogleAuthScopes::Array(vec!["test-scope".into()])), + ..Default::default() + }), + ..Default::default() + }) + .unwrap(); + + provider + .call("xai/grok-4.1-fast-reasoning") + .do_generate(&prompt()) + .await + .unwrap(); + + assert_eq!(fetch.count("GET", METADATA_TOKEN_PATH), 1); + let seen = fetch.seen(); + let metadata = seen + .iter() + .find(|seen| seen.url.contains(METADATA_TOKEN_PATH)) + .unwrap(); + assert!( + metadata.url.ends_with("?scopes=test-scope"), + "{}", + metadata.url + ); + let chat = seen + .iter() + .find(|seen| seen.url.ends_with("/chat/completions")) + .unwrap(); + assert_eq!(chat.headers["authorization"], "Bearer mock-auth-token"); + } + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider-node.test.ts › google-vertex-xai-provider-node › should merge custom headers with auth header + #[tokio::test] + async fn should_merge_custom_headers_with_auth_header() { + let fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&fetch); + let settings = VertexXaiProviderSettings { + headers: Some(Resolvable::Value( + [("X-Custom".to_string(), Some("header-value".to_string()))].into(), + )), + ..settings + }; + + let seen = request(settings, &fetch, &prompt()).await; + + assert_eq!( + sent_headers(&seen, &["user-agent"]), + headers(&[ + ("content-type", "application/json"), + ("x-custom", "header-value"), + ("authorization", "Bearer mock-auth-token"), + ]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider-node.test.ts › google-vertex-xai-provider-node › should use custom fetch when provided + #[tokio::test] + async fn should_use_custom_fetch_when_provided() { + let custom_fetch = MockFetch::new(vec![chat_response()]); + let (settings, _) = settings(&custom_fetch); + + let seen = request(settings, &custom_fetch, &prompt()).await; + + assert_eq!( + sent_headers(&seen, &["content-type", "user-agent"]), + headers(&[("authorization", "Bearer mock-auth-token")]) + ); + } + + /// upstream: @ai-sdk/google-vertex src/xai/google-vertex-xai-provider-node.test.ts › google-vertex-xai-provider-node › should export default googleVertexXai instance + #[test] + fn should_export_default_google_vertex_xai_instance() { + let model = google_vertex_xai().call("xai/grok-4.1-fast-reasoning"); + assert_eq!(model.provider(), "googleVertex.xai.chat"); + } +} diff --git a/aimux-providers/tests/google_vertex_image_test.rs b/aimux-providers/tests/google_vertex_image_test.rs index 0722a0766..0424571f5 100644 --- a/aimux-providers/tests/google_vertex_image_test.rs +++ b/aimux-providers/tests/google_vertex_image_test.rs @@ -36,6 +36,7 @@ fn options(prompt: &str) -> ImageCallOptions { ImageCallOptions::new(prompt.to_string()) } +/// partial: `should use the language model endpoint and extract generated images`; full port in src/vertex/tests/image_model.rs. #[tokio::test] async fn imagen_should_extract_generated_images() { let server = MockServer::start().await; @@ -53,6 +54,7 @@ async fn imagen_should_extract_generated_images() { } } +/// partial: `should send response modalities, aspect ratio, seed, and headers`; full port in src/vertex/tests/image_model.rs. #[tokio::test] async fn imagen_should_send_aspect_ratio() { let server = MockServer::start().await; @@ -75,6 +77,7 @@ async fn imagen_should_send_aspect_ratio() { ); } +/// partial: `should return a warning for the unsupported size option`; full port in src/vertex/tests/image_model.rs. #[tokio::test] async fn imagen_should_warn_for_size() { let server = MockServer::start().await; @@ -92,6 +95,7 @@ async fn imagen_should_warn_for_size() { assert!(result.warnings.iter().any(|w| w.feature() == "size")); } +/// partial: `should include input images for editing`; full port in src/vertex/tests/image_model.rs. #[tokio::test] async fn imagen_should_support_editing() { let server = MockServer::start().await; @@ -117,6 +121,7 @@ async fn imagen_should_support_editing() { ); } +/// partial: `should use the language model endpoint and extract generated images`; full port in src/vertex/tests/image_model.rs. #[tokio::test] async fn gemini_should_extract_image() { let server = MockServer::start().await; @@ -136,6 +141,7 @@ async fn gemini_should_extract_image() { } } +/// partial: `should send response modalities, aspect ratio, seed, and headers`; full port in src/vertex/tests/image_model.rs. #[tokio::test] async fn gemini_should_send_response_modalities() { let server = MockServer::start().await; @@ -157,6 +163,7 @@ async fn gemini_should_send_response_modalities() { ); } +/// partial: `should send response modalities, aspect ratio, seed, and headers`; full port in src/vertex/tests/image_model.rs. #[tokio::test] async fn gemini_should_pass_aspect_ratio() { let server = MockServer::start().await; @@ -179,6 +186,7 @@ async fn gemini_should_pass_aspect_ratio() { ); } +/// partial: `should reject unsupported URL editing input and masks`; full port in src/vertex/tests/image_model.rs. #[tokio::test] async fn gemini_should_throw_for_mask() { let server = MockServer::start().await; @@ -206,6 +214,7 @@ impl WarningFeature for aimux_core::types::Warning { } } +/// partial: `advertises file=$supportsFileInputs and mask=$supportsMaskInputs for $modelId`; full port in src/vertex/tests/image_model.rs. #[test] fn gemini_image_models_advertise_file_inputs_and_no_masks() { let provider = create_google_vertex(VertexProviderSettings { diff --git a/aimux-providers/tests/vertex_location_validation_test.rs b/aimux-providers/tests/vertex_location_validation_test.rs index 7f4c2bd50..7978f5b93 100644 --- a/aimux-providers/tests/vertex_location_validation_test.rs +++ b/aimux-providers/tests/vertex_location_validation_test.rs @@ -115,54 +115,58 @@ async fn rejects(kind: Kind, location: Option<&str>, base_url: Option<&str>) -> rejected } -/// TS: "rejects %j before fetching" -#[tokio::test] -async fn rejects_a_location_that_is_not_one_dns_label() { - for kind in KINDS { - for location in INVALID { - assert!( - rejects(kind, Some(location), None).await, - "{kind:?} {location:?}" - ); +mod name_location_validation { + use super::*; + + /// upstream: @ai-sdk/google-vertex src/google-vertex-location-validation.test.ts › ${…} › rejects %j before fetching + #[tokio::test] + async fn rejects_a_location_that_is_not_one_dns_label() { + for kind in KINDS { + for location in INVALID { + assert!( + rejects(kind, Some(location), None).await, + "{kind:?} {location:?}" + ); + } } } -} -/// TS: "accepts %j" -#[tokio::test] -async fn accepts_regional_and_multi_region_locations() { - for kind in KINDS { - for location in VALID { - assert!( - !rejects(kind, Some(location), None).await, - "{kind:?} {location:?}" - ); + /// upstream: @ai-sdk/google-vertex src/google-vertex-location-validation.test.ts › ${…} › accepts %j + #[tokio::test] + async fn accepts_regional_and_multi_region_locations() { + for kind in KINDS { + for location in VALID { + assert!( + !rejects(kind, Some(location), None).await, + "{kind:?} {location:?}" + ); + } } } -} -/// TS: "uses the environment location" -#[tokio::test] -#[serial] -async fn validates_the_environment_location() { - let _location = EnvVar::set("GOOGLE_VERTEX_LOCATION", Some("user@internal/#")); - for kind in KINDS { - assert!(rejects(kind, None, None).await, "{kind:?}"); + /// upstream: @ai-sdk/google-vertex src/google-vertex-location-validation.test.ts › ${…} › uses the environment location + #[tokio::test] + #[serial] + async fn validates_the_environment_location() { + let _location = EnvVar::set("GOOGLE_VERTEX_LOCATION", Some("user@internal/#")); + for kind in KINDS { + assert!(rejects(kind, None, None).await, "{kind:?}"); + } } -} -/// TS: "allows a custom URL with an unused invalid location" -#[tokio::test] -async fn a_custom_url_ignores_an_invalid_location() { - for kind in KINDS { - assert!( - !rejects(kind, Some("not a region"), Some("https://proxy.example/v1")).await, - "{kind:?}" - ); + /// upstream: @ai-sdk/google-vertex src/google-vertex-location-validation.test.ts › ${…} › allows a custom URL with an unused invalid location + #[tokio::test] + async fn a_custom_url_ignores_an_invalid_location() { + for kind in KINDS { + assert!( + !rejects(kind, Some("not a region"), Some("https://proxy.example/v1")).await, + "{kind:?}" + ); + } } } -/// TS: "allows a custom URL without a location" +/// partial: `allows a custom URL without a location`; full port in src/vertex/tests/location_validation.rs. #[tokio::test] #[serial] async fn a_custom_url_needs_no_location() { @@ -175,7 +179,7 @@ async fn a_custom_url_needs_no_location() { } } -/// TS: "does not validate location in express mode" +/// upstream: @ai-sdk/google-vertex src/google-vertex-location-validation.test.ts › does not validate location in express mode #[tokio::test] async fn express_mode_ignores_the_location() { let fetch = MockFetch::new(vec![]); @@ -197,7 +201,7 @@ async fn express_mode_ignores_the_location() { assert_eq!(fetch.seen().len(), 1); } -/// TS: "validates location for the separate transcription host of %s". +/// partial: `validates location for the separate transcription host of %s`; full port in src/vertex/tests/location_validation.rs. /// Upstream's Gemini case is the live model; aimux's Gemini transcription /// models, unary and live, resolve the project and location the same way, and /// the unary one runs without a socket. diff --git a/aimux-providers/tests/vertex_transcription_test.rs b/aimux-providers/tests/vertex_transcription_test.rs index 23212dd7d..fd2d14334 100644 --- a/aimux-providers/tests/vertex_transcription_test.rs +++ b/aimux-providers/tests/vertex_transcription_test.rs @@ -60,6 +60,7 @@ async fn mock_response(server: &MockServer, body: &Value, headers: &[(&str, &str .await; } +/// partial: `should send the model, languageCodes, features and base64 content`; full port in src/vertex/tests/transcription_model.rs. #[tokio::test] async fn should_send_model_language_codes_features_and_base64() { let server = MockServer::start().await; @@ -86,6 +87,7 @@ async fn should_send_model_language_codes_features_and_base64() { assert_eq!(body["content"], "AQIDBAUGBwg="); } +/// partial: `should extract text, segments, language and duration`; full port in src/vertex/tests/transcription_model.rs. #[tokio::test] async fn should_extract_text_segments_language_and_duration() { let server = MockServer::start().await; @@ -111,6 +113,7 @@ async fn should_extract_text_segments_language_and_duration() { assert_eq!(result.duration_in_seconds, Some(1.0)); } +/// partial: `should pass headers and user agent`; full port in src/vertex/tests/transcription_model.rs. #[tokio::test] async fn should_pass_headers() { let server = MockServer::start().await; @@ -158,6 +161,7 @@ async fn should_pass_headers() { ); } +/// partial: `should include response data with timestamp, modelId and headers`; full port in src/vertex/tests/transcription_model.rs. #[tokio::test] async fn should_include_response_data() { let server = MockServer::start().await; @@ -184,6 +188,7 @@ async fn should_include_response_data() { ); } +/// partial: `should return empty text and segments when there are no results`; full port in src/vertex/tests/transcription_model.rs. #[tokio::test] async fn should_handle_no_results() { let server = MockServer::start().await; diff --git a/aimux-providers/tests/vertex_upstream_e2e_test.rs b/aimux-providers/tests/vertex_upstream_e2e_test.rs index d26c82586..b7a1ef54e 100644 --- a/aimux-providers/tests/vertex_upstream_e2e_test.rs +++ b/aimux-providers/tests/vertex_upstream_e2e_test.rs @@ -33,7 +33,7 @@ fn provider(fetch: &Arc) -> VertexProvider { .unwrap() } -/// TS: "should preserve local JSON Schema references in tool requests" (@ai-sdk/google-vertex/src/google-vertex-language-model.test.ts) +/// partial: `should preserve local JSON Schema references in tool requests`; full port in src/vertex/tests/language_model.rs. #[tokio::test] async fn language_generate_preserves_tool_schema_references() { let response = json!({ @@ -104,9 +104,9 @@ async fn language_generate_preserves_tool_schema_references() { assert_eq!(result.response.unwrap().body.unwrap(), response); } -/// TS: "should pass the model parameters correctly" (@ai-sdk/google-vertex/src/google-vertex-embedding-model.test.ts) -/// TS: "should extract embeddings" (@ai-sdk/google-vertex/src/google-vertex-embedding-model.test.ts) -/// TS: "should extract usage" (@ai-sdk/google-vertex/src/google-vertex-embedding-model.test.ts) +/// partial: `should pass the model parameters correctly`; full port in src/vertex/tests/embedding_model.rs. +/// partial: `should extract embeddings`; full port in src/vertex/tests/embedding_model.rs. +/// partial: `should extract usage`; full port in src/vertex/tests/embedding_model.rs. #[tokio::test] async fn embedding_request_and_result() { let response = json!({"predictions": [ @@ -162,9 +162,9 @@ async fn embedding_request_and_result() { assert_eq!(result.response.unwrap().body.unwrap(), response); } -/// TS: "should send response modalities, aspect ratio, seed, and headers" (@ai-sdk/google-vertex/src/google-vertex-image-model.test.ts) -/// TS: "should use the language model endpoint and extract generated images" (@ai-sdk/google-vertex/src/google-vertex-image-model.test.ts) -/// TS: "should include usage and response metadata" (@ai-sdk/google-vertex/src/google-vertex-image-model.test.ts) +/// partial: `should send response modalities, aspect ratio, seed, and headers`; full port in src/vertex/tests/image_model.rs. +/// partial: `should use the language model endpoint and extract generated images`; full port in src/vertex/tests/image_model.rs. +/// partial: `should include usage and response metadata`; full port in src/vertex/tests/image_model.rs. #[tokio::test] async fn image_generate_request_and_result() { let response = json!({ @@ -247,8 +247,8 @@ async fn image_generate_request_and_result() { ); } -/// TS: "should send the model, languageCodes, features and base64 content" (@ai-sdk/google-vertex/src/google-vertex-transcription-model.test.ts) -/// TS: "should extract text, segments, language and duration" (@ai-sdk/google-vertex/src/google-vertex-transcription-model.test.ts) +/// partial: `should send the model, languageCodes, features and base64 content`; full port in src/vertex/tests/transcription_model.rs. +/// partial: `should extract text, segments, language and duration`; full port in src/vertex/tests/transcription_model.rs. #[tokio::test] async fn transcription_generate_request_and_result() { let fetch = MockFetch::new(vec![Canned::json(&json!({ @@ -300,9 +300,9 @@ async fn transcription_generate_request_and_result() { assert_eq!(result.duration_in_seconds, Some(1.0)); } -/// TS: "should pass correct request body" (@ai-sdk/google-vertex/src/google-vertex-video-model.test.ts) -/// TS: "should return operation with operationName" (@ai-sdk/google-vertex/src/google-vertex-video-model.test.ts) -/// TS: "should return completed with video data when done" (@ai-sdk/google-vertex/src/google-vertex-video-model.test.ts) +/// partial: `should pass correct request body`; full port in src/vertex/tests/video_model.rs. +/// partial: `should return operation with operationName`; full port in src/vertex/tests/video_model.rs. +/// partial: `should return completed with video data when done`; full port in src/vertex/tests/video_model.rs. #[tokio::test] async fn video_start_and_completed_result() { let fetch = MockFetch::new(vec![ @@ -355,7 +355,7 @@ async fn video_start_and_completed_result() { ); } -/// TS: "should throw when no operation name returned" (@ai-sdk/google-vertex/src/google-vertex-video-model.test.ts) +/// partial: `should throw when no operation name returned`; full port in src/vertex/tests/video_model.rs. #[tokio::test] async fn video_missing_operation_name_is_an_error() { let fetch = MockFetch::new(vec![Canned::json(&json!({"done": false}))]); @@ -377,8 +377,7 @@ async fn video_missing_operation_name_is_an_error() { ); } -/// TS: "transcribes audio via Vertex generateContent with audioTranscriptionConfig" -/// (google-vertex/src/gemini-transcription/google-vertex-gemini-transcription-model.test.ts) +/// partial: `transcribes audio via Vertex generateContent with audioTranscriptionConfig`; full port in src/vertex/tests/gemini_transcription/gemini_transcription_model.rs. #[tokio::test] async fn gemini_transcription_generate_request_and_result() { let usage = json!({"promptTokenCount": 10, "candidatesTokenCount": 4}); @@ -492,6 +491,7 @@ async fn executable_credential_without_the_opt_in_fails_and_spawns_nothing() { assert!(fetch.seen().is_empty(), "a request was sent"); } +/// partial: `should accept the legacy vertex options key and prefer googleVertex`; full port in src/vertex/tests/image_model.rs. /// Upstream: @ai-sdk/google-vertex 5.0.100 src/google-vertex-image-model.ts /// reads the image options under `googleVertex`, then the legacy `vertex`, /// hands them to the inner language model under both keys (:190-194) and @@ -553,6 +553,7 @@ fn partner_settings( } } +/// partial: `should create a provider with correct base URL for regional location`, `should default to global location when not specified`, `should default Llama 4 max tokens without overriding explicit or other model settings`, `should generate auth token and add to request headers`; full ports in src/vertex/tests/maas/maas_provider.rs and src/vertex/tests/maas/maas_provider_node.rs. /// Upstream: @ai-sdk/google-vertex 5.0.100 src/maas/google-vertex-maas-provider.ts:79-130 /// (the project/location URL, `global` by default, the `vertex.maas` name, /// the llama-4 `max_tokens` default at :14-25) and @@ -595,6 +596,7 @@ async fn maas_sub_provider_builds_on_openai_compatible() { assert_eq!(seen[1].json_body()["max_tokens"], 100); } +/// partial: `should create a provider with correct base URL for regional location`, `should strip reasoning_effort from request bodies`, `should count Grok reasoning tokens separately from completion tokens`, `should support HTTP image URLs`, `should throw NoSuchModelError for embedding and image models`; full ports in src/vertex/tests/xai/xai_provider.rs. /// Upstream: @ai-sdk/google-vertex 5.0.100 src/xai/google-vertex-xai-provider.ts: /// the global host whatever the location (:159-164), the `googleVertex.xai` /// name with usage in streams and structured outputs (:172-184), no @@ -669,6 +671,7 @@ async fn partner_sub_providers_require_a_project() { assert!(fetch.seen().is_empty()); } +/// partial: `should merge custom headers with auth header`, `should preserve headers from init parameter`; full ports in src/vertex/tests/maas/maas_provider_node.rs (the first also in src/vertex/tests/xai/xai_provider_node.rs). /// Upstream: @ai-sdk/google-vertex 5.0.100 maas/google-vertex-maas-provider-node.ts:34-58 /// and xai/google-vertex-xai-provider-node.ts:33-55 hand no headers to /// `createOpenAICompatible`; their fetch spreads the provider headers, then diff --git a/aimux-providers/tests/vertex_video_test.rs b/aimux-providers/tests/vertex_video_test.rs index 42a556d62..fb7f779e1 100644 --- a/aimux-providers/tests/vertex_video_test.rs +++ b/aimux-providers/tests/vertex_video_test.rs @@ -68,6 +68,7 @@ async fn mock_predict_and_poll(server: &MockServer, result: &Value) { .await; } +/// partial: `should return completed with GCS URI video`; full port in src/vertex/tests/video_model.rs. #[tokio::test] async fn should_generate_video() { let server = MockServer::start().await; @@ -96,6 +97,7 @@ async fn should_generate_video() { } } +/// partial: `should pass correct request body`; full port in src/vertex/tests/video_model.rs. #[tokio::test] async fn should_pass_prompt() { let server = MockServer::start().await; @@ -123,6 +125,7 @@ async fn should_pass_prompt() { assert_eq!(poll_body, json!({"operationName": "operations/test-op"})); } +/// partial: `should include warnings and response metadata`; full port in src/vertex/tests/video_model.rs. #[tokio::test] async fn should_include_response_data() { let server = MockServer::start().await;