From 59a2e65cefac6ec6d39cc4ec2f297616b020e154 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 12 Oct 2026 09:33:59 +0000 Subject: [PATCH] fix(providers): vendor behaviour follows the upstream packages Closes the vendor differences left after the move to factories, checked against each pinned upstream package: - Factories, request settings, base URLs, empty keys and provider names follow upstream; vendor option schemas live in each package and vendor namespaces are spelled through the package constants. - Vendor features upstream has and aimux lacked are ported, including the image-model middleware. - Stream error events become error parts; Cohere and Mistral detect media types as upstream does; Bedrock converter errors and Anthropic ids follow upstream; the Azure user-agent uses the package namespace. - The OpenAI and OpenAI-compatible chat models drive the tool-call tracker with their package's options and flush pending tool calls through it. - API requests follow cross-origin redirects and drop credentials on the way. A local runtime's optional key has its own credential kind. - create_provider lists every built-in name when a name misses. - Decision providers follow call-level retries and identity-only replay; the migrated recording fixtures are marked schema 3. Tests: the upstream-ported model tests (sixteen files, plus the Bedrock and Anthropic remaining cases) are kept, AI SDK protocol samples are recorded under a mocked fetch, the DeepSeek cassette replay is kept, and the hand-written tests the factories replace are deleted. An ignored live SigV4 test checks against STS, Bedrock and Polly. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- CHANGELOG.md | 16 + aimux-core/src/lib.rs | 6 +- aimux-core/src/replay.rs | 1 + aimux-core/src/result.rs | 4 +- aimux-core/src/retry.rs | 11 - aimux-ffi/tests/exports_smoke_test.rs | 33 +- aimux-provider-utils/src/ws.rs | 20 + aimux-providers/Cargo.toml | 1 + .../src/anthropic/cache_control.rs | 4 + aimux-providers/src/anthropic/config.rs | 32 +- aimux-providers/src/anthropic/convert.rs | 954 +++++++++-- aimux-providers/src/anthropic/files.rs | 11 +- aimux-providers/src/anthropic/mod.rs | 42 +- aimux-providers/src/anthropic/model.rs | 74 +- .../src/anthropic/prepare_tools.rs | 165 +- aimux-providers/src/anthropic/skills.rs | 7 +- aimux-providers/src/anthropic/stream.rs | 577 +++++-- .../src/anthropic/tool_name_mapping.rs | 48 +- aimux-providers/src/anthropic/types.rs | 77 +- aimux-providers/src/anthropic/usage.rs | 46 +- aimux-providers/src/azure/mod.rs | 156 +- aimux-providers/src/azure/options.rs | 6 +- aimux-providers/src/azure/transcription.rs | 374 ++++ aimux-providers/src/bedrock/convert.rs | 1512 +++++++++++++---- aimux-providers/src/bedrock/embedding.rs | 126 +- aimux-providers/src/bedrock/event_stream.rs | 129 +- aimux-providers/src/bedrock/image.rs | 121 +- aimux-providers/src/bedrock/mod.rs | 137 +- aimux-providers/src/bedrock/model.rs | 473 ++++-- aimux-providers/src/bedrock/options.rs | 14 +- aimux-providers/src/bedrock/reranking.rs | 21 +- aimux-providers/src/bedrock/tools.rs | 282 +++ aimux-providers/src/bedrock/types.rs | 47 +- aimux-providers/src/cohere/convert.rs | 94 +- aimux-providers/src/cohere/embedding.rs | 122 +- aimux-providers/src/cohere/mod.rs | 63 +- aimux-providers/src/cohere/model.rs | 435 ++--- aimux-providers/src/cohere/options.rs | 2 +- aimux-providers/src/cohere/reranking.rs | 39 +- aimux-providers/src/cohere/types.rs | 110 +- aimux-providers/src/deepseek/convert.rs | 119 +- aimux-providers/src/deepseek/model.rs | 17 +- aimux-providers/src/deepseek/prepare_tools.rs | 2 +- aimux-providers/src/default_providers.rs | 33 +- aimux-providers/src/google/convert.rs | 1136 ++++++++----- aimux-providers/src/google/embedding.rs | 293 ++-- aimux-providers/src/google/files.rs | 90 +- aimux-providers/src/google/image.rs | 581 ++----- aimux-providers/src/google/mod.rs | 19 +- aimux-providers/src/google/model.rs | 387 +++-- aimux-providers/src/google/options.rs | 46 +- aimux-providers/src/google/types.rs | 12 +- aimux-providers/src/google/utils.rs | 71 +- aimux-providers/src/google/video.rs | 244 ++- aimux-providers/src/groq/convert.rs | 15 +- aimux-providers/src/groq/model.rs | 25 +- aimux-providers/src/groq/prepare_tools.rs | 19 +- aimux-providers/src/huggingface/options.rs | 2 +- aimux-providers/src/huggingface/responses.rs | 4 +- aimux-providers/src/lib.rs | 7 +- aimux-providers/src/mistral/convert.rs | 463 ++--- aimux-providers/src/mistral/embedding.rs | 57 +- aimux-providers/src/mistral/mod.rs | 40 +- aimux-providers/src/mistral/model.rs | 230 +-- aimux-providers/src/mistral/options.rs | 31 +- aimux-providers/src/mistral/types.rs | 45 +- aimux-providers/src/open_responses.rs | 4 +- aimux-providers/src/openai/config.rs | 13 + aimux-providers/src/openai/convert.rs | 271 ++- aimux-providers/src/openai/convert_common.rs | 36 +- aimux-providers/src/openai/embedding.rs | 123 +- aimux-providers/src/openai/files.rs | 25 +- aimux-providers/src/openai/image.rs | 140 +- aimux-providers/src/openai/mod.rs | 114 +- aimux-providers/src/openai/model.rs | 78 +- .../src/openai/responses/convert.rs | 1324 ++++++++++++--- aimux-providers/src/openai/responses/mod.rs | 204 ++- .../src/openai/responses/provider_events.rs | 314 ++++ .../src/openai/responses/responses_convert.rs | 663 +++++--- .../src/openai/responses/tool_args.rs | 437 +++++ aimux-providers/src/openai/responses/types.rs | 2 - aimux-providers/src/openai/speech.rs | 34 +- aimux-providers/src/openai/tools.rs | 713 ++++++++ aimux-providers/src/openai/transcription.rs | 14 +- aimux-providers/src/openai/types.rs | 26 +- aimux-providers/src/openai_compatible/chat.rs | 418 ++--- .../src/openai_compatible/config.rs | 13 +- .../src/openai_compatible/convert.rs | 284 ++-- .../src/openai_compatible/embedding.rs | 87 +- .../src/openai_compatible/image.rs | 184 +- aimux-providers/src/openai_compatible/mod.rs | 98 +- .../src/openai_compatible/types.rs | 8 +- aimux-providers/src/preset.rs | 32 +- aimux-providers/src/provider.rs | 3 - aimux-providers/src/shared/mod.rs | 3 - aimux-providers/src/vertex/anthropic_model.rs | 6 +- .../src/vertex/anthropic_provider.rs | 116 ++ aimux-providers/src/vertex/auth.rs | 1090 ++++++++++++ .../src/vertex/auth_certificate.rs | 194 +++ aimux-providers/src/vertex/auth_external.rs | 383 +++++ aimux-providers/src/vertex/embedding.rs | 71 +- .../src/vertex/gemini_transcription.rs | 459 +++++ aimux-providers/src/vertex/image.rs | 355 ++-- aimux-providers/src/vertex/mod.rs | 273 ++- aimux-providers/src/vertex/model.rs | 778 +++++---- aimux-providers/src/vertex/transcription.rs | 24 +- aimux-providers/src/vertex/video.rs | 245 ++- aimux-providers/src/xai/convert.rs | 68 +- aimux-providers/src/xai/mod.rs | 68 +- aimux-providers/src/xai/options.rs | 2 +- aimux-providers/src/xai/responses/convert.rs | 326 ++-- aimux-providers/src/xai/responses/mod.rs | 328 +++- aimux-providers/src/xai/responses/types.rs | 28 +- aimux-providers/tests/alibaba_test.rs | 60 +- .../tests/amazon_bedrock_image_test.rs | 4 +- .../tests/anthropic_remaining_test.rs | 524 ++++++ .../tests/anthropic_upstream_e2e_test.rs | 253 +++ .../tests/azure_upstream_e2e_test.rs | 373 ++++ aimux-providers/tests/bedrock_convert_test.rs | 58 +- .../tests/bedrock_embedding_test.rs | 12 +- .../tests/bedrock_remaining_test.rs | 1215 +++++++++++++ .../tests/bedrock_upstream_e2e_test.rs | 469 +++++ .../tests/cartesia_transcription_test.rs | 59 + .../tests/cassette_multimodal_test.rs | 45 - aimux-providers/tests/codex_test.rs | 2 - .../tests/cohere_upstream_e2e_test.rs | 480 ++++++ aimux-providers/tests/common/replay.rs | 4 +- aimux-providers/tests/conformance_test.rs | 28 - .../tests/data_loss_regression_test.rs | 190 ++- aimux-providers/tests/dataforseo_test.rs | 2 - .../tests/deepgram_transcription_test.rs | 82 + aimux-providers/tests/deepseek_chat_test.rs | 148 +- aimux-providers/tests/google_factory_test.rs | 6 - aimux-providers/tests/google_pse_test.rs | 2 - .../tests/google_upstream_e2e_test.rs | 404 +++++ aimux-providers/tests/groq_test.rs | 72 +- .../tests/huggingface_responses_test.rs | 9 +- aimux-providers/tests/jina_ai_test.rs | 2 - aimux-providers/tests/linkup_test.rs | 2 - .../tests/mistral_upstream_e2e_test.rs | 253 +++ aimux-providers/tests/moonshotai_test.rs | 44 - aimux-providers/tests/open_responses_test.rs | 171 +- .../tests/openai_compatible_factory_test.rs | 2 +- .../openai_compatible_upstream_e2e_test.rs | 394 +++++ aimux-providers/tests/openai_factory_test.rs | 6 +- .../tests/openai_upstream_e2e_test.rs | 574 +++++++ aimux-providers/tests/openrouter_test.rs | 4 - aimux-providers/tests/parallel_ai_test.rs | 2 - aimux-providers/tests/provider_error_test.rs | 99 -- aimux-providers/tests/recraft_test.rs | 4 - aimux-providers/tests/replicate_image_test.rs | 26 + aimux-providers/tests/searxng_test.rs | 4 - .../tests/sse_terminal_error_test.rs | 8 +- aimux-providers/tests/tinyfish_test.rs | 2 - .../tests/vertex_upstream_e2e_test.rs | 426 +++++ .../tests/voyage_embedding_test.rs | 62 + .../tests/xai_upstream_e2e_test.rs | 266 +++ aimux-providers/tests/you_com_test.rs | 2 - docs/api/providers.md | 2 +- 159 files changed, 21239 insertions(+), 6476 deletions(-) create mode 100644 aimux-providers/src/azure/transcription.rs create mode 100644 aimux-providers/src/bedrock/tools.rs create mode 100644 aimux-providers/src/openai/responses/provider_events.rs create mode 100644 aimux-providers/src/openai/responses/tool_args.rs create mode 100644 aimux-providers/src/openai/tools.rs create mode 100644 aimux-providers/src/vertex/anthropic_provider.rs create mode 100644 aimux-providers/src/vertex/auth.rs create mode 100644 aimux-providers/src/vertex/auth_certificate.rs create mode 100644 aimux-providers/src/vertex/auth_external.rs create mode 100644 aimux-providers/src/vertex/gemini_transcription.rs create mode 100644 aimux-providers/tests/anthropic_remaining_test.rs create mode 100644 aimux-providers/tests/anthropic_upstream_e2e_test.rs create mode 100644 aimux-providers/tests/azure_upstream_e2e_test.rs create mode 100644 aimux-providers/tests/bedrock_remaining_test.rs create mode 100644 aimux-providers/tests/bedrock_upstream_e2e_test.rs create mode 100644 aimux-providers/tests/cohere_upstream_e2e_test.rs create mode 100644 aimux-providers/tests/google_upstream_e2e_test.rs create mode 100644 aimux-providers/tests/mistral_upstream_e2e_test.rs create mode 100644 aimux-providers/tests/openai_compatible_upstream_e2e_test.rs create mode 100644 aimux-providers/tests/openai_upstream_e2e_test.rs create mode 100644 aimux-providers/tests/vertex_upstream_e2e_test.rs create mode 100644 aimux-providers/tests/xai_upstream_e2e_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index d185890db..326678d19 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `available_providers`. Node and Python OpenAI constructors use Responses by default. +- Vendor packages follow their upstream counterparts more closely. Visible + changes: Bedrock percent-encodes the model id in request URLs and writes + provider metadata under both `amazonBedrock` and `bedrock`; OpenAI embeddings + request `encoding_format: "float"` and no longer decode base64 vectors; the + Google provider result no longer carries a model id unless the API returns + one; an Anthropic mid-stream `error` event is forwarded as an error part + without a synthetic finish; each package sends its own `user-agent` suffix; + Mistral has no `transform_request_body` setting; the deprecated + `text_embedding` and `text_embedding_model` provider methods are removed. + - Public surface (`aimux-core`, `aimux-provider-utils`): removed `recording::init_recording_from_env`; made private the retry preparation helpers and default constants, `composite::{add_usage, @@ -156,6 +166,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `tool-call-N` id instead of an empty string; a new call without a function name ends the stream with `InvalidResponseData` (previously it started a call with an empty name). `DeltaToolCall.index` is now `Option`. +- Once compatible servers moved to `openai_compatible`, the OpenAI (and + Azure) chat stream drives the tracker with `@ai-sdk/openai`'s options: + `type` is validated when present, ids missing from the wire come from + `generate_id` instead of `tool-call` / `tool-call-N`, and + `DeltaToolCall.index` is required again, so a chunk without it fails to + parse. **Rust (provider factory, RFC-0036)** diff --git a/aimux-core/src/lib.rs b/aimux-core/src/lib.rs index 09b7470b5..4d9db9856 100644 --- a/aimux-core/src/lib.rs +++ b/aimux-core/src/lib.rs @@ -129,9 +129,9 @@ pub mod prelude { init_session_store, list_sessions, session_calls, }; pub use crate::shared::{ - AspectRatio, FileBytes, FileData, GeneratedFileData, JsonObject, SharedHeaders, - SharedProviderMetadata, SharedProviderOptions, SharedProviderReference, Size, - provider_namespace, + AspectRatio, FileBytes, FileData, GeneratedFileData, JsonObject, ResponseInfo, + SharedHeaders, SharedProviderMetadata, SharedProviderOptions, SharedProviderReference, + Size, StreamResponseInfo, provider_namespace, }; pub use crate::skills_model::{SkillFile, Skills, UploadSkillCallOptions, UploadSkillResult}; pub use crate::speech_model::{SpeechCallOptions, SpeechModel, SpeechResult, generate_speech}; diff --git a/aimux-core/src/replay.rs b/aimux-core/src/replay.rs index ebc162e09..934aadc83 100644 --- a/aimux-core/src/replay.rs +++ b/aimux-core/src/replay.rs @@ -952,6 +952,7 @@ pub struct ReplayOverrides { /// - `recording.input.options` 中的 headers/provider_options 已脱敏, /// 重发会用脱敏后的 `[REDACTED]` 值——需要真实头时用 `overrides` 或 /// 重建 provider 时补充。 +/// - 逐消息和逐内容部分的 `provider_options` 原样保留。 /// /// # Errors /// diff --git a/aimux-core/src/result.rs b/aimux-core/src/result.rs index 678947ce3..5f0d17c8b 100644 --- a/aimux-core/src/result.rs +++ b/aimux-core/src/result.rs @@ -6,7 +6,9 @@ use std::pin::Pin; use ts_rs::TS; use crate::error::AiMuxError; -use crate::shared::{GeneratedFileData, RequestInfo, ResponseInfo, StreamResponseInfo}; +pub use crate::shared::GeneratedFileData; +use crate::shared::{RequestInfo, ResponseInfo, StreamResponseInfo}; + use crate::stream_part::StreamPart; use crate::tool::{RawToolCall, ToolResult}; use crate::types::{FinishReason, ProviderMetadata, Usage, Warning}; diff --git a/aimux-core/src/retry.rs b/aimux-core/src/retry.rs index da3f99c06..497ba03a7 100644 --- a/aimux-core/src/retry.rs +++ b/aimux-core/src/retry.rs @@ -297,17 +297,6 @@ mod tests { assert_eq!(attempts.load(Ordering::SeqCst), 0); } - #[test] - fn prepared_retries_use_the_fixed_delay_and_backoff() { - let retries = prepare_retries(Some(4), None); - - assert_eq!(retries.max_retries, 4); - assert_eq!(retries.initial_delay_ms, DEFAULT_INITIAL_DELAY_MS); - assert_eq!(retries.backoff_factor, u64::from(DEFAULT_BACKOFF_FACTOR)); - assert_eq!(prepare_retries(None, None).max_retries, DEFAULT_MAX_RETRIES); - assert_eq!(prepare_retries(Some(0), None).max_retries, 0); - } - #[tokio::test] async fn first_non_retryable_error_is_not_wrapped() { let error = retry_with_exponential_backoff( diff --git a/aimux-ffi/tests/exports_smoke_test.rs b/aimux-ffi/tests/exports_smoke_test.rs index f75625171..4a5bf8701 100644 --- a/aimux-ffi/tests/exports_smoke_test.rs +++ b/aimux-ffi/tests/exports_smoke_test.rs @@ -68,16 +68,16 @@ use aimux_ffi::{ aimux_openai_image_new, aimux_openai_image_new_with_base, aimux_openai_new, aimux_openai_new_with_base, aimux_openai_speech_new, aimux_openai_speech_new_with_base, aimux_openai_transcription_new, aimux_openai_transcription_new_with_base, - aimux_provider_from_env, aimux_provider_handle_new, aimux_provider_list_models, - aimux_provider_model, aimux_provider_new, aimux_recording_flush, aimux_recording_stop, - aimux_recording_try_flush, aimux_register_providers, aimux_rerank, aimux_router_new, - aimux_search, aimux_session_calls, aimux_session_infer_init, aimux_session_store_init, - aimux_speech_generate, aimux_stream_text, aimux_stream_text_as_openai, - aimux_stream_text_as_openai_with_abort, aimux_stream_text_with_abort, aimux_tavily_search_new, - aimux_tavily_search_new_with_base, aimux_tool_call_repair_context, aimux_trace_aggregate, - aimux_trace_clear, aimux_trace_export_jsonl, aimux_trace_new, aimux_trace_new_audited, - aimux_trace_session_chain, aimux_trace_session_trajectory, aimux_transcription_generate, - aimux_transcription_input_done, aimux_transcription_next_part, aimux_transcription_push_audio, + aimux_provider_handle_new, aimux_provider_list_models, aimux_provider_model, + aimux_provider_new, aimux_recording_flush, aimux_recording_stop, aimux_recording_try_flush, + aimux_register_providers, aimux_rerank, aimux_router_new, aimux_search, aimux_session_calls, + aimux_session_infer_init, aimux_session_store_init, aimux_speech_generate, aimux_stream_text, + aimux_stream_text_as_openai, aimux_stream_text_as_openai_with_abort, + aimux_stream_text_with_abort, aimux_tavily_search_new, aimux_tavily_search_new_with_base, + aimux_tool_call_repair_context, aimux_trace_aggregate, aimux_trace_clear, + aimux_trace_export_jsonl, aimux_trace_new, aimux_trace_new_audited, aimux_trace_session_chain, + aimux_trace_session_trajectory, aimux_transcription_generate, aimux_transcription_input_done, + aimux_transcription_next_part, aimux_transcription_push_audio, aimux_transcription_session_drop, aimux_transcription_session_new, aimux_vertex_new, aimux_vertex_new_with_base, aimux_video_generate, aimux_xai_new, aimux_xai_new_with_base, }; @@ -569,19 +569,6 @@ fn constructor_exports_build_and_release_handles() { } } -/// `provider_from_env` defers reading the key until request time. -#[test] -fn provider_from_env_handle() { - let mut h = 0; - let e = aimux_provider_from_env( - c("groq").as_ptr(), - c("llama-3.3-70b-versatile").as_ptr(), - &mut h, - ); - expect_handle(e, h, "provider_from_env"); - aimux_drop_handle(h); -} - // ── session class: text generation (4 exports) ────────────────────────────── #[test] diff --git a/aimux-provider-utils/src/ws.rs b/aimux-provider-utils/src/ws.rs index c45458c6a..f11f408c9 100644 --- a/aimux-provider-utils/src/ws.rs +++ b/aimux-provider-utils/src/ws.rs @@ -57,6 +57,26 @@ pub struct WebSocketRequest { pub connector: Option>, } +impl WebSocketRequest { + /// Build a transcription connection with the operation's cancellation and timeouts. + #[must_use] + pub fn for_transcription( + url: String, + headers: Vec<(String, String)>, + options: &aimux_core::transcription_model::TranscriptionStreamOptions, + connector: Option>, + ) -> Self { + Self { + url, + headers, + subprotocols: Vec::new(), + abort_signal: options.abort_signal.clone(), + timeout: options.timeout, + connector, + } + } +} + /// Opens WebSocket connections. The built-in implementation is /// tungstenite plus the RFC-0034 proxy tunnel; a custom one returns a /// [`WsConnection`] built with [`WsConnection::from_transport`]. diff --git a/aimux-providers/Cargo.toml b/aimux-providers/Cargo.toml index bfca8c5d6..5bbd030f5 100644 --- a/aimux-providers/Cargo.toml +++ b/aimux-providers/Cargo.toml @@ -31,6 +31,7 @@ chrono = "0.4" url = { workspace = true } percent-encoding = "2" regex = "1" +ring = "0.17" [features] default = ["realtime"] diff --git a/aimux-providers/src/anthropic/cache_control.rs b/aimux-providers/src/anthropic/cache_control.rs index 31e82e210..faa2f35a2 100644 --- a/aimux-providers/src/anthropic/cache_control.rs +++ b/aimux-providers/src/anthropic/cache_control.rs @@ -37,6 +37,7 @@ pub(crate) fn extract_cache_control_for( let anthropic = anthropic_options(provider_options, options_name)?; anthropic .get("cacheControl") + .filter(|value| !value.is_null()) .or_else(|| anthropic.get("cache_control")) .cloned() } @@ -97,6 +98,9 @@ impl CacheControlValidator { let cache_control_value = extract_cache_control_for(provider_options, &self.options_name); let cc = cache_control_value?; + if cc.is_null() || cc == false || cc == 0 || cc == "" { + return None; + } if !can_cache { self.warnings.push(Warning::Unsupported { diff --git a/aimux-providers/src/anthropic/config.rs b/aimux-providers/src/anthropic/config.rs index d0b1a36d7..2e72a921b 100644 --- a/aimux-providers/src/anthropic/config.rs +++ b/aimux-providers/src/anthropic/config.rs @@ -150,8 +150,7 @@ impl AnthropicModelConfig { } /// Resolve the provider headers, add the `anthropic-beta` header for - /// `betas`, layer the per-call headers over both (case-insensitive, later - /// wins, `None` removes) and return the list that goes on the wire. + /// `betas`, merge beta values with per-call headers (case-insensitive) and return the list that goes on the wire. /// /// # Errors /// @@ -164,13 +163,6 @@ impl AnthropicModelConfig { betas: &BTreeSet, ) -> Result, AiMuxError> { let provider = self.headers.resolve().await?; - let mut beta_layer = HeaderMapOpt::new(); - if !betas.is_empty() { - beta_layer.insert( - "anthropic-beta".to_string(), - Some(betas.iter().cloned().collect::>().join(",")), - ); - } let call: HeaderMapOpt = call_headers .map(|headers| { headers @@ -179,10 +171,30 @@ impl AnthropicModelConfig { .collect() }) .unwrap_or_default(); + let mut merged_betas = betas.clone(); + for headers in [&provider, &call] { + for (name, value) in headers { + if name.eq_ignore_ascii_case("anthropic-beta") { + for beta in value.as_deref().unwrap_or_default().split(',') { + let beta = beta.trim().to_ascii_lowercase(); + if !beta.is_empty() { + merged_betas.insert(beta); + } + } + } + } + } + let mut beta_layer = HeaderMapOpt::new(); + if !merged_betas.is_empty() { + beta_layer.insert( + "anthropic-beta".to_string(), + Some(merged_betas.into_iter().collect::>().join(",")), + ); + } Ok(normalize_headers(combine_headers(&[ &provider, - &beta_layer, &call, + &beta_layer, ]))) } diff --git a/aimux-providers/src/anthropic/convert.rs b/aimux-providers/src/anthropic/convert.rs index 0a63c7ed5..123780d3d 100644 --- a/aimux-providers/src/anthropic/convert.rs +++ b/aimux-providers/src/anthropic/convert.rs @@ -25,17 +25,19 @@ use aimux_core::language_model_message::{ AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, }; -use aimux_core::options::{CallOptions, Tool}; +use aimux_core::options::{CallOptions, FunctionTool, ResponseFormat, Tool, ToolChoice}; use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; use aimux_core::types::{FinishReason, FinishReasonUnified, ReasoningEffort, Warning}; -use aimux_provider_utils::{get_top_level_media_type, resolve_full_media_type}; +use aimux_provider_utils::{ + MediaTypeData, detect_media_type, get_top_level_media_type, is_full_media_type, +}; use serde_json::{Map, Value, json}; use crate::anthropic::cache_control::CacheControlValidator; use crate::anthropic::options::{ CANONICAL, PROVIDER_SKILL_TYPE, anthropic_options, anthropic_options_in, }; -use crate::anthropic::prepare_tools::{AnthropicTool, prepare_tools_for}; +use crate::anthropic::prepare_tools::{AnthropicTool, prepare_tools_with_validator}; use crate::anthropic::tool_name_mapping::ToolNameMapping; /// Beta header emitted when a PDF file part is present. @@ -119,6 +121,23 @@ pub(crate) fn convert_prompt_for( send_reasoning: bool, tool_names: &ToolNameMapping, options_name: &str, +) -> Result { + let mut validator = CacheControlValidator::for_options_name(options_name); + convert_prompt_with_validator( + prompt, + send_reasoning, + tool_names, + options_name, + &mut validator, + ) +} + +fn convert_prompt_with_validator( + prompt: &LanguageModelPrompt, + send_reasoning: bool, + tool_names: &ToolNameMapping, + options_name: &str, + validator: &mut CacheControlValidator, ) -> Result { // Ids of tool calls that were executed over MCP. Their results must be sent // back as `mcp_tool_result`, not as a provider-tool result block. Upstream @@ -130,7 +149,6 @@ pub(crate) fn convert_prompt_for( let mut messages: Vec = Vec::new(); let mut betas: BTreeSet = BTreeSet::new(); let mut warnings: Vec = Vec::new(); - let mut validator = CacheControlValidator::for_options_name(options_name); #[derive(Clone, Copy, PartialEq, Eq)] enum Eff { @@ -150,25 +168,139 @@ pub(crate) fn convert_prompt_for( Eff::User => "user", Eff::Assistant => "assistant", }; - messages.push(json!({ "role": role_str, "content": std::mem::take(acc) })); + let mut content = std::mem::take(acc); + if role == Eff::Assistant { + let mut ordered = Vec::new(); + let mut segment = Vec::new(); + let flush_segment = |segment: &mut Vec, ordered: &mut Vec| { + let (tool_uses, other): (Vec<_>, Vec<_>) = + std::mem::take(segment).into_iter().partition(|part| { + part.get("type").and_then(Value::as_str) == Some("tool_use") + }); + ordered.extend(other); + ordered.extend(tool_uses); + }; + for part in content { + if matches!( + part.get("type").and_then(Value::as_str), + Some("thinking" | "redacted_thinking") + ) { + flush_segment(&mut segment, &mut ordered); + ordered.push(part); + } else { + segment.push(part); + } + } + flush_segment(&mut segment, &mut ordered); + content = ordered; + } + messages.push(json!({ "role": role_str, "content": content })); } } - for msg in prompt { + let mut has_system = false; + for (index, msg) in prompt.iter().enumerate() { let eff = match msg { - LanguageModelMessage::System { - content, - provider_options, - } => { + LanguageModelMessage::System { .. } => { + if index > 0 && matches!(prompt[index - 1], LanguageModelMessage::System { .. }) { + continue; + } flush(&mut messages, &mut last, &mut acc); - let cc = - validator.get_cache_control(provider_options.as_ref(), "system message", true); - let blocks = vec![apply_cc(json!({ "type": "text", "text": content }), cc)]; - if !seen_non_system { - system.extend(blocks); + let mut converted = Vec::new(); + for message in prompt[index..] + .iter() + .take_while(|message| matches!(message, LanguageModelMessage::System { .. })) + { + let LanguageModelMessage::System { + content, + provider_options, + } = message + else { + unreachable!() + }; + let options = anthropic_options(provider_options.as_ref(), options_name) + .unwrap_or_default(); + let invalid = + || AiMuxError::InvalidArgument("invalid anthropic provider options".into()); + let clear_at = options.get("clearAt"); + if clear_at.is_some_and(|value| value.as_str() != Some("next_user_message")) { + return Err(invalid()); + } + let effort = options.get("effort"); + if effort.is_some_and(|value| { + !matches!( + value.as_str(), + Some("low" | "medium" | "high" | "xhigh" | "max") + ) + }) { + return Err(invalid()); + } + let changes = match options.get("toolChanges") { + Some(value) => value.as_array().ok_or_else(invalid)?.as_slice(), + None => &[], + }; + let mut blocks = Vec::new(); + if !content.is_empty() + || (changes.is_empty() && clear_at.is_none() && effort.is_none()) + { + let cc = validator.get_cache_control( + provider_options.as_ref(), + "system message", + true, + ); + blocks.push(apply_cc(json!({ "type": "text", "text": content }), cc)); + } + for change in changes { + let kind = change + .get("type") + .and_then(Value::as_str) + .filter(|kind| matches!(*kind, "tool_addition" | "tool_removal")) + .ok_or_else(invalid)?; + let name = change + .get("toolName") + .and_then(Value::as_str) + .ok_or_else(invalid)?; + blocks.push(json!({ "type": kind, "tool": { + "type": "tool_reference", "name": tool_names.to_provider_tool_name(name) + }})); + } + converted.push((blocks, clear_at.cloned(), effort.cloned(), changes.len())); + } + let change_count: usize = converted.iter().map(|message| message.3).sum(); + let has_controls = converted + .iter() + .any(|message| message.1.is_some() || message.2.is_some()); + if !seen_non_system || (!has_system && change_count == 0 && !has_controls) { + if change_count > 0 { + warnings.push(Warning::Other { message: "tool changes on the initial system message are not supported by Anthropic. Configure the initial tool set via the tools option instead. The tool changes have been ignored.".into() }); + } + for (blocks, clear_at, effort, _) in converted { + if blocks.is_empty() && clear_at.is_none() && effort.is_some() { + messages.push(json!({ "role": "system", "content": [], "output_config": { "effort": effort } })); + betas.insert("mid-conversation-output-config-2026-07-01".into()); + } else if clear_at.is_some() || effort.is_some() { + warnings.push(Warning::Other { message: "clearAt and effort on this initial system message are not supported by Anthropic. Use a separate effort-only system message with empty content to set effort. These options have been ignored.".into() }); + } + system.extend(blocks.into_iter().filter(|block| block["type"] == "text")); + } + has_system = true; } else { - messages.push(json!({ "role": "system", "content": blocks })); betas.insert(BETA_MID_CONVERSATION_SYSTEM.to_string()); + for (blocks, clear_at, effort, changes) in converted { + let mut message = json!({ "role": "system", "content": blocks }); + if let Some(clear_at) = clear_at { + message["clear_at"] = clear_at; + betas.insert("mid-conversation-system-clear-at-2026-08-21".into()); + } + if let Some(effort) = effort { + message["output_config"] = json!({ "effort": effort }); + betas.insert("mid-conversation-output-config-2026-07-01".into()); + } + if changes > 0 { + betas.insert("mid-conversation-tool-changes-2026-07-01".into()); + } + messages.push(message); + } } continue; } @@ -193,7 +325,7 @@ pub(crate) fn convert_prompt_for( provider_options: part_options, }) => { let cc = resolve_cache_control( - &mut validator, + validator, part_options.as_ref(), provider_options.as_ref(), is_last, @@ -205,7 +337,7 @@ pub(crate) fn convert_prompt_for( UserPart::File(file) => convert_file_part( file, &mut betas, - &mut validator, + validator, provider_options.as_ref(), is_last, "user message part", @@ -226,11 +358,13 @@ pub(crate) fn convert_prompt_for( }; acc.push(convert_tool_result( part, - &mut validator, + validator, &mut betas, &mut warnings, idx + 1 == content.len(), provider_options.as_ref(), + options_name, + tool_names, )?); } } @@ -246,9 +380,10 @@ pub(crate) fn convert_prompt_for( tool_names, &mcp_tool_use_ids, &mut warnings, - &mut validator, + validator, is_last, provider_options.as_ref(), + options_name, ) } else { convert_assistant_part( @@ -257,7 +392,7 @@ pub(crate) fn convert_prompt_for( tool_names, &mut betas, &mut warnings, - &mut validator, + validator, is_last, provider_options.as_ref(), options_name, @@ -290,11 +425,7 @@ pub(crate) fn convert_prompt_for( // Merge any cache_control validation warnings. warnings.extend(validator.take_warnings()); - let system_opt = if system.is_empty() { - None - } else { - Some(system) - }; + let system_opt = has_system.then_some(system); Ok(AnthropicPromptConversion { system: system_opt, messages, @@ -391,7 +522,19 @@ fn convert_file_part( filename, provider_options, } = file; - let block = match data { + let resolve_inline_media_type = |data| -> Result { + if is_full_media_type(media_type) { + return Ok(media_type.clone()); + } + if let Some(detected) = detect_media_type(data, Some(get_top_level_media_type(media_type)))? + { + return Ok(detected.into()); + } + Err(AiMuxError::UnsupportedFunctionality(format!( + "file of media type \"{media_type}\" must specify subtype since it could not be auto-detected" + ))) + }; + let mut block = match data { FileData::Data { data: FileBytes::Binary(bytes), } => { @@ -399,7 +542,7 @@ fn convert_file_part( get_top_level_media_type(media_type), "image" | "application" ) { - resolve_full_media_type(file)? + resolve_inline_media_type(MediaTypeData::Bytes(bytes))? } else { media_type.clone() }; @@ -416,7 +559,7 @@ fn convert_file_part( get_top_level_media_type(media_type), "image" | "application" ) { - resolve_full_media_type(file)? + resolve_inline_media_type(MediaTypeData::Base64(data))? } else { media_type.clone() }; @@ -470,6 +613,26 @@ fn convert_file_part( block } }; + if block.get("type").and_then(Value::as_str) == Some("document") + && !matches!(data, FileData::Reference { .. }) + && let Some(metadata) = anthropic_options(provider_options.as_ref(), options_name) + { + for field in ["title", "context"] { + if let Some(value) = metadata.get(field).filter(|value| { + value.is_string() && (field == "title" || value.as_str() != Some("")) + }) { + block[field] = value.clone(); + } + } + if metadata + .get("citations") + .and_then(|v| v.get("enabled")) + .and_then(Value::as_bool) + == Some(true) + { + block["citations"] = json!({ "enabled": true }); + } + } let cc = resolve_cache_control( validator, provider_options.as_ref(), @@ -481,6 +644,7 @@ fn convert_file_part( Ok(apply_cc(block, cc)) } +#[allow(clippy::too_many_arguments)] fn convert_tool_result( part: &ToolResultPart, validator: &mut CacheControlValidator, @@ -488,29 +652,31 @@ fn convert_tool_result( warnings: &mut Vec, is_last_part: bool, message_provider_options: Option<&SharedProviderOptions>, + options_name: &str, + tool_names: &ToolNameMapping, ) -> Result { let ToolResultPart { tool_call_id, + tool_name, output, provider_options, .. } = part; - if let ToolResultOutput::Content { value } = output { - for item in value { - if let ToolResultContent::File(file) = item - && let FileData::Data { .. } = &file.data - { - resolve_full_media_type(file)?; - } - } - } - let (content, is_error) = resolve_tool_result_output(output, betas, warnings); + let (content, is_error) = resolve_tool_result_output(output, betas, warnings, options_name)?; let mut block = json!({ "type": "tool_result", "tool_use_id": tool_call_id, "content": content, }); + if let Some(name) = toolset_name( + tool_names, + tool_name, + provider_options.as_ref(), + options_name, + ) { + block["toolset_name"] = json!(name); + } if is_error { block["is_error"] = json!(true); } @@ -560,14 +726,31 @@ fn convert_assistant_part( "assistant message", ) }; - Ok(Some(match part { AssistantPart::Text(TextPart { text, provider_options, }) => { let cc = resolve_cc(validator, provider_options.as_ref()); - apply_cc(json!({ "type": "text", "text": text }), cc) + let metadata = anthropic_options(provider_options.as_ref(), options_name); + let mut block = json!({ "type": "text", "text": text }); + if let Some(metadata) = metadata { + if metadata.get("type").and_then(Value::as_str) == Some("compaction") { + if text.is_empty() { + return Ok(None); + } + block = json!({ "type": "compaction", "content": text }); + if let Some(signature) = + metadata.get("signature").filter(|value| value.is_string()) + { + block["signature"] = signature.clone(); + betas.insert("compact-2026-09-04".to_string()); + } + } else if let Some(citations) = metadata.get("citations").filter(|v| !v.is_null()) { + block["citations"] = citations.clone(); + } + } + apply_cc(block, cc) } AssistantPart::File(file) => convert_file_part( @@ -580,7 +763,6 @@ fn convert_assistant_part( "assistant message", options_name, )?, - AssistantPart::Reasoning(ReasoningPart { text, provider_options, @@ -673,15 +855,14 @@ fn convert_assistant_part( return Ok(None); }; - return Ok(Some(apply_cc( - json!({ - "type": "server_tool_use", - "id": tool_call_id, - "name": server_name, - "input": server_input, - }), - cc, - ))); + let mut block = json!({ + "type": "server_tool_use", "id": tool_call_id, + "name": server_name, "input": server_input, + }); + if let Some(caller) = tool_caller(provider_options.as_ref(), options_name) { + block["caller"] = caller; + } + return Ok(Some(apply_cc(block, cc))); } // Anthropic requires `input` to be a JSON object. The SDK wraps any @@ -692,24 +873,38 @@ fn convert_assistant_part( } else { json!({ "rawInvalidInput": input }) }; - let caller = anthropic_options(provider_options.as_ref(), options_name) - .as_ref() - .and_then(|anthropic| anthropic.get("caller")) - .and_then(|caller| { - let caller_type = caller.get("type")?.as_str()?; - match caller_type { - "code_execution_20250825" | "code_execution_20260120" => { - let tool_id = caller.get("toolId")?.as_str()?; - Some(json!({ "type": caller_type, "tool_id": tool_id })) - } - "direct" => Some(json!({ "type": "direct" })), - _ => None, - } + let caller = tool_caller(provider_options.as_ref(), options_name); + if let Some(toolset_name) = toolset_name( + tool_names, + tool_name, + provider_options.as_ref(), + options_name, + ) { + let Some(action) = input_val.get("action").and_then(Value::as_str) else { + warnings.push(Warning::Other { + message: format!( + "toolset tool call for tool {tool_name} is missing the action" + ), + }); + return Ok(None); + }; + let mut block = json!({ + "type": "tool_use", + "id": tool_call_id, + "name": action, + "toolset_name": toolset_name, + "input": input_val, }); + block["input"].as_object_mut().unwrap().remove("action"); + if let Some(caller) = caller { + block["caller"] = caller; + } + return Ok(Some(apply_cc(block, cc))); + } let mut block = json!({ "type": "tool_use", "id": tool_call_id, - "name": tool_names.to_provider_tool_name(tool_name), + "name": tool_name, "input": input_val, }); if let Some(caller) = caller { @@ -718,7 +913,7 @@ fn convert_assistant_part( apply_cc(block, cc) } - AssistantPart::Custom(_) | AssistantPart::ReasoningFile(_) => return Ok(None), + AssistantPart::ReasoningFile(_) | AssistantPart::Custom(_) => return Ok(None), AssistantPart::ToolResult(_) => unreachable!(), })) } @@ -759,6 +954,45 @@ fn collect_mcp_tool_use_ids<'a>( ids } +fn toolset_name( + tool_names: &ToolNameMapping, + tool_name: &str, + provider_options: Option<&SharedProviderOptions>, + options_name: &str, +) -> Option { + tool_names + .toolset_name(tool_name) + .map(str::to_string) + .or_else(|| { + anthropic_options(provider_options, options_name).and_then(|options| { + options + .get("toolsetName") + .and_then(Value::as_str) + .map(str::to_string) + }) + }) +} + +fn tool_caller( + provider_options: Option<&SharedProviderOptions>, + options_name: &str, +) -> Option { + anthropic_options(provider_options, options_name) + .as_ref() + .and_then(|anthropic| anthropic.get("caller")) + .and_then(|caller| { + let caller_type = caller.get("type")?.as_str()?; + match caller_type { + "code_execution_20250825" | "code_execution_20260120" => { + let tool_id = caller.get("toolId")?.as_str()?; + Some(json!({ "type": caller_type, "tool_id": tool_id })) + } + "direct" => Some(json!({ "type": "direct" })), + _ => None, + } + }) +} + /// Clone `key` out of `value`, defaulting to `null`. fn field(value: &Value, key: &str) -> Value { value.get(key).cloned().unwrap_or(Value::Null) @@ -770,15 +1004,19 @@ fn field(value: &Value, key: &str) -> Value { /// through unmapped keep `error_code`. Both are accepted so a replayed result /// round-trips either way. fn result_error_code(value: &Value, fallback: &str) -> String { + let parsed = value + .as_str() + .and_then(|value| serde_json::from_str::(value).ok()); + let value = parsed.as_ref().unwrap_or(value); value .get("errorCode") .or_else(|| value.get("error_code")) - .and_then(|c| c.as_str()) + .and_then(Value::as_str) .unwrap_or(fallback) .to_string() } -/// Convert an assistant-role `ToolResultPart` into the matching +/// Convert an assistant-role `AssistantPart::ToolResult` into the matching /// Anthropic provider-executed result block. /// /// Anthropic only accepts a bare `tool_result` block inside a **user** message; @@ -789,6 +1027,7 @@ fn result_error_code(value: &Value, fallback: &str) -> String { /// rather than sent as a bare `tool_result`. /// /// Returns `None` when the part is skipped. +#[allow(clippy::too_many_arguments)] fn convert_assistant_tool_result( part: &ToolResultPart, tool_names: &ToolNameMapping, @@ -797,6 +1036,7 @@ fn convert_assistant_tool_result( validator: &mut CacheControlValidator, is_last_part: bool, message_provider_options: Option<&SharedProviderOptions>, + options_name: &str, ) -> Option { let ToolResultPart { tool_call_id, @@ -806,24 +1046,14 @@ fn convert_assistant_tool_result( .. } = part; - // cache_control: part ?? output ?? (is_last_part ? message) — the same - // resolution order the bare `tool_result` path uses. - let cache_control = validator - .get_cache_control(provider_options.as_ref(), "assistant message part", true) - .or_else(|| { - validator.get_cache_control( - extract_tool_result_output_provider_options(output), - "tool result output", - true, - ) - }) - .or_else(|| { - is_last_part - .then(|| { - validator.get_cache_control(message_provider_options, "assistant message", true) - }) - .flatten() - }); + let cache_control = resolve_cache_control( + validator, + provider_options.as_ref(), + message_provider_options, + is_last_part, + "assistant message part", + "assistant message", + ); let raw_value = match output { ToolResultOutput::Json { value, .. } | ToolResultOutput::ErrorJson { value, .. } => { @@ -914,7 +1144,11 @@ fn convert_assistant_tool_result( ), "web_search" => ( "web_search_tool_result", - assistant_web_search_content(value), + if is_error { + json!({ "type": "web_search_tool_result_error", "error_code": result_error_code(value, "unavailable") }) + } else { + assistant_web_search_content(value) + }, ), "tool_search_tool_regex" | "tool_search_tool_bm25" => ( "tool_search_tool_result", @@ -940,6 +1174,11 @@ fn convert_assistant_tool_result( }) }; + if matches!(provider_tool_name, "web_fetch" | "web_search") + && let Some(caller) = tool_caller(provider_options.as_ref(), options_name) + { + block["caller"] = caller; + } if let Some(cc) = cache_control { block["cache_control"] = cc; } @@ -1329,63 +1568,100 @@ fn resolve_anthropic_reference( if let Some(id) = reference.get(CANONICAL) { return Ok(id.to_string()); } - let mut providers: Vec<&str> = reference.keys().map(String::as_str).collect(); - providers.sort_unstable(); + let providers: Vec<&str> = reference.keys().map(String::as_str).collect(); Err(AiMuxError::InvalidArgument(format!( "No provider reference found for provider '{CANONICAL}'. Available providers: {}", providers.join(", ") ))) } -/// Resolve a `ToolResultPart` `output` value into the Anthropic -/// `tool_result.content` (and whether it is an error), mirroring the TS SDK. -/// -/// The V4 `tool-result` `output` is a discriminated `{ type, value }` object: -/// - `json` → `JSON.stringify(value)` (a string), -/// - `text` → `value` as-is, -/// - `error` → `value` as-is with `is_error: true`, -/// - `content` → `value` as-is (an array of content blocks), -/// - anything else → the `output` is passed through unchanged. +/// Serialize a unified tool-result payload into the Anthropic content string. fn resolve_tool_result_output( output: &ToolResultOutput, betas: &mut BTreeSet, warnings: &mut Vec, -) -> (Value, bool) { - match output { - ToolResultOutput::Text { value, .. } => (json!(value), false), - ToolResultOutput::ErrorText { value, .. } => (json!(value), true), - ToolResultOutput::Json { value, .. } => (json!(value.to_string()), false), - ToolResultOutput::ErrorJson { value, .. } => (json!(value.to_string()), true), - ToolResultOutput::ExecutionDenied { reason, .. } => (json!(reason.as_deref().unwrap_or("Tool call execution denied.")), false), - ToolResultOutput::Content { value } => (json!(value.iter().filter_map(|part| match part { - ToolResultContent::Text(part) => Some(json!({ "type": "text", "text": part.text })), - ToolResultContent::File(part) => { - use base64::Engine; - let top_level = part.media_type.split('/').next(); - let source = match &part.data { - FileData::Url { url, .. } => json!({ "type": "url", "url": url }), - FileData::Data { data } => { - let data = match data { - FileBytes::Binary(bytes) => base64::engine::general_purpose::STANDARD.encode(bytes), - FileBytes::Base64(data) => data.clone(), + options_name: &str, +) -> Result<(Value, bool), AiMuxError> { + let value = match output { + ToolResultOutput::Text { value, .. } | ToolResultOutput::ErrorText { value, .. } => { + json!(value) + } + ToolResultOutput::Json { value, .. } | ToolResultOutput::ErrorJson { value, .. } => { + json!(value.to_string()) + } + ToolResultOutput::ExecutionDenied { reason, .. } => { + json!(reason.as_deref().unwrap_or("Tool call execution denied.")) + } + ToolResultOutput::Content { value } => { + let mut content = Vec::new(); + for part in value { + match part { + ToolResultContent::Text(part) => { + content.push(json!({ "type": "text", "text": part.text })) + } + ToolResultContent::File(part) => { + let is_image = get_top_level_media_type(&part.media_type) == "image"; + let source = match &part.data { + FileData::Url { url, .. } => json!({ "type": "url", "url": url }), + FileData::Data { data } => { + use base64::Engine; + let media_type = + aimux_provider_utils::resolve_full_media_type(part)?; + if !is_image && media_type != "application/pdf" { + warnings.push(Warning::Other { message: format!("unsupported tool content part type: file with media type: {}", part.media_type) }); + continue; + } + if !is_image { + betas.insert(BETA_PDFS.to_string()); + } + let data = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + json!({ "type": "base64", "media_type": media_type, "data": data }) + } + data => { + let data_type = match data { + FileData::Reference { .. } => "reference", + FileData::Text { .. } => "text", + _ => unreachable!(), + }; + warnings.push(Warning::Other { message: format!("unsupported tool content part type: file with data type: {data_type}") }); + continue; + } }; - let full = resolve_full_media_type(part).ok()?; - if top_level != Some("image") && full != "application/pdf" { warnings.push(Warning::Other { message: "unsupported tool content file media type".to_string() }); return None; } - if full == "application/pdf" { betas.insert("pdfs-2024-09-25".to_string()); } - json!({ "type": "base64", "media_type": full, "data": data }) - }, - _ => { warnings.push(Warning::Other { message: "unsupported tool content file part".to_string() }); return None; }, - }; - Some(json!({ "type": if top_level == Some("image") { "image" } else { "document" }, "source": source })) - }, - ToolResultContent::Custom { provider_options } => { - if let Some(options) = provider_options.as_ref().and_then(|options| options.get(CANONICAL)) - && options.get("type").and_then(Value::as_str) == Some("tool-reference") { - Some(json!({ "type": "tool_reference", "tool_name": options.get("toolName") })) - } else { warnings.push(Warning::Other { message: "unsupported custom tool content part".to_string() }); None } - }, - }).collect::>()), false), - } + content.push(json!({ "type": if is_image { "image" } else { "document" }, "source": source })); + } + ToolResultContent::Custom { provider_options } => { + if let Some(options) = + anthropic_options(provider_options.as_ref(), options_name) + && options.get("type").and_then(Value::as_str) == Some("tool-reference") + { + let mut part = json!({ "type": "tool_reference" }); + if let Some(tool_name) = options.get("toolName") { + part["tool_name"] = tool_name.clone(); + } + content.push(part); + } else { + warnings.push(Warning::Other { + message: "unsupported custom tool content part".to_string(), + }); + } + } + } + } + json!(content) + } + }; + Ok(( + value, + matches!( + output, + ToolResultOutput::ErrorText { .. } | ToolResultOutput::ErrorJson { .. } + ), + )) } // ── model capabilities & reasoning config ─────────────────────────────────── @@ -1399,6 +1675,9 @@ struct ModelCapabilities { rejects_sampling_parameters: bool, supports_xhigh_effort: bool, rejects_thinking_disabled_above_high_effort: bool, + rejects_thinking_disabled: bool, + rejects_forced_tool_use: bool, + supports_between_tools_thinking: bool, /// Whether the model supports native structured outputs (and therefore /// strict tool schemas). Mirrors the TS `supportsStructuredOutput` from /// `getModelCapabilities`; `supportsStrictTools` tracks it 1:1 because the @@ -1453,18 +1732,32 @@ fn get_model_capabilities(model_id: &str) -> ModelCapabilities { rejects_sampling_parameters, supports_xhigh_effort, rejects_thinking_disabled_above_high_effort, + rejects_thinking_disabled: false, + rejects_forced_tool_use: false, + supports_between_tools_thinking: false, supports_structured_output, is_known_model, } } - if model_id.contains("claude-opus-5") { - caps(128000, true, true, true, true, true, true) - } else if model_id.contains("claude-opus-4-8") - || model_id.contains("claude-opus-4-7") - || model_id.contains("claude-fable-5") - || model_id.contains("claude-sonnet-5") - { + if let Some((_, suffix)) = model_id.split_once("claude-opus-5") { + let mut result = caps(128000, true, true, true, true, true, true); + result.rejects_thinking_disabled = suffix.starts_with("-5"); + result.rejects_forced_tool_use = result.rejects_thinking_disabled; + result + } else if let Some((_, suffix)) = model_id.split_once("claude-fable-5") { + let mut result = caps(128000, true, true, true, false, true, true); + result.rejects_thinking_disabled = true; + result.rejects_forced_tool_use = suffix.starts_with("-1"); + result + } else if let Some((_, suffix)) = model_id.split_once("claude-sonnet-5") { + let mut result = caps(128000, true, true, true, false, true, true); + result.supports_between_tools_thinking = suffix.starts_with("-5"); + result.rejects_thinking_disabled_above_high_effort = result.supports_between_tools_thinking; + result.rejects_thinking_disabled = result.supports_between_tools_thinking; + result.rejects_forced_tool_use = result.supports_between_tools_thinking; + result + } else if model_id.contains("claude-opus-4-8") || model_id.contains("claude-opus-4-7") { caps(128000, true, true, true, false, true, true) } else if model_id.contains("claude-sonnet-4-6") || model_id.contains("claude-opus-4-6") { caps(128000, true, false, false, false, true, true) @@ -1610,7 +1903,7 @@ fn resolve_anthropic_reasoning_config( if supports_adaptive_thinking { let effort = map_reasoning_to_effort(reasoning, supports_xhigh_effort, warnings)?; return Some(ReasoningConfig { - thinking: json!({ "type": "adaptive" }), + thinking: json!({ "type": "adaptive", "display": "summarized" }), effort: Some(effort), }); } @@ -1626,6 +1919,7 @@ fn resolve_anthropic_reasoning_config( /// beta headers required by the request. #[derive(Debug, Clone)] pub struct RequestBodyResult { + pub(crate) uses_json_response_tool: bool, pub body: Value, pub warnings: Vec, /// Beta headers (e.g. `code-execution-2025-08-25`, `mcp-client-2025-04-04`) @@ -1932,6 +2226,38 @@ fn strip_anthropic_sampling_params( warnings: &mut Vec, ) -> (Option, Option, Option) { let mut temperature = options.temperature; + for (present, feature) in [ + (options.frequency_penalty.is_some(), "frequencyPenalty"), + (options.presence_penalty.is_some(), "presencePenalty"), + (options.seed.is_some(), "seed"), + ] { + if present { + warnings.push(Warning::Unsupported { + feature: feature.to_string(), + details: None, + }); + } + } + if let Some(value) = temperature + && !(0.0..=1.0).contains(&value) + { + let (limit, details) = if value > 1.0 { + ( + 1.0, + format!("{value} exceeds anthropic maximum of 1.0. clamped to 1.0"), + ) + } else { + ( + 0.0, + format!("{value} is below anthropic minimum of 0. clamped to 0"), + ) + }; + temperature = Some(limit); + warnings.push(Warning::Unsupported { + feature: "temperature".to_string(), + details: Some(details), + }); + } let mut top_p = options.top_p; let mut top_k = options.top_k; @@ -2014,6 +2340,61 @@ fn resolve_anthropic_thinking( } } + if options.reasoning == Some(ReasoningEffort::None) + && anthropic_option(&options.provider_options, profile, "thinking").is_none() + && anthropic_option(&options.provider_options, profile, "effort").is_none() + && caps.rejects_thinking_disabled + { + if caps.supports_between_tools_thinking { + thinking_config = Some(json!({ "type": "between_tools" })); + } else { + thinking_config = None; + effort = Some("low".to_string()); + warnings.push(Warning::Compatibility { feature: "reasoning".to_string(), details: Some(format!("reasoning 'none' is not supported by {model_id}; it always uses adaptive thinking. Using effort 'low' to minimize thinking instead.")) }); + } + } + + if caps.rejects_thinking_disabled { + match thinking_config + .as_ref() + .and_then(|v| v.get("type")) + .and_then(Value::as_str) + { + Some("disabled") => { + let details = if caps.supports_between_tools_thinking { + thinking_config = Some(json!({ "type": "between_tools" })); + format!( + "thinking cannot be disabled for {model_id}. Using 'between_tools' thinking, the lowest thinking setting, instead." + ) + } else { + thinking_config = None; + format!( + "thinking cannot be disabled for {model_id}; it always uses adaptive thinking. The thinking setting has been removed. Lower 'effort' to reduce thinking." + ) + }; + warnings.push(Warning::Unsupported { + feature: "providerOptions.anthropic.thinking".to_string(), + details: Some(details), + }); + } + Some("enabled") => { + thinking_config = Some(json!({ "type": "adaptive" })); + warnings.push(Warning::Unsupported { feature: "providerOptions.anthropic.thinking".to_string(), details: Some(format!("budget-based thinking is not supported by {model_id}; it always uses adaptive thinking. Using adaptive thinking instead. Use 'effort' to control how much the model thinks.")) }); + } + _ => {} + } + } + if thinking_config + .as_ref() + .and_then(|v| v.get("type")) + .and_then(Value::as_str) + == Some("between_tools") + && matches!(effort.as_deref(), Some("xhigh" | "max")) + { + warnings.push(Warning::Unsupported { feature: "providerOptions.anthropic.effort".to_string(), details: Some(format!("effort '{}' is not supported with 'between_tools' thinking. The effort has been lowered to 'high'.", effort.as_deref().unwrap_or_default())) }); + effort = Some("high".to_string()); + } + // Newer models only allow disabling thinking at effort ≤ high; lower the // effort to 'high' with a warning (TS L451-464). if caps.rejects_thinking_disabled_above_high_effort { @@ -2051,8 +2432,10 @@ fn derive_anthropic_thinking( .and_then(|t| t.as_str()) .map(std::string::ToString::to_string); - let is_thinking = - thinking_type.as_deref() == Some("enabled") || thinking_type.as_deref() == Some("adaptive"); + let is_thinking = matches!( + thinking_type.as_deref(), + Some("enabled" | "adaptive" | "between_tools") + ); let send_thinking = is_thinking || thinking_type.as_deref() == Some("disabled"); let thinking_budget: Option = if thinking_type.as_deref() == Some("enabled") { @@ -2120,8 +2503,10 @@ fn apply_anthropic_thinking_post_processing( top_p: &mut Option, warnings: &mut Vec, ) -> u32 { - let is_thinking = - thinking_type.as_deref() == Some("enabled") || thinking_type.as_deref() == Some("adaptive"); + let is_thinking = matches!( + thinking_type.as_deref(), + Some("enabled" | "adaptive" | "between_tools") + ); if !is_thinking { return max_tokens; } @@ -2164,7 +2549,7 @@ fn apply_anthropic_thinking_post_processing( *top_p = None; } - max_tokens + thinking_budget.unwrap_or(0) + max_tokens.saturating_add(thinking_budget.unwrap_or(0)) } /// Insert the surviving sampling params + stop sequences into the body. @@ -2192,17 +2577,20 @@ fn insert_anthropic_sampling( /// Map user tools to `tools` / `tool_choice`, delegating to /// `prepare_tools_with_provider` (provider-defined tools, required beta /// headers, and tool warnings). +#[allow(clippy::too_many_arguments)] fn apply_anthropic_tools( body: &mut Value, options: &CallOptions, profile: &RequestProfile, stream: bool, + json_schema: Option<&Value>, caps: &ModelCapabilities, betas: &mut BTreeSet, warnings: &mut Vec, + cache_validator: &mut CacheControlValidator, ) { - let disable_parallel_tool_use = - anthropic_option(&options.provider_options, profile, "disableParallelToolUse") + let disable_parallel_tool_use = json_schema.is_some() + || anthropic_option(&options.provider_options, profile, "disableParallelToolUse") .and_then(|v| v.as_bool()) .unwrap_or(false); let default_eager_input_streaming = stream @@ -2210,7 +2598,7 @@ fn apply_anthropic_tools( .and_then(|v| v.as_bool()) .unwrap_or(true); - let anthropic_tools: Vec = match &options.tools { + let mut anthropic_tools: Vec = match &options.tools { Some(tools) => tools .iter() .map(|t| match t { @@ -2225,18 +2613,50 @@ fn apply_anthropic_tools( None => Vec::new(), }; - let prepared = prepare_tools_for( - if options.tools.is_some() { + if let Some(schema) = json_schema { + anthropic_tools.push(AnthropicTool::Function( + FunctionTool::new("json", schema.clone()) + .with_description("Respond with a JSON object."), + )); + } + let requested_choice = if json_schema.is_some() { + Some(&ToolChoice::Required) + } else { + options.tool_choice.as_ref() + }; + let tool_choice = if caps.rejects_forced_tool_use + && !anthropic_tools.is_empty() + && matches!( + requested_choice, + Some(ToolChoice::Required | ToolChoice::Tool { .. }) + ) { + warnings.push(Warning::Unsupported { + feature: "toolChoice".to_string(), + details: Some(match requested_choice { + Some(ToolChoice::Required) => "toolChoice 'required' is not supported by this model because it rejects forced tool use. Using 'auto' instead. Instruct the model to use a tool in the prompt and verify that a tool call was made.".to_string(), + Some(ToolChoice::Tool { tool_name }) => format!("toolChoice 'tool' is not supported by this model because it rejects forced tool use. Only the '{tool_name}' tool is sent with 'auto' tool choice. Instruct the model to use the tool in the prompt and verify that a tool call was made."), + _ => unreachable!(), + }), + }); + Some(&ToolChoice::Auto) + } else { + requested_choice + }; + let prepared = prepare_tools_with_validator( + if options.tools.is_some() || json_schema.is_some() { Some(&anthropic_tools) } else { None }, - options.tool_choice.as_ref(), + tool_choice, disable_parallel_tool_use, - caps.supports_structured_output && profile.supports_native_structured_output, + json_schema.is_none() + && caps.supports_structured_output + && profile.supports_native_structured_output, caps.supports_structured_output && profile.supports_strict_tools, default_eager_input_streaming, &profile.options_name, + cache_validator, ); warnings.extend(prepared.tool_warnings); @@ -2247,12 +2667,26 @@ fn apply_anthropic_tools( // request body matches the TS behaviour (JSON.stringify drops // `undefined` args on provider-defined tools such as web_search). let mut defs = tool_defs; + if caps.rejects_forced_tool_use + && let Some(ToolChoice::Tool { tool_name }) = requested_choice + { + defs.retain(|tool| { + tool.get("name").and_then(Value::as_str) == Some(tool_name.as_str()) + }); + } for def in defs.iter_mut() { strip_null_fields(def); } body["tools"] = json!(defs); } - if let Some(tool_choice) = prepared.tool_choice { + if let Some(mut tool_choice) = prepared.tool_choice { + if json_schema.is_none() + && let Some(value) = + anthropic_option(&options.provider_options, profile, "disableParallelToolUse") + .filter(Value::is_boolean) + { + tool_choice["disable_parallel_tool_use"] = value; + } body["tool_choice"] = tool_choice; } } @@ -2475,19 +2909,17 @@ pub(crate) fn build_request_body_for( let max_tokens = options.max_output_tokens.unwrap_or(caps.max_output_tokens); - // Extended-thinking multi-turn: when thinking is enabled/adaptive, the - // prior assistant turns' reasoning parts (carrying their signatures) - // must be echoed back as thinking blocks — Anthropic rejects tool-use - // continuations that omit them, and the reasoning context is lost - // otherwise. With thinking disabled the blocks would be rejected, so - // the parts are omitted with a warning instead. let send_input_reasoning = - matches!(thinking_type.as_deref(), Some("enabled") | Some("adaptive")); - let conversion = convert_prompt_for( + anthropic_option(&options.provider_options, profile, "sendReasoning") + .and_then(|value| value.as_bool()) + .unwrap_or(true); + let mut cache_validator = CacheControlValidator::for_options_name(&profile.options_name); + let conversion = convert_prompt_with_validator( &options.prompt, send_input_reasoning, &ToolNameMapping::new(options.tools.as_deref()), &profile.options_name, + &mut cache_validator, )?; let system = conversion.system; let messages = conversion.messages; @@ -2517,9 +2949,18 @@ pub(crate) fn build_request_body_for( &thinking_effort, ); + if let Some(binding) = thinking_config.as_ref().and_then(|v| v.get("blockBinding")) { + if body.get("thinking").is_none() { + body["thinking"] = json!({}); + } + body["thinking"]["block_binding"] = + json!({ "prefix_mismatch_behavior": binding.get("prefixMismatchBehavior") }); + betas.insert("thinking-binding-controls-2026-08-01".to_string()); + } + // Thinking-enabled post-processing (TS L651-696): default budget, // sampling-parameter stripping, `max_tokens` adjustment. - let adjusted_max_tokens = apply_anthropic_thinking_post_processing( + let mut adjusted_max_tokens = apply_anthropic_thinking_post_processing( &mut body, &thinking_type, &mut thinking_budget, @@ -2529,12 +2970,87 @@ pub(crate) fn build_request_body_for( &mut top_p, &mut warnings, ); + if caps.is_known_model && adjusted_max_tokens > caps.max_output_tokens { + if options.max_output_tokens.is_some() { + warnings.push(Warning::Unsupported { + feature: "maxOutputTokens".to_string(), + details: Some(format!( + "{} (maxOutputTokens + thinkingBudget) is greater than {} {} max output tokens. The max output tokens have been limited to {}.", + adjusted_max_tokens, model_id, caps.max_output_tokens, caps.max_output_tokens + )), + }); + } + adjusted_max_tokens = caps.max_output_tokens; + } if adjusted_max_tokens != max_tokens { body["max_tokens"] = json!(adjusted_max_tokens); } + if !matches!( + thinking_type.as_deref(), + Some("enabled" | "adaptive" | "between_tools") + ) && (caps.is_known_model || is_legacy_claude(model_id)) + && temperature.is_some() + && top_p.is_some() + { + warnings.push(Warning::Unsupported { + feature: "topP".to_string(), + details: Some( + "topP is not supported when temperature is set. topP is ignored.".to_string(), + ), + }); + top_p = None; + } insert_anthropic_sampling(&mut body, temperature, top_p, top_k, options); + let mut json_tool_schema = None; + if let Some(ResponseFormat::Json { schema, .. }) = &options.response_format { + if let Some(schema) = schema { + let mode = anthropic_option(&options.provider_options, profile, "structuredOutputMode"); + let mut use_native = mode.as_ref().and_then(Value::as_str) == Some("outputFormat") + || (mode.as_ref().and_then(Value::as_str).unwrap_or("auto") == "auto" + && caps.supports_structured_output + && profile.supports_native_structured_output); + if !use_native + && caps.rejects_forced_tool_use + && caps.supports_structured_output + && profile.supports_native_structured_output + { + warnings.push(Warning::Unsupported { + feature: "providerOptions.anthropic.structuredOutputMode".to_string(), + details: Some(format!("structuredOutputMode 'jsonTool' is not supported by {model_id} because it rejects forced tool use. Using 'outputFormat' instead.")), + }); + use_native = true; + } + if !use_native { + json_tool_schema = Some(schema); + if anthropic_option(&options.provider_options, profile, "disableParallelToolUse") + .and_then(|v| v.as_bool()) + == Some(false) + { + warnings.push(Warning::Unsupported { + feature: "providerOptions.anthropic.disableParallelToolUse".to_string(), + details: Some("`disableParallelToolUse: false` is ignored when using the JSON response tool. Parallel tool use is disabled to ensure a single coherent JSON tool call.".to_string()), + }); + } + } + if use_native { + if body.get("output_config").is_none() { + body["output_config"] = json!({}); + } + body["output_config"]["format"] = json!({ "type": "json_schema", "schema": super::sanitize_json_schema::sanitize_json_schema(schema) }); + } + } else { + warnings.push(Warning::Unsupported { + feature: "responseFormat".to_string(), + details: Some( + "JSON response format requires a schema. The response format is ignored." + .to_string(), + ), + }); + } + } + // providerOptions.anthropic.metadata.userId -> metadata.user_id. if let Some(user_id) = anthropic_option(&options.provider_options, profile, "metadata") .and_then(|metadata| metadata.get("userId").cloned()) @@ -2550,10 +3066,13 @@ pub(crate) fn build_request_body_for( options, profile, stream, + json_tool_schema, &caps, &mut betas, &mut warnings, + &mut cache_validator, ); + warnings.extend(cache_validator.take_warnings()); // providerOptions.anthropic.mcpServers → mcp_servers + beta header. append_anthropic_mcp_servers(&mut body, options, profile, &mut betas); @@ -2561,10 +3080,110 @@ pub(crate) fn build_request_body_for( // providerOptions.anthropic.container — programmatic tool calling / skills. append_anthropic_container(&mut body, options, profile, &mut betas, &mut warnings)?; - // providerOptions.anthropic.contextManagement → context_management is not - // yet implemented (build_context_management pending). + let provider_options = + anthropic_options_in(options.provider_options.as_ref(), &profile.options_name); + if let Some(provider_options) = provider_options { + if provider_options + .get("compaction") + .is_some_and(|v| !v.is_null()) + && provider_options + .get("contextManagement") + .is_some_and(|v| !v.is_null()) + { + return Err(AiMuxError::InvalidArgument("Anthropic provider options `compaction` and `contextManagement` cannot be used together.".to_string())); + } + for (source, target) in [ + ("speed", "speed"), + ("serviceTier", "service_tier"), + ("inferenceGeo", "inference_geo"), + ("cacheControl", "cache_control"), + ("compaction", "compaction"), + ] { + if let Some(value) = provider_options.get(source).filter(|v| !v.is_null()) { + body[target] = value.clone(); + } + } + if let Some(value) = provider_options.get("fallbacks").filter(|v| { + v.as_str() == Some("default") || v.as_array().is_some_and(|a| !a.is_empty()) + }) { + body["fallbacks"] = value.clone(); + betas.insert( + if value.as_str() == Some("default") { + "server-side-fallback-2026-07-01" + } else { + "server-side-fallback-2026-06-01" + } + .to_string(), + ); + } + if let Some(value) = provider_options.get("taskBudget").filter(|v| !v.is_null()) { + if body.get("output_config").is_none() { + body["output_config"] = json!({}); + } + body["output_config"]["task_budget"] = value.clone(); + betas.insert("task-budgets-2026-03-13".to_string()); + } + if body.get("speed").and_then(Value::as_str) == Some("fast") { + betas.insert("fast-mode-2026-02-01".to_string()); + } + if body.get("compaction").is_some() { + betas.insert("compact-2026-09-04".to_string()); + } + if let Some(edits) = provider_options + .get("contextManagement") + .and_then(|v| v.get("edits")) + .and_then(Value::as_array) + { + let mut mapped = Vec::new(); + for edit in edits { + let strategy = edit.get("type").and_then(Value::as_str).unwrap_or_default(); + let fields: &[(&str, &str)] = match strategy { + "clear_tool_uses_20250919" => &[ + ("trigger", "trigger"), + ("keep", "keep"), + ("clearAtLeast", "clear_at_least"), + ("clearToolInputs", "clear_tool_inputs"), + ("excludeTools", "exclude_tools"), + ], + "clear_thinking_20251015" => &[("keep", "keep")], + "compact_20260112" => { + betas.insert("compact-2026-01-12".to_string()); + &[ + ("trigger", "trigger"), + ("pauseAfterCompaction", "pause_after_compaction"), + ("instructions", "instructions"), + ] + } + _ => { + warnings.push(Warning::Other { + message: format!("Unknown context management strategy: {strategy}"), + }); + continue; + } + }; + let mut mapped_edit = json!({ "type": strategy }); + for (source, target) in fields { + if let Some(value) = edit.get(*source) { + mapped_edit[*target] = value.clone(); + } + } + mapped.push(mapped_edit); + } + body["context_management"] = json!({ "edits": mapped }); + betas.insert("context-management-2025-06-27".to_string()); + } + } + if let Some(values) = anthropic_option(&options.provider_options, profile, "anthropicBeta") + .and_then(|value| value.as_array().cloned()) + { + betas.extend(values.iter().filter_map(Value::as_str).map(str::to_string)); + } + if thinking_display.as_ref().and_then(Value::as_str) == Some("updates") { + betas.insert("thinking-display-updates-2026-08-18".to_string()); + } Ok(RequestBodyResult { + uses_json_response_tool: json_tool_schema.is_some(), body, warnings, betas, @@ -2575,8 +3194,9 @@ pub(crate) fn build_request_body_for( #[must_use] pub fn parse_stop_reason(s: &str) -> FinishReason { let unified = match s { - "end_turn" => FinishReasonUnified::Stop, - "max_tokens" => FinishReasonUnified::Length, + "end_turn" | "stop_sequence" | "pause_turn" => FinishReasonUnified::Stop, + "max_tokens" | "model_context_window_exceeded" => FinishReasonUnified::Length, + "refusal" => FinishReasonUnified::ContentFilter, "tool_use" => FinishReasonUnified::ToolCalls, _ => FinishReasonUnified::Other, }; diff --git a/aimux-providers/src/anthropic/files.rs b/aimux-providers/src/anthropic/files.rs index 59ee58cb1..e728b99a9 100644 --- a/aimux-providers/src/anthropic/files.rs +++ b/aimux-providers/src/anthropic/files.rs @@ -70,7 +70,7 @@ impl AnthropicFiles { #[async_trait] impl Files for AnthropicFiles { - /// `"{name}.files"` with the `.messages` suffix removed. + /// The configured provider name. fn provider(&self) -> &str { &self.config.provider } @@ -96,7 +96,10 @@ impl Files for AnthropicFiles { let betas = std::collections::BTreeSet::from([FILES_BETA_HEADER.to_string()]); let config = self.config.resolved().await?; - let header_list = config.request_headers(None, &betas).await?; + let mut header_list = config.request_headers(None, &betas).await?; + // The Files endpoint replaces the provider beta header with its own. + header_list.retain(|(name, _)| !name.eq_ignore_ascii_case("anthropic-beta")); + header_list.push(("anthropic-beta".to_string(), FILES_BETA_HEADER.to_string())); // `send()` returns Ok only for 2xx; non-2xx responses are mapped to an // error internally using the shared error structure. The multipart body @@ -148,9 +151,7 @@ impl Files for AnthropicFiles { provider_reference: provider_ref, media_type: Some(result_media_type), filename: result_filename, - provider_metadata: Some( - std::iter::once((self.config.provider_options_name.clone(), metadata)).collect(), - ), + provider_metadata: Some(std::iter::once((CANONICAL.to_string(), metadata)).collect()), warnings: Vec::new(), }) } diff --git a/aimux-providers/src/anthropic/mod.rs b/aimux-providers/src/anthropic/mod.rs index d3a96fc5f..8351f7629 100644 --- a/aimux-providers/src/anthropic/mod.rs +++ b/aimux-providers/src/anthropic/mod.rs @@ -74,6 +74,11 @@ const DEFAULT_NAME: &str = "anthropic.messages"; /// The first-party URL without its version segment means the versioned one /// (`normalizeBaseURL` in `createAnthropic`); anything else is taken as given. fn normalize_base_url(url: &str) -> Result { + if url.trim().is_empty() { + return Err(AiMuxError::InvalidArgument( + "baseURL must be a non-empty string.".into(), + )); + } let url = validate_base_url(url)?; Ok(if url == API_ORIGIN { DEFAULT_BASE_URL.to_string() @@ -103,7 +108,7 @@ fn supported_urls() -> SupportedUrls { #[derive(Clone, Default)] pub struct AnthropicProviderSettings { /// Base URL for the API calls, version segment included. Default - /// `https://api.anthropic.com/v1`; a trailing slash is removed. The bare + /// `ANTHROPIC_BASE_URL` or `https://api.anthropic.com/v1`; a trailing slash is removed. The bare /// `https://api.anthropic.com` means the default. pub base_url: Option, /// The API key, sent as `x-api-key`. `None` (with no `auth_token`) loads @@ -111,8 +116,9 @@ pub struct AnthropicProviderSettings { /// `AiMuxError::LoadApiKey` if it is unset. An explicit value is used as /// given, `""` included: it never falls back to the environment. pub api_key: Option, - /// A non-empty bearer token is sent as `Authorization: Bearer` instead - /// of `x-api-key`. Giving non-empty values for both credentials is an error. + /// A bearer token, sent as `Authorization: Bearer`. When set, no + /// `x-api-key` header is sent. Giving it together with `api_key` is an + /// error. pub auth_token: Option, /// Extra headers on every request. A `None` value removes the header, /// including one of the fixed ones. Per-call headers win over these. @@ -125,7 +131,7 @@ pub struct AnthropicProviderSettings { /// The transport: a mock, a signing decorator, a proxy-aware client. /// `None` uses the process default, resolved per request. pub fetch: Option, - /// Generator for source IDs. Defaults to the provider utility generator. + /// Generates identifiers for returned sources. pub generate_id: Option String + Send + Sync>>, } @@ -151,8 +157,8 @@ impl std::fmt::Debug for AnthropicProviderSettings { /// /// # Errors /// -/// Returns `AiMuxError::InvalidArgument` when `base_url` is not an `http(s)` -/// URL with a host, or when both `api_key` and `auth_token` are given. Those +/// Returns `AiMuxError::InvalidArgument` when the configured base URL is not an +/// `http(s)` URL with a host, or when both credentials are nonempty. Those /// are the only ways this fails: credentials are loaded per request, not here. pub fn create_anthropic( settings: AnthropicProviderSettings, @@ -170,7 +176,8 @@ pub fn create_anthropic( .to_string(), )); } - let (credential, scheme) = match settings.auth_token.filter(is_truthy) { + let auth_token = settings.auth_token.filter(is_truthy); + let (credential, scheme) = match auth_token { Some(token) => ( Credential::Explicit(Resolvable::Value(token)), AuthScheme::Bearer, @@ -194,21 +201,24 @@ pub fn create_anthropic( vec![("anthropic-version".to_string(), API_VERSION.to_string())], settings.headers, ), - crate::anthropic::options::CANONICAL, + options::CANONICAL, "4.0.68", ), fetch: settings.fetch, - generate_id: settings.generate_id, supported_urls: supported_urls(), + generate_id: settings.generate_id, }) } /// The default provider: `create_anthropic` with default settings, created on -/// first use. The base URL is loaded at creation; the key is loaded per request. +/// first use. Creation reads `ANTHROPIC_BASE_URL`; a missing key surfaces +/// from the first request instead. pub fn anthropic() -> &'static AnthropicProvider { static DEFAULT: OnceLock = OnceLock::new(); DEFAULT.get_or_init(|| { - create_anthropic(AnthropicProviderSettings::default()).expect("invalid ANTHROPIC_BASE_URL") + // An invalid base URL environment setting fails at creation. + create_anthropic(AnthropicProviderSettings::default()) + .expect("Anthropic base URL must be valid") }) } @@ -219,8 +229,8 @@ pub struct AnthropicProvider { base_url: String, headers: HeadersFn, fetch: Option, - generate_id: Option String + Send + Sync>>, supported_urls: SupportedUrls, + generate_id: Option String + Send + Sync>>, } impl AnthropicProvider { @@ -258,7 +268,13 @@ impl AnthropicProvider { .with_generate_id(self.generate_id.clone()) } - /// The files interface; `provider()` is the provider name. + /// An alias for the Messages model. + #[must_use] + pub fn chat(&self, model_id: &str) -> AnthropicMessagesModel { + self.messages(model_id) + } + + /// The files interface, retaining the configured provider name. #[must_use] pub fn files(&self) -> AnthropicFiles { AnthropicFiles::from_config(self.model_config(self.name.clone())) diff --git a/aimux-providers/src/anthropic/model.rs b/aimux-providers/src/anthropic/model.rs index efd4edac3..bc157bbae 100644 --- a/aimux-providers/src/anthropic/model.rs +++ b/aimux-providers/src/anthropic/model.rs @@ -6,11 +6,14 @@ //! and hands the result to the shared [`super::stream`] core, which sends it //! and parses the response or the SSE stream. +use std::sync::Arc; + use async_trait::async_trait; use serde_json::Value; use aimux_core::error::AiMuxError; use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::language_model_message::{LanguageModelMessage, UserPart}; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateResult, StreamResult}; use aimux_core::types::Warning; @@ -18,7 +21,8 @@ use aimux_provider_utils::HttpRequest; use super::config::AnthropicModelConfig; use super::convert::build_request_body_for; -use super::stream::{anthropic_generate_core, anthropic_stream_core}; +use super::options::CANONICAL; +use super::stream::{CitationDocument, anthropic_generate_core, anthropic_stream_core}; use super::tool_name_mapping::ToolNameMapping; /// One call, ready to send. @@ -26,6 +30,7 @@ struct PreparedCall { http: HttpRequest, body: Value, warnings: Vec, + uses_json_response_tool: bool, } /// An Anthropic Messages model (e.g. `claude-sonnet-4-20250514`), created by @@ -33,7 +38,7 @@ struct PreparedCall { pub struct AnthropicMessagesModel { model_id: String, config: AnthropicModelConfig, - generate_id: std::sync::Arc String + Send + Sync>, + generate_id: Arc String + Send + Sync>, } impl AnthropicMessagesModel { @@ -42,13 +47,14 @@ impl AnthropicMessagesModel { Self { model_id, config, - generate_id: std::sync::Arc::new(aimux_provider_utils::generate_id), + generate_id: Arc::new(aimux_provider_utils::generate_id), } } + #[must_use] pub(crate) fn with_generate_id( mut self, - generate_id: Option String + Send + Sync>>, + generate_id: Option String + Send + Sync>>, ) -> Self { if let Some(generate_id) = generate_id { self.generate_id = generate_id; @@ -56,6 +62,49 @@ impl AnthropicMessagesModel { self } + fn citation_documents(options: &CallOptions) -> Vec { + options + .prompt + .iter() + .filter_map(|message| match message { + LanguageModelMessage::User { content, .. } => Some(content), + _ => None, + }) + .flatten() + .filter_map(|part| match part { + UserPart::File(file) + if matches!(file.media_type.as_str(), "application/pdf" | "text/plain") + && file + .provider_options + .as_ref() + .and_then(|options| options.get(CANONICAL)) + .and_then(|options| options.get("citations")) + .and_then(|citations| citations.get("enabled")) + .and_then(Value::as_bool) + == Some(true) => + { + Some(CitationDocument { + title: file + .filename + .clone() + .unwrap_or_else(|| "Untitled Document".to_string()), + filename: file.filename.clone(), + media_type: file.media_type.clone(), + }) + } + _ => None, + }) + .collect() + } + + fn uses_custom_options(&self, options: &CallOptions) -> bool { + self.config.provider_options_name != CANONICAL + && options + .provider_options + .as_ref() + .is_some_and(|value| value.contains_key(&self.config.provider_options_name)) + } + /// Build the request for one call: the body (host preparation and the /// provider's `transform_request_body` applied), its warnings, and the /// URL/headers/transport the exchange goes through. @@ -80,6 +129,7 @@ impl AnthropicMessagesModel { http, body, warnings: built.warnings, + uses_json_response_tool: built.uses_json_response_tool, }) } } @@ -106,12 +156,11 @@ impl LanguageModel for AnthropicMessagesModel { call.body, call.warnings, &self.config, - options - .provider_options - .as_ref() - .is_some_and(|options| options.contains_key(&self.config.provider_options_name)), &ToolNameMapping::new(options.tools.as_deref()), - self.generate_id.as_ref(), + self.uses_custom_options(options), + call.uses_json_response_tool, + Self::citation_documents(options), + self.generate_id.clone(), ) .await } @@ -123,11 +172,10 @@ impl LanguageModel for AnthropicMessagesModel { call.body, call.warnings, &self.config, - options - .provider_options - .as_ref() - .is_some_and(|options| options.contains_key(&self.config.provider_options_name)), ToolNameMapping::new(options.tools.as_deref()), + self.uses_custom_options(options), + call.uses_json_response_tool, + Self::citation_documents(options), self.generate_id.clone(), ) .await diff --git a/aimux-providers/src/anthropic/prepare_tools.rs b/aimux-providers/src/anthropic/prepare_tools.rs index c7ff89a33..6954fa0f5 100644 --- a/aimux-providers/src/anthropic/prepare_tools.rs +++ b/aimux-providers/src/anthropic/prepare_tools.rs @@ -5,8 +5,6 @@ //! plus provider-defined tools (computer use, web search, code execution, ...) //! via [`prepare_tools_with_provider`]. //! -//! Cache-control resolution is out of scope: it would require a -//! `CacheControlValidator` plus provider options on individual parts. use std::collections::BTreeSet; @@ -15,6 +13,7 @@ use aimux_core::tool::ToolChoice; use aimux_core::types::Warning; use serde_json::{Map, Value, json}; +use super::cache_control::CacheControlValidator; use super::options::{CANONICAL, anthropic_options_in}; /// Result of [`prepare_tools`] / [`prepare_tools_with_provider`]. @@ -48,63 +47,21 @@ pub fn prepare_tools( supports_strict_tools: bool, default_eager_input_streaming: bool, ) -> PreparedTools { - // Empty arrays are coerced to "no tools" to match the TS behaviour. - let non_empty = tools.filter(|&t| !t.is_empty()); - - let mut tool_warnings: Vec = Vec::new(); - let mut betas: BTreeSet = BTreeSet::new(); - - let tools_opt = match non_empty { - None => None, - Some(tools) => { - let mut anthropic_tools: Vec = Vec::new(); - for tool in tools { - anthropic_tools.push(prepare_function_tool( - tool, - supports_structured_output, - supports_strict_tools, - default_eager_input_streaming, - CANONICAL, - &mut betas, - &mut tool_warnings, - )); - } - Some(anthropic_tools) - } - }; - - // If the tools were dropped (e.g. there never were any), there is no - // tool_choice to emit either. - let tools_opt = match tools_opt { - None => { - return PreparedTools { - tools: None, - tool_choice: None, - tool_warnings, - betas, - }; - } - Some(t) => Some(t), - }; - - // Anthropic does not support 'none' tool choice, so the tools are removed. - if matches!(tool_choice, Some(ToolChoice::None)) { - return PreparedTools { - tools: None, - tool_choice: None, - tool_warnings, - betas, - }; - } - - let tool_choice_opt = build_tool_choice(tool_choice, disable_parallel_tool_use); - - PreparedTools { - tools: tools_opt, - tool_choice: tool_choice_opt, - tool_warnings, - betas, - } + let tools = tools.map(|tools| { + tools + .iter() + .cloned() + .map(AnthropicTool::Function) + .collect::>() + }); + prepare_tools_with_provider( + tools.as_deref(), + tool_choice, + disable_parallel_tool_use, + supports_structured_output, + supports_strict_tools, + default_eager_input_streaming, + ) } // ============================================================================= @@ -169,6 +126,32 @@ pub(crate) fn prepare_tools_for( supports_strict_tools: bool, default_eager_input_streaming: bool, options_name: &str, +) -> PreparedTools { + let mut validator = CacheControlValidator::for_options_name(options_name); + let mut result = prepare_tools_with_validator( + tools, + tool_choice, + disable_parallel_tool_use, + supports_structured_output, + supports_strict_tools, + default_eager_input_streaming, + options_name, + &mut validator, + ); + result.tool_warnings.extend(validator.take_warnings()); + result +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn prepare_tools_with_validator( + tools: Option<&[AnthropicTool]>, + tool_choice: Option<&ToolChoice>, + disable_parallel_tool_use: bool, + supports_structured_output: bool, + supports_strict_tools: bool, + default_eager_input_streaming: bool, + options_name: &str, + validator: &mut CacheControlValidator, ) -> PreparedTools { // Empty arrays are coerced to "no tools" to match the TS behaviour. let non_empty = tools.filter(|&t| !t.is_empty()); @@ -183,7 +166,7 @@ pub(crate) fn prepare_tools_for( for tool in tools { match tool { AnthropicTool::Function(ft) => { - anthropic_tools.push(prepare_function_tool( + let mut def = prepare_function_tool( ft, supports_structured_output, supports_strict_tools, @@ -191,7 +174,15 @@ pub(crate) fn prepare_tools_for( options_name, &mut betas, &mut tool_warnings, - )); + ); + if let Some(cache_control) = validator.get_cache_control( + ft.provider_options.as_ref(), + "tool definition", + true, + ) { + def["cache_control"] = cache_control; + } + anthropic_tools.push(def); } AnthropicTool::Provider { id, name: _, args } => { if let Some(def) = prepare_provider_tool(id, args, &mut betas) { @@ -350,7 +341,7 @@ pub(crate) fn prepare_provider_tool( args: &Map, betas: &mut BTreeSet, ) -> Option { - Some(match id { + let mut def = match id { "anthropic.computer_20241022" => { betas.insert("computer-use-2024-10-22".to_string()); json!({ @@ -382,6 +373,29 @@ pub(crate) fn prepare_provider_tool( "enable_zoom": arg(args, "enableZoom"), }) } + "anthropic.computer_toolset_20260801" => { + let mut def = json!({ "type": "computer_toolset_20260801" }); + if let Some(configs) = args.get("configs").and_then(Value::as_object) { + let mapped: serde_json::Map = configs + .iter() + .map(|(member, config)| { + let mut value = json!({}); + for (source, target) in + [("enabled", "enabled"), ("deferLoading", "defer_loading")] + { + if let Some(option) = + config.get(source).filter(|value| !value.is_null()) + { + value[target] = option.clone(); + } + } + (member.clone(), value) + }) + .collect(); + def["configs"] = Value::Object(mapped); + } + def + } "anthropic.text_editor_20241022" => { betas.insert("computer-use-2024-10-22".to_string()); json!({ "name": "str_replace_editor", "type": "text_editor_20241022" }) @@ -469,6 +483,24 @@ pub(crate) fn prepare_provider_tool( "max_content_tokens": arg(args, "maxContentTokens"), }) } + "anthropic.web_search_20260318" => json!({ + "type": "web_search_20260318", "name": "web_search", + "max_uses": arg(args, "maxUses"), + "allowed_domains": arg(args, "allowedDomains"), + "blocked_domains": arg(args, "blockedDomains"), + "user_location": arg(args, "userLocation"), + "response_inclusion": arg(args, "responseInclusion"), + }), + "anthropic.web_fetch_20260318" => json!({ + "type": "web_fetch_20260318", "name": "web_fetch", + "max_uses": arg(args, "maxUses"), + "allowed_domains": arg(args, "allowedDomains"), + "blocked_domains": arg(args, "blockedDomains"), + "citations": arg(args, "citations"), + "max_content_tokens": arg(args, "maxContentTokens"), + "use_cache": arg(args, "useCache"), + "response_inclusion": arg(args, "responseInclusion"), + }), "anthropic.tool_search_regex_20251119" => { json!({ "name": "tool_search_tool_regex", "type": "tool_search_tool_regex_20251119" }) } @@ -485,11 +517,18 @@ pub(crate) fn prepare_provider_tool( if args.get("maxUses").is_some() { def["max_uses"] = arg(args, "maxUses"); } + if args.get("maxTokens").is_some() { + def["max_tokens"] = arg(args, "maxTokens"); + } if args.get("caching").is_some() { def["caching"] = arg(args, "caching"); } def } _ => return None, - }) + }; + if let Some(fields) = def.as_object_mut() { + fields.retain(|_, value| !value.is_null()); + } + Some(def) } diff --git a/aimux-providers/src/anthropic/skills.rs b/aimux-providers/src/anthropic/skills.rs index b10fffb00..f4c7cb871 100644 --- a/aimux-providers/src/anthropic/skills.rs +++ b/aimux-providers/src/anthropic/skills.rs @@ -124,9 +124,14 @@ impl Skills for AnthropicSkills { } body.extend_from_slice(format!("--{boundary}--\r\n").as_bytes()); let config = self.config.resolved().await?; - let headers = config + let mut headers = config .request_headers(None, &BTreeSet::from(["skills-2025-10-02".to_string()])) .await?; + headers.retain(|(name, _)| !name.eq_ignore_ascii_case("anthropic-beta")); + headers.push(( + "anthropic-beta".to_string(), + "skills-2025-10-02".to_string(), + )); let response = aimux_provider_utils::post_to_api( config.with_transport(HttpRequest { url: config.url("/skills"), diff --git a/aimux-providers/src/anthropic/stream.rs b/aimux-providers/src/anthropic/stream.rs index 4ba960f05..1449bbd70 100644 --- a/aimux-providers/src/anthropic/stream.rs +++ b/aimux-providers/src/anthropic/stream.rs @@ -29,6 +29,7 @@ use serde_json::{Value, json}; use super::config::AnthropicModelConfig; use super::convert::parse_stop_reason; +use super::options::CANONICAL; use super::tool_name_mapping::ToolNameMapping; use super::types::{AnthropicResponse, ContentBlock, StreamErrorData, StreamEvent, ToolCallCaller}; @@ -72,13 +73,16 @@ fn str_field(v: &Value, key: &str) -> Option { /// `web_search_result` → the camel-cased shape upstream exposes (:1243-1249). fn map_web_search_result(result: &Value) -> Value { - json!({ - "url": result.get("url").cloned().unwrap_or(Value::Null), - "title": result.get("title").cloned().unwrap_or(Value::Null), + let mut mapped = json!({ + "url": result["url"], "pageAge": result.get("page_age").cloned().unwrap_or(Value::Null), - "encryptedContent": result.get("encrypted_content").cloned().unwrap_or(Value::Null), - "type": result.get("type").cloned().unwrap_or(Value::Null), - }) + "encryptedContent": result["encrypted_content"], + "type": result["type"], + }); + if let Some(title) = result.get("title").filter(|title| !title.is_null()) { + mapped["title"] = title.clone(); + } + mapped } /// `web_fetch_tool_result.content` → `(result, is_error)` (upstream :1196-1231). @@ -301,6 +305,28 @@ pub(crate) fn tool_call_caller_metadata( ) } +fn toolset_member_input(member_name: &str, input: &Value) -> Value { + let mut value = json!({ "action": member_name }); + if let Some(input) = input.as_object() { + value.as_object_mut().unwrap().extend(input.clone()); + } + value +} + +fn tool_call_metadata( + caller: Option<&ToolCallCaller>, + toolset_name: Option<&str>, + options_name: &str, +) -> Option { + let mut metadata = tool_call_caller_metadata(caller, options_name); + if let Some(name) = toolset_name { + let metadata = + metadata.get_or_insert_with(|| provider_namespace(options_name, json!({})).unwrap()); + metadata.get_mut(options_name).unwrap()["toolsetName"] = json!(name); + } + metadata +} + fn is_tool_search_provider_name(name: &str) -> bool { matches!(name, "tool_search_tool_regex" | "tool_search_tool_bm25") } @@ -340,6 +366,11 @@ pub(crate) fn stream_parts_for_result_block( server_tool_calls: &HashMap, generate_id: &(dyn Fn() -> String + Send + Sync), ) -> Vec { + let caller = match block { + ContentBlock::WebSearchToolResult { caller, .. } + | ContentBlock::WebFetchToolResult { caller, .. } => caller.as_ref(), + _ => None, + }; let tool_result = |tool_name: String, (result, is_error): (Value, Option), tool_use_id: &str| { StreamPart::ToolResult(ToolResult { @@ -349,7 +380,7 @@ pub(crate) fn stream_parts_for_result_block( is_error, preliminary: None, dynamic: None, - provider_metadata: None, + provider_metadata: tool_call_caller_metadata(caller, CANONICAL), }) }; @@ -357,6 +388,7 @@ pub(crate) fn stream_parts_for_result_block( ContentBlock::WebSearchToolResult { tool_use_id, content, + .. } => { let name = names.to_custom_tool_name("web_search").to_string(); let Some(results) = content.as_array() else { @@ -370,7 +402,7 @@ pub(crate) fn stream_parts_for_result_block( is_error: Some(true), preliminary: None, dynamic: None, - provider_metadata: None, + provider_metadata: tool_call_caller_metadata(caller, CANONICAL), })]; }; // The tool result, then one Source per hit — the same pair the @@ -404,6 +436,7 @@ pub(crate) fn stream_parts_for_result_block( ContentBlock::WebFetchToolResult { tool_use_id, content, + .. } => vec![tool_result( names.to_custom_tool_name("web_fetch").to_string(), map_web_fetch_result(content), @@ -412,6 +445,7 @@ pub(crate) fn stream_parts_for_result_block( ContentBlock::CodeExecutionToolResult { tool_use_id, content, + .. } => vec![tool_result( names.to_custom_tool_name("code_execution").to_string(), map_code_execution_result(content), @@ -424,6 +458,7 @@ pub(crate) fn stream_parts_for_result_block( | ContentBlock::TextEditorCodeExecutionToolResult { tool_use_id, content, + .. } => vec![tool_result( names.to_custom_tool_name("code_execution").to_string(), (content.clone(), None), @@ -432,6 +467,7 @@ pub(crate) fn stream_parts_for_result_block( ContentBlock::ToolSearchToolResult { tool_use_id, content, + .. } => vec![tool_result( names .to_custom_tool_name(tool_search_provider_name( @@ -445,6 +481,7 @@ pub(crate) fn stream_parts_for_result_block( ContentBlock::AdvisorToolResult { tool_use_id, content, + .. } => vec![tool_result( names.to_custom_tool_name("advisor").to_string(), map_advisor_result(content), @@ -476,6 +513,95 @@ pub(crate) fn stream_parts_for_result_block( } } +fn citation_metadata(citations: &[Value]) -> Option { + let citations: Vec<_> = citations + .iter() + .filter(|citation| citation["type"] == "web_search_result_location") + .cloned() + .collect(); + (!citations.is_empty()).then(|| { + provider_namespace(CANONICAL, json!({ "citations": citations })) + .expect("provider metadata must be an object") + }) +} + +pub(crate) struct CitationDocument { + pub title: String, + pub filename: Option, + pub media_type: String, +} + +fn citation_source( + citation: &Value, + documents: &[CitationDocument], + generate_id: &(dyn Fn() -> String + Send + Sync), +) -> Option { + match citation["type"].as_str()? { + "web_search_result_location" => Some(Source::Url { + id: generate_id(), + url: str_field(citation, "url")?, + title: str_field(citation, "title"), + provider_metadata: Some( + provider_namespace( + CANONICAL, + json!({ + "citedText": citation["cited_text"], + "encryptedIndex": citation["encrypted_index"], + }), + ) + .expect("provider metadata must be an object"), + ), + }), + kind @ ("page_location" | "char_location") => { + let index = usize::try_from(citation["document_index"].as_u64()?).ok()?; + let document = documents.get(index)?; + let metadata = if kind == "page_location" { + json!({ + "citedText": citation["cited_text"], + "startPageNumber": citation["start_page_number"], + "endPageNumber": citation["end_page_number"], + }) + } else { + json!({ + "citedText": citation["cited_text"], + "startCharIndex": citation["start_char_index"], + "endCharIndex": citation["end_char_index"], + }) + }; + Some(Source::Document { + id: generate_id(), + media_type: document.media_type.clone(), + title: str_field(citation, "document_title") + .unwrap_or_else(|| document.title.clone()), + filename: document.filename.clone(), + provider_metadata: Some( + provider_namespace(CANONICAL, metadata) + .expect("provider metadata must be an object"), + ), + }) + } + _ => None, + } +} + +fn web_fetch_document(payload: &Value) -> Option { + (payload["type"] == "web_fetch_result").then(|| CitationDocument { + title: str_field(&payload["content"], "title") + .or_else(|| str_field(payload, "url")) + .unwrap_or_default(), + filename: None, + media_type: str_field(&payload["content"]["source"], "media_type").unwrap_or_default(), + }) +} + +fn compaction_metadata(signature: Option<&str>) -> ProviderMetadata { + let mut metadata = json!({ "type": "compaction" }); + if let Some(signature) = signature { + metadata["signature"] = json!(signature); + } + provider_namespace(CANONICAL, metadata).expect("provider metadata must be an object") +} + /// Map Anthropic response content blocks into `GenerateContent` items. /// /// Text / tool_use / thinking / server_tool_use blocks are surfaced, and every @@ -486,11 +612,14 @@ pub(crate) fn stream_parts_for_result_block( /// `names` maps Anthropic's wire tool names back to the names the caller used; /// pass the mapping built from `CallOptions.tools`. pub(crate) fn parse_anthropic_content( - options_name: &str, + _options_name: &str, blocks: &[ContentBlock], names: &ToolNameMapping, + uses_json_response_tool: bool, + mut citation_documents: Vec, generate_id: &(dyn Fn() -> String + Send + Sync), ) -> Vec { + let options_name = CANONICAL; let mut content = Vec::new(); // Result blocks inherit information from their matching calls. Index the // complete response first so ordering does not affect non-stream parsing. @@ -515,9 +644,44 @@ pub(crate) fn parse_anthropic_content( for block in blocks { match block { - ContentBlock::Text { text } => { + ContentBlock::ContainerUpload { file_id } => { + content.push(GenerateContent::Custom { + kind: "anthropic.container_upload".to_string(), + provider_metadata: Some( + provider_namespace(CANONICAL, json!({ "fileId": file_id })) + .expect("provider metadata must be an object"), + ), + }); + } + ContentBlock::Text { .. } if uses_json_response_tool => {} + ContentBlock::Text { text, citations } => { + content.push(GenerateContent::Text { + text: text.clone(), + provider_metadata: citation_metadata(citations), + }); + content.extend( + citations + .iter() + .filter_map(|citation| { + citation_source(citation, &citation_documents, generate_id) + }) + .map(GenerateContent::Source), + ); + } + ContentBlock::Compaction { + content: Some(text), + signature, + } if !text.is_empty() => { content.push(GenerateContent::Text { text: text.clone(), + provider_metadata: Some(compaction_metadata(signature.as_deref())), + }); + } + ContentBlock::ToolUse { name, input, .. } + if uses_json_response_tool && name == "json" => + { + content.push(GenerateContent::Text { + text: input.to_string(), provider_metadata: None, }); } @@ -525,15 +689,26 @@ pub(crate) fn parse_anthropic_content( id, name, input, + toolset_name, caller, } => { content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id.clone(), - tool_name: names.to_custom_tool_name(name).to_string(), - input: input.to_string(), + tool_name: toolset_name.as_deref().map_or_else( + || name.clone(), + |name| names.to_custom_tool_name(name).to_string(), + ), + input: toolset_name.as_ref().map_or_else( + || input.to_string(), + |_| toolset_member_input(name, input).to_string(), + ), provider_executed: None, dynamic: None, - provider_metadata: tool_call_caller_metadata(caller.as_ref(), options_name), + provider_metadata: tool_call_metadata( + caller.as_ref(), + toolset_name.as_deref(), + options_name, + ), })); } ContentBlock::Thinking { @@ -550,7 +725,12 @@ pub(crate) fn parse_anthropic_content( } // Provider-executed (server-side) tool calls are surfaced as tool // calls so they round-trip on follow-up turns. - ContentBlock::ServerToolUse { id, name, input } => { + ContentBlock::ServerToolUse { + id, + name, + input, + caller, + } => { let provider_name = server_tool_provider_name(name); content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id.clone(), @@ -560,7 +740,7 @@ pub(crate) fn parse_anthropic_content( dynamic: (provider_name == "code_execution" && names.mark_code_execution_dynamic()) .then_some(true), - provider_metadata: None, + provider_metadata: tool_call_caller_metadata(caller.as_ref(), CANONICAL), })); } // MCP tool use — provider-executed + dynamic (upstream :1166-1182). @@ -603,6 +783,7 @@ pub(crate) fn parse_anthropic_content( ContentBlock::WebSearchToolResult { tool_use_id, content: payload, + caller, } => { // Success is an array of results; anything else is the error // object (upstream branches on `Array.isArray`). @@ -617,7 +798,10 @@ pub(crate) fn parse_anthropic_content( is_error: None, preliminary: None, dynamic: None, - provider_metadata: None, + provider_metadata: tool_call_caller_metadata( + caller.as_ref(), + CANONICAL, + ), })); // Each result also becomes a `Source` — that is how the // URLs and titles reach `result.sources`. @@ -650,14 +834,18 @@ pub(crate) fn parse_anthropic_content( is_error: Some(true), preliminary: None, dynamic: None, - provider_metadata: None, + provider_metadata: tool_call_caller_metadata(caller.as_ref(), CANONICAL), })), } } ContentBlock::WebFetchToolResult { tool_use_id, content: payload, + caller, } => { + if let Some(document) = web_fetch_document(payload) { + citation_documents.push(document); + } let (result, is_error) = map_web_fetch_result(payload); content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: tool_use_id.clone(), @@ -666,7 +854,7 @@ pub(crate) fn parse_anthropic_content( is_error, preliminary: None, dynamic: None, - provider_metadata: None, + provider_metadata: tool_call_caller_metadata(caller.as_ref(), CANONICAL), })); } ContentBlock::CodeExecutionToolResult { @@ -778,14 +966,17 @@ pub(crate) fn parse_anthropic_content( /// full version, so every host reports the same detailed token accounting. /// `config` supplies the providerOptions key the metadata is written under and /// the host's error shape. +#[allow(clippy::too_many_arguments)] pub(crate) async fn anthropic_generate_core( request: HttpRequest, body: serde_json::Value, warnings: Vec, config: &AnthropicModelConfig, - used_custom_provider_key: bool, tool_names: &ToolNameMapping, - generate_id: &(dyn Fn() -> String + Send + Sync), + used_custom_options_key: bool, + uses_json_response_tool: bool, + citation_documents: Vec, + generate_id: std::sync::Arc String + Send + Sync>, ) -> Result { let resp = aimux_provider_utils::post_json_to_api( request, @@ -796,18 +987,23 @@ pub(crate) async fn anthropic_generate_core( .await?; let response_body = resp.raw_value; - let response_headers = resp.response_headers; - let data: AnthropicResponse = resp.value; + let is_json_response_from_tool = uses_json_response_tool + && data + .content + .iter() + .any(|part| matches!(part, ContentBlock::ToolUse { name, .. } if name == "json")); let content = parse_anthropic_content( &config.provider_options_name, &data.content, tool_names, - generate_id, + uses_json_response_tool, + citation_documents, + generate_id.as_ref(), ); - let finish_reason = data + let mut finish_reason = data .stop_reason .as_deref() .map(parse_stop_reason) @@ -816,18 +1012,28 @@ pub(crate) async fn anthropic_generate_core( raw: None, }); + if is_json_response_from_tool && finish_reason.unified == FinishReasonUnified::ToolCalls { + finish_reason.unified = FinishReasonUnified::Stop; + } + // RFC-0015 P0-2: fill cache fields + raw; total = input + cache_read + // cache_creation (Anthropic's input_tokens excludes cache). Output side // breakdown (text/reasoning) comes from output_tokens_details. let usage = super::usage::usage_from_anthropic(&data.usage); - let provider_metadata = super::usage::result_provider_metadata( + let mut provider_metadata = super::usage::result_provider_metadata( &config.provider_options_name, - used_custom_provider_key, &serde_json::to_value(&data.usage).unwrap_or(Value::Null), data.stop_sequence.as_deref(), data.container.as_ref(), data.context_management.as_ref(), + used_custom_options_key, + ); + super::usage::extend_result_metadata( + &mut provider_metadata, + data.stop_details.as_ref(), + data.input_transformations.as_ref(), + data.safeguard_results.as_ref(), ); Ok(GenerateResult { @@ -840,7 +1046,7 @@ pub(crate) async fn anthropic_generate_core( id: Some(data.id), timestamp: None, model_id: Some(data.model), - headers: Some(response_headers), + headers: Some(resp.response_headers), body: response_body, }), request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), @@ -849,13 +1055,12 @@ pub(crate) async fn anthropic_generate_core( /// Per-content-block state during streaming. /// -/// Text blocks track whether a `TextStart` has been emitted; tool_use blocks -/// accumulate the `input_json_delta` partial-json fragments so the final -/// `ToolCall` can carry the parsed JSON object; thinking blocks track whether a -/// `ReasoningStart` has been emitted. +/// Text blocks collect web citations for `TextEnd`; tool-use blocks accumulate +/// partial JSON for the final tool call. Every text/reasoning block starts as +/// soon as its provider content-block-start event arrives. enum BlockState { Text { - started: bool, + citations: Vec, }, ToolUse { id: String, @@ -867,12 +1072,9 @@ enum BlockState { provider_tool_input_type: Option, provider_metadata: Option, first_delta: bool, + toolset_member_name: Option, }, - Thinking { - started: bool, - /// Accumulated `signature_delta` pieces for the thinking block. - signature: Option, - }, + Thinking, } /// Shared Anthropic streaming core. @@ -881,13 +1083,16 @@ enum BlockState { /// resolved by the model), then runs the Anthropic SSE event loop to produce a /// `StreamResult`. `config` supplies the providerOptions key the metadata is /// written under and the host's error shape. +#[allow(clippy::too_many_arguments)] pub(crate) async fn anthropic_stream_core( request: HttpRequest, body: serde_json::Value, warnings: Vec, config: &AnthropicModelConfig, - used_custom_provider_key: bool, tool_names: ToolNameMapping, + used_custom_options_key: bool, + uses_json_response_tool: bool, + mut citation_documents: Vec, generate_id: std::sync::Arc String + Send + Sync>, ) -> Result { let endpoint = request.url.clone(); @@ -924,17 +1129,24 @@ pub(crate) async fn anthropic_stream_core( let mut sse = futures::stream::iter(first_event.into_iter()).chain(sse_stream); let mut blocks: HashMap = HashMap::new(); + let mut is_json_response_from_tool = false; let mut final_usage = Usage::default(); let mut final_finish_reason: Option = None; // Result-level providerMetadata: the raw usage (`message_start`'s, // updated by every `message_delta`'s), the stop sequence, the // container and the context-management edits. let mut raw_usage = Value::Object(serde_json::Map::new()); + let mut effective_usage = raw_usage.clone(); let mut stop_sequence: Option = None; let mut container: Option = None; let mut context_management: Option = None; + let mut stop_details: Option = None; + let mut input_transformations: Option = None; + let mut safeguard_results: Option = None; let mut response_meta_emitted = false; - let mut stream_errored = false; + let mut active_message_id: Option = None; + let mut message_stopped = false; + // id → (tool name, server name), so `mcp_tool_result` can inherit them // from the `mcp_tool_use` it answers. let mut mcp_tool_calls: HashMap = HashMap::new(); @@ -947,12 +1159,31 @@ pub(crate) async fn anthropic_stream_core( Ok(stream_event) => { match stream_event { StreamEvent::MessageStart { message } => { + if let Some(active_id) = &active_message_id { + if active_id == &message.id { + continue; + } + yield Ok(StreamPart::Error { + error: AiMuxError::InvalidResponseData(format!( + "Received message_start for message {:?} while message {:?} is still open.", + message.id, active_id, + )), + }); + return; + } + active_message_id = Some(message.id.clone()); + input_transformations = message.input_transformations.or(input_transformations); + container = message.container.or(container); + if let Some(reason) = message.stop_reason.as_deref() { + final_finish_reason = Some(parse_stop_reason(reason)); + } if let Some(usage) = &message.usage { // RFC-0015 P0-2: full input side incl. cache // fields + raw (Anthropic reports cache only // in message_start). final_usage = super::usage::usage_from_anthropic(usage); raw_usage = serde_json::to_value(usage).unwrap_or(raw_usage); + effective_usage = raw_usage.clone(); } if !response_meta_emitted { yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { @@ -965,27 +1196,42 @@ pub(crate) async fn anthropic_stream_core( } StreamEvent::ContentBlockStart { index, content_block } => { match content_block { + ContentBlock::Text { .. } if uses_json_response_tool => {} ContentBlock::Text { .. } => { - blocks.insert(index, BlockState::Text { started: false }); + blocks.insert(index, BlockState::Text { citations: Vec::new() }); + yield Ok(StreamPart::TextStart { id: index.to_string(), provider_metadata: None }); + } + ContentBlock::Compaction { content, signature } => { + blocks.insert(index, BlockState::Text { citations: Vec::new() }); + yield Ok(StreamPart::TextStart { id: index.to_string(), provider_metadata: Some(compaction_metadata(signature.as_deref())) }); + if signature.is_some() + && let Some(text) = content.filter(|text| !text.is_empty()) { + yield Ok(StreamPart::TextDelta { id: index.to_string(), delta: text, provider_metadata: None }); + } } ContentBlock::Thinking { .. } => { + yield Ok(StreamPart::ReasoningStart { id: index.to_string(), provider_metadata: None }); blocks.insert( index, - BlockState::Thinking { - started: false, - signature: None, - }, + BlockState::Thinking, ); } + ContentBlock::ToolUse { name, .. } if uses_json_response_tool && name == "json" => { + is_json_response_from_tool = true; + blocks.insert(index, BlockState::Text { citations: Vec::new() }); + yield Ok(StreamPart::TextStart { id: index.to_string(), provider_metadata: None }); + } ContentBlock::ToolUse { id, name, input, + toolset_name, caller, } => { - let custom_name = tool_names - .to_custom_tool_name(&name) - .to_string(); + let custom_name = toolset_name.as_deref().map_or_else( + || name.clone(), + |name| tool_names.to_custom_tool_name(name).to_string(), + ); let initial_input = initial_tool_input(&input); yield Ok(StreamPart::ToolInputStart { id: id.clone(), @@ -1004,13 +1250,14 @@ pub(crate) async fn anthropic_stream_core( dynamic: None, provider_tool_name: None, provider_tool_input_type: None, - provider_metadata: tool_call_caller_metadata(caller.as_ref(), options_name.as_str()), + provider_metadata: tool_call_metadata(caller.as_ref(), toolset_name.as_deref(), CANONICAL), + toolset_member_name: toolset_name.map(|_| name), }); } // Server-side tool use follows the same input // lifecycle as client tools because some code // execution inputs arrive entirely via deltas. - ContentBlock::ServerToolUse { id, name, input } => { + ContentBlock::ServerToolUse { id, name, input, caller } => { if is_tool_search_provider_name(&name) { server_tool_calls.insert(id.clone(), name.clone()); } @@ -1047,7 +1294,8 @@ pub(crate) async fn anthropic_stream_core( } _ => None, }, - provider_metadata: None, + provider_metadata: tool_call_caller_metadata(caller.as_ref(), CANONICAL), + toolset_member_name: None, }); } // MCP tool use — provider-executed + dynamic. @@ -1060,7 +1308,7 @@ pub(crate) async fn anthropic_stream_core( input: input.to_string(), provider_executed: Some(true), dynamic: Some(true), - provider_metadata: Some(provider_namespace(options_name.as_str(), json!({ + provider_metadata: Some(provider_namespace(CANONICAL, json!({ "type": "mcp-tool-use", "serverName": server_name, })).expect("provider metadata must be an object")), @@ -1071,14 +1319,11 @@ pub(crate) async fn anthropic_stream_core( let id = index.to_string(); yield Ok(StreamPart::ReasoningStart { id: id.clone(), - provider_metadata: Some(provider_namespace(options_name.as_str(), json!({ "redactedData": data })).expect("provider metadata must be an object")), + provider_metadata: Some(provider_namespace(CANONICAL, json!({ "redactedData": data })).expect("provider metadata must be an object")), }); blocks.insert( index, - BlockState::Thinking { - started: true, - signature: None, - }, + BlockState::Thinking, ); } // Server-tool result blocks arrive whole on @@ -1087,8 +1332,12 @@ pub(crate) async fn anthropic_stream_core( // mirrors its `doGenerate` switch here too // (anthropic-language-model.ts:1901-2178). other => { + if let ContentBlock::WebFetchToolResult { content, .. } = &other + && let Some(document) = web_fetch_document(content) { + citation_documents.push(document); + } for part in stream_parts_for_result_block( - options_name.as_str(), + CANONICAL, &other, &tool_names, &mcp_tool_calls, @@ -1101,39 +1350,35 @@ pub(crate) async fn anthropic_stream_core( } } StreamEvent::ContentBlockDelta { index, delta } => { - if let Some(text) = delta.text { - // Start the text segment on the first delta. The - // text id is the stringified content-block - // index, matching the TS SDK. - let start_id: Option = match blocks.get_mut(&index) { - Some(BlockState::Text { started: false }) => { - if let Some(BlockState::Text { started }) = - blocks.get_mut(&index) - { - *started = true; - } - Some(index.to_string()) - } - _ => None, - }; - if let Some(id) = start_id { - yield Ok(StreamPart::TextStart { - id, - provider_metadata: None, - }); - } + if !uses_json_response_tool && let Some(text) = delta.text { yield Ok(StreamPart::TextDelta { id: index.to_string(), delta: text, provider_metadata: None, }); } + if let Some(text) = delta.content { + yield Ok(StreamPart::TextDelta { id: index.to_string(), delta: text, provider_metadata: None }); + } + if let Some(citation) = delta.citation { + if let Some(BlockState::Text { citations }) = blocks.get_mut(&index) { + citations.push(citation.clone()); + } + if let Some(source) = citation_source(&citation, &citation_documents, generate_id.as_ref()) { + yield Ok(StreamPart::Source(source)); + } + } if let Some(partial) = delta.partial_json { // Accumulate the partial JSON fragment and emit // a ToolInputDelta. Empty fragments (the // leading `input_json_delta` with // `partial_json: ""`) are skipped, matching the // TS SDK. + if is_json_response_from_tool && matches!(blocks.get(&index), Some(BlockState::Text { .. })) && !partial.is_empty() { + yield Ok(StreamPart::TextDelta { id: index.to_string(), delta: partial, provider_metadata: None }); + continue; + } + if is_json_response_from_tool { continue; } let delta_event: Option<(String, String)> = match blocks.get_mut(&index) { Some(BlockState::ToolUse { @@ -1141,6 +1386,7 @@ pub(crate) async fn anthropic_stream_core( accumulated_json, provider_tool_input_type, first_delta, + toolset_member_name, .. }) if !partial.is_empty() => { let emitted_delta = if *first_delta { @@ -1157,7 +1403,7 @@ pub(crate) async fn anthropic_stream_core( }; accumulated_json.push_str(&emitted_delta); *first_delta = false; - Some((id.clone(), emitted_delta)) + toolset_member_name.is_none().then(|| (id.clone(), emitted_delta)) } _ => None, }; @@ -1170,92 +1416,64 @@ pub(crate) async fn anthropic_stream_core( } } if let Some(thinking) = delta.thinking { - // Start the reasoning segment on the first - // thinking delta. The id is the stringified - // content-block index, matching the TS SDK. - let start_id: Option = match blocks.get_mut(&index) { - Some(BlockState::Thinking { started: false, .. }) => { - if let Some(BlockState::Thinking { started, .. }) = - blocks.get_mut(&index) - { - *started = true; - } - Some(index.to_string()) - } - _ => None, - }; - if let Some(id) = start_id { - yield Ok(StreamPart::ReasoningStart { - id, - provider_metadata: None, - }); - } yield Ok(StreamPart::ReasoningDelta { id: index.to_string(), delta: thinking, provider_metadata: None, }); } - if let Some(sig) = delta.signature { - // `signature_delta`: incremental pieces of the - // thinking-block signature. Accumulated on the - // block state and attached to ReasoningEnd, so - // the aggregated Reasoning part can echo it - // back on extended-thinking multi-turn — same - // shape as the non-streaming path. - // - // Edge: a thinking block carrying only a - // signature (no text deltas) stays - // `started: false` and emits nothing — the - // protocol always sends text first, and the - // core aggregator gates on non-empty text - // anyway, so there is nothing to attach to. - if let Some(BlockState::Thinking { signature, .. }) = - blocks.get_mut(&index) - { - signature.get_or_insert_with(String::new).push_str(&sig); + if let Some(sig) = delta.signature + && let Some(BlockState::Thinking) = blocks.get(&index) { + yield Ok(StreamPart::ReasoningDelta { + id: index.to_string(), delta: String::new(), + provider_metadata: Some(provider_namespace(CANONICAL, json!({ "signature": sig })).expect("provider metadata must be an object")), + }); } - } } StreamEvent::ContentBlockStop { index } => { // Removing the block releases the borrow before any // yield. if let Some(state) = blocks.remove(&index) { match state { - BlockState::Text { started: true } => { + BlockState::Text { citations } => { yield Ok(StreamPart::TextEnd { id: index.to_string(), - provider_metadata: None, + provider_metadata: citation_metadata(&citations), }); } - BlockState::Text { started: false } => { - // Text block with no deltas — nothing to emit. - } - BlockState::Thinking { - started: true, - signature, - } => { + BlockState::Thinking => { yield Ok(StreamPart::ReasoningEnd { id: index.to_string(), - provider_metadata: signature.map(|s| { - provider_namespace(options_name.as_str(), json!({ "signature": s })).expect("provider metadata must be an object") - }), + provider_metadata: None, }); } - BlockState::Thinking { started: false, .. } => { - // Thinking block with no deltas — nothing to emit. - } BlockState::ToolUse { id, name, - accumulated_json, + mut accumulated_json, provider_executed, dynamic, provider_tool_name, provider_tool_input_type, provider_metadata, + toolset_member_name, .. } => { + if let Some(member_name) = toolset_member_name { + let parsed = if accumulated_json.is_empty() { + Ok(json!({})) + } else { + serde_json::from_str::(&accumulated_json) + }; + if let Ok(input) = parsed { + accumulated_json = toolset_member_input(&member_name, &input).to_string(); + } + yield Ok(StreamPart::ToolInputDelta { + id: id.clone(), + delta: accumulated_json.clone(), + provider_metadata: None, + }); + } yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: None, @@ -1278,41 +1496,36 @@ pub(crate) async fn anthropic_stream_core( } } } - StreamEvent::MessageDelta { delta, usage, context_management: edits } => { + StreamEvent::MessageDelta { delta, usage, context_management: edits, input_transformations: transformations } => { if let Some(reason) = delta.stop_reason { final_finish_reason = Some(parse_stop_reason(&reason)); } - stop_sequence = delta.stop_sequence.or(stop_sequence); - container = delta.container.or(container); + stop_sequence = delta.stop_sequence; + container = delta.container; context_management = edits.or(context_management); + stop_details = delta.stop_details; + input_transformations = transformations.or(input_transformations); + safeguard_results = delta.safeguard_results.or(safeguard_results); if let Some(u) = usage { if let (Value::Object(raw), Ok(Value::Object(update))) = (&mut raw_usage, serde_json::to_value(&u)) { + if let Value::Object(effective) = &mut effective_usage { + effective.extend(update.iter().filter(|(_, value)| !value.is_null()).map(|(key, value)| (key.clone(), value.clone()))); + } raw.extend(update); } - final_usage.raw = raw_usage.as_object().cloned(); - let reasoning_tokens = u - .output_tokens_details - .as_ref() - .and_then(|d| d.thinking_tokens); - let output_total = u.output_tokens; - let text_tokens = reasoning_tokens - .zip(output_total) - .map(|(r, t)| t.saturating_sub(r)); - final_usage.output_tokens = aimux_core::types::OutputTokenUsage { - total: output_total, - text: text_tokens, - reasoning: reasoning_tokens, - }; + if let Ok(usage) = serde_json::from_value(effective_usage.clone()) { + final_usage = super::usage::usage_from_anthropic(&usage); + final_usage.raw = raw_usage.as_object().cloned(); + } } } - StreamEvent::MessageStop => break, + StreamEvent::MessageStop => { + message_stopped = true; + break; + }, StreamEvent::Error { error } => { - // Surface Anthropic in-stream errors (e.g. - // overloaded_error) as a `StreamPart::Error` and - // stop the stream, mirroring the TS "forward error - // chunks" / "forward overloaded error" behaviour. yield Ok(StreamPart::Error { error: anthropic_stream_error( &error, @@ -1321,19 +1534,14 @@ pub(crate) async fn anthropic_stream_core( stream_response_headers.clone(), ), }); - // Finish is the final-chunk contract: it must - // still terminate the stream after an in-stream - // error (openai and google do the same), so - // downstream aggregators see a well-formed end. - stream_errored = true; - break; + } _ => {} } } Err(error) => { let recoverable = error.is_recoverable_stream_error(); - yield Err(error); + yield Ok(StreamPart::Error { error }); if !recoverable { return; } @@ -1341,33 +1549,32 @@ pub(crate) async fn anthropic_stream_core( } } - // Final part: Finish. + // The provider completes a message explicitly; EOF after an error + // does not synthesize a finish event. + if !message_stopped { + return; + } + if is_json_response_from_tool && let Some(reason) = &mut final_finish_reason + && reason.unified == FinishReasonUnified::ToolCalls { + reason.unified = FinishReasonUnified::Stop; + } yield Ok(StreamPart::Finish { - finish_reason: if stream_errored { - FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - } - } else { - final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }) - }, - usage: if stream_errored { - Usage::default() - } else { - final_usage - }, - provider_metadata: (!stream_errored).then(|| { - super::usage::result_provider_metadata( + finish_reason: final_finish_reason.unwrap_or(FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }), + usage: final_usage, + provider_metadata: Some({ + let mut metadata = super::usage::result_provider_metadata( options_name.as_str(), - used_custom_provider_key, - &raw_usage, + &raw_usage, stop_sequence.as_deref(), container.as_ref(), context_management.as_ref(), - ) + used_custom_options_key, + ); + super::usage::extend_result_metadata(&mut metadata, stop_details.as_ref(), input_transformations.as_ref(), safeguard_results.as_ref()); + metadata }), }); }; diff --git a/aimux-providers/src/anthropic/tool_name_mapping.rs b/aimux-providers/src/anthropic/tool_name_mapping.rs index 9bf67eaa0..ee07e2d71 100644 --- a/aimux-providers/src/anthropic/tool_name_mapping.rs +++ b/aimux-providers/src/anthropic/tool_name_mapping.rs @@ -33,15 +33,20 @@ fn provider_tool_name(id: &str) -> Option<&'static str> { | "anthropic.code_execution_20260120" => "code_execution", "anthropic.computer_20241022" | "anthropic.computer_20250124" - | "anthropic.computer_20251124" => "computer", + | "anthropic.computer_20251124" + | "anthropic.computer_toolset_20260801" => "computer", "anthropic.text_editor_20241022" | "anthropic.text_editor_20250124" => "str_replace_editor", "anthropic.text_editor_20250429" | "anthropic.text_editor_20250728" => { "str_replace_based_edit_tool" } "anthropic.bash_20241022" | "anthropic.bash_20250124" => "bash", "anthropic.memory_20250818" => "memory", - "anthropic.web_search_20250305" | "anthropic.web_search_20260209" => "web_search", - "anthropic.web_fetch_20250910" | "anthropic.web_fetch_20260209" => "web_fetch", + "anthropic.web_search_20250305" + | "anthropic.web_search_20260209" + | "anthropic.web_search_20260318" => "web_search", + "anthropic.web_fetch_20250910" + | "anthropic.web_fetch_20260209" + | "anthropic.web_fetch_20260318" => "web_fetch", "anthropic.tool_search_regex_20251119" => "tool_search_tool_regex", "anthropic.tool_search_bm25_20251119" => "tool_search_tool_bm25", "anthropic.advisor_20260301" => "advisor", @@ -58,6 +63,7 @@ pub struct ToolNameMapping { custom_to_provider: HashMap, provider_to_custom: HashMap, mark_code_execution_dynamic: bool, + toolset_names: HashMap, } impl ToolNameMapping { @@ -69,9 +75,17 @@ impl ToolNameMapping { let mut has_code_execution = false; for tool in tools.unwrap_or(&[]) { if let Tool::Provider(pt) = tool { + if pt.id == "anthropic.computer_toolset_20260801" { + mapping + .toolset_names + .insert(pt.name.clone(), "computer".to_string()); + } has_web_tool_20260209 |= matches!( pt.id.as_str(), - "anthropic.web_search_20260209" | "anthropic.web_fetch_20260209" + "anthropic.web_search_20260209" + | "anthropic.web_fetch_20260209" + | "anthropic.web_search_20260318" + | "anthropic.web_fetch_20260318" ); has_code_execution |= pt.id.starts_with("anthropic.code_execution_"); @@ -89,6 +103,10 @@ impl ToolNameMapping { mapping } + pub(crate) fn toolset_name(&self, custom_name: &str) -> Option<&str> { + self.toolset_names.get(custom_name).map(String::as_str) + } + /// Caller's name → Anthropic's wire name. Unmapped names pass through. pub fn to_provider_tool_name<'a>(&'a self, custom_name: &'a str) -> &'a str { self.custom_to_provider @@ -111,25 +129,3 @@ impl ToolNameMapping { self.mark_code_execution_dynamic } } - -#[cfg(test)] -mod tests { - use super::*; - use aimux_core::tool::FunctionTool; - use serde_json::json; - - #[test] - fn unmapped_names_pass_through() { - let mapping = ToolNameMapping::new(None); - assert_eq!(mapping.to_custom_tool_name("web_search"), "web_search"); - assert_eq!(mapping.to_provider_tool_name("get_weather"), "get_weather"); - } - - #[test] - fn function_tools_are_ignored() { - let tools = vec![Tool::Function(FunctionTool::new("web_search", json!({})))]; - let mapping = ToolNameMapping::new(Some(&tools)); - // A function tool named `web_search` must not hijack the provider name. - assert_eq!(mapping.to_custom_tool_name("web_search"), "web_search"); - } -} diff --git a/aimux-providers/src/anthropic/types.rs b/aimux-providers/src/anthropic/types.rs index 066f09519..cd40ccee8 100644 --- a/aimux-providers/src/anthropic/types.rs +++ b/aimux-providers/src/anthropic/types.rs @@ -29,6 +29,12 @@ pub struct AnthropicResponse { /// `providerMetadata.anthropic.contextManagement`. #[serde(default)] pub context_management: Option, + #[serde(default)] + pub stop_details: Option, + #[serde(default)] + pub input_transformations: Option, + #[serde(default)] + pub safeguard_results: Option, } /// Origin of an Anthropic programmatic tool call. @@ -47,13 +53,28 @@ pub enum ToolCallCaller { #[serde(tag = "type")] pub enum ContentBlock { #[serde(rename = "text")] - Text { text: String }, + Text { + text: String, + #[serde(default)] + citations: Vec, + }, + #[serde(rename = "container_upload")] + ContainerUpload { file_id: String }, + #[serde(rename = "compaction")] + Compaction { + #[serde(default)] + content: Option, + #[serde(default)] + signature: Option, + }, #[serde(rename = "tool_use")] ToolUse { id: String, name: String, input: Value, #[serde(default)] + toolset_name: Option, + #[serde(default)] caller: Option, }, /// Anthropic extended-thinking block. Carries the reasoning text and an @@ -81,6 +102,8 @@ pub enum ContentBlock { name: String, #[serde(default)] input: Value, + #[serde(default)] + caller: Option, }, /// Result of a server-side web search. `content` is either an array of /// `web_search_result` objects or a `web_search_tool_result_error` object. @@ -90,6 +113,8 @@ pub enum ContentBlock { tool_use_id: String, #[serde(default)] content: Value, + #[serde(default)] + caller: Option, }, /// Result of a server-side web fetch. #[serde(rename = "web_fetch_tool_result")] @@ -98,6 +123,8 @@ pub enum ContentBlock { tool_use_id: String, #[serde(default)] content: Value, + #[serde(default)] + caller: Option, }, /// Result of server-side code execution (20250522). #[serde(rename = "code_execution_tool_result")] @@ -174,35 +201,11 @@ pub enum ContentBlock { Other, } -/// `output_tokens_details` nested inside an Anthropic usage object, carrying -/// the count of thinking/reasoning tokens. -#[derive(Debug, Default, Serialize, Deserialize)] -pub struct OutputTokenDetails { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub thinking_tokens: Option, -} - -/// The `usage` object of a response, `message_start` or `message_delta`. -/// -/// Serializing it reproduces what the API sent: absent fields stay absent and -/// every field this type does not name (`iterations`, `service_tier`, ...) is -/// kept in `extra`, so the raw usage and `providerMetadata.anthropic.usage` -/// are the provider's own object. +/// Preserve the provider usage object verbatim, including explicit nulls and +/// unknown fields. Token accounting uses the separate typed view in `usage`. #[derive(Debug, Default, Serialize, Deserialize)] -pub struct AnthropicUsage { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub input_tokens: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub output_tokens: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub output_tokens_details: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cache_creation_input_tokens: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cache_read_input_tokens: Option, - #[serde(flatten)] - pub extra: serde_json::Map, -} +#[serde(transparent)] +pub struct AnthropicUsage(pub Value); // ── Streaming events ── @@ -227,6 +230,8 @@ pub enum StreamEvent { usage: Option, #[serde(default)] context_management: Option, + #[serde(default)] + input_transformations: Option, }, #[serde(rename = "message_stop")] MessageStop, @@ -254,6 +259,12 @@ pub struct MessageStartData { pub model: String, #[serde(default)] pub usage: Option, + #[serde(default)] + pub input_transformations: Option, + #[serde(default)] + pub container: Option, + #[serde(default)] + pub stop_reason: Option, } #[derive(Debug, Deserialize)] @@ -268,6 +279,10 @@ pub struct DeltaBlock { /// `signature_delta` payload — incremental signature text. #[serde(default)] pub signature: Option, + #[serde(default)] + pub content: Option, + #[serde(default)] + pub citation: Option, } #[derive(Debug, Deserialize)] @@ -278,4 +293,8 @@ pub struct MessageDeltaBody { pub stop_sequence: Option, #[serde(default)] pub container: Option, + #[serde(default)] + pub stop_details: Option, + #[serde(default)] + pub safeguard_results: Option, } diff --git a/aimux-providers/src/anthropic/usage.rs b/aimux-providers/src/anthropic/usage.rs index 9568f71cc..072673a20 100644 --- a/aimux-providers/src/anthropic/usage.rs +++ b/aimux-providers/src/anthropic/usage.rs @@ -147,11 +147,11 @@ pub fn convert_anthropic_usage(usage: &Value, raw_usage: Option<&Value>) -> Anth /// are re-keyed in camelCase; every absent piece is `null`. pub(crate) fn result_provider_metadata( options_name: &str, - used_custom_provider_key: bool, usage: &Value, stop_sequence: Option<&str>, container: Option<&Value>, context_management: Option<&Value>, + used_custom_options_key: bool, ) -> ProviderMetadata { let field = |value: &Value, key: &str| value.get(key).cloned().unwrap_or(Value::Null); let iterations = @@ -168,17 +168,18 @@ pub(crate) fn result_provider_metadata( "inputTokens": field(i, "input_tokens"), "outputTokens": field(i, "output_tokens"), }); - if let Some(model) = i.get("model").filter(|model| !model.is_null()) { + if let Some(model) = i.get("model").filter(|value| !value.is_null()) { iteration["model"] = model.clone(); } - for (source, target) in [ + for (wire, key) in [ ("cache_creation_input_tokens", "cacheCreationInputTokens"), ("cache_read_input_tokens", "cacheReadInputTokens"), ] { - if let Some(tokens) = i.get(source).filter(|tokens| { - tokens.as_u64().is_some_and(|value| value != 0) - }) { - iteration[target] = tokens.clone(); + if let Some(value) = i + .get(wire) + .filter(|value| value.as_u64().is_some_and(|value| value != 0)) + { + iteration[key] = value.clone(); } } iteration @@ -235,12 +236,41 @@ pub(crate) fn result_provider_metadata( }); let mut result = provider_namespace(CANONICAL_KEY, metadata).expect("provider metadata must be an object"); - if used_custom_provider_key && options_name != CANONICAL_KEY { + if used_custom_options_key && options_name != CANONICAL_KEY { result.insert(options_name.to_string(), result[CANONICAL_KEY].clone()); } result } +pub(crate) fn extend_result_metadata( + metadata: &mut ProviderMetadata, + stop_details: Option<&Value>, + input_transformations: Option<&Value>, + safeguard_results: Option<&Value>, +) { + for namespace in metadata.values_mut() { + if let Some(details) = stop_details.filter(|details| !details.is_null()) { + let mut mapped = json!({ "type": details["type"] }); + for (wire, key) in [ + ("category", "category"), + ("explanation", "explanation"), + ("recommended_model", "recommendedModel"), + ] { + if let Some(value) = details.get(wire).filter(|value| !value.is_null()) { + mapped[key] = value.clone(); + } + } + namespace.insert("stopDetails".to_string(), mapped); + } + if let Some(value) = input_transformations.filter(|value| !value.is_null()) { + namespace.insert("inputTransformations".to_string(), value.clone()); + } + if let Some(value) = safeguard_results.filter(|value| !value.is_null()) { + namespace.insert("safeguardResults".to_string(), value.clone()); + } + } +} + /// `cleared_input_tokens` -> `clearedInputTokens`. fn snake_to_camel(key: &str) -> String { let mut out = String::with_capacity(key.len()); diff --git a/aimux-providers/src/azure/mod.rs b/aimux-providers/src/azure/mod.rs index 5ab024c63..04fbe7eb7 100644 --- a/aimux-providers/src/azure/mod.rs +++ b/aimux-providers/src/azure/mod.rs @@ -1,7 +1,7 @@ //! # Azure OpenAI provider //! //! [`create_azure`] is the Rust form of the AI SDK's `createAzure`: it takes -//! [`AzureOpenAIProviderSettings`], rejects an `api_key` given together with a +//! [`AzureOpenAIProviderSettings`], rejects a non-empty `api_key` given with a //! `token_provider`, validates the base URL and returns an //! [`AzureOpenAIProvider`]. Nothing is read from the environment there: the //! API key (`AZURE_API_KEY`), the Entra ID token and the resource name @@ -19,12 +19,12 @@ //! providerOptions namespace and `assistant-` file-id prefix. //! //! Not ported from the AI SDK package: Azure-hosted DeepSeek (`deepseek`), -//! the legacy completion model, MAI-Transcribe / MAI-Voice (`speechBaseURL`, -//! `maiBaseURL`, `webSocket`), and Foundry-project message item types. +//! the legacy completion model and MAI-Voice (`maiBaseURL`, `webSocket`). pub(crate) mod options; +mod transcription; -pub use crate::shared::TransformRequestBody; +pub use transcription::AzureTranscriptionModel; use std::sync::{Arc, OnceLock}; @@ -39,8 +39,9 @@ use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_core::speech_model::SpeechModel; use aimux_core::transcription_model::TranscriptionModel; use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, Resolvable, combine_headers, load_setting, - validate_base_url, without_trailing_slash, + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, HeadersFn, + HttpRequest, Resolvable, default_fetch, load_setting, validate_base_url, + without_trailing_slash, }; use crate::openai::config::OpenAIModelConfig; @@ -86,11 +87,14 @@ pub struct AzureOpenAIProviderSettings { /// used as is; a non-Azure gateway gets `{base_url}{path}`. A trailing /// slash is removed. pub base_url: Option, + /// Azure Speech endpoint prefix. Independent of `base_url` and `api_version`. + /// Defaults to the resource's Cognitive Services endpoint. + pub speech_base_url: Option, /// The API key, sent as `api-key`. `None` (with no `token_provider`) loads /// `AZURE_API_KEY` when a request is made and fails that request with /// `AiMuxError::LoadApiKey` if it is unset. An explicit value is used as - /// given, `""` included. Giving it together with `token_provider` is an - /// error. + /// given, `""` included. Giving a non-empty key together with + /// `token_provider` is an error. pub api_key: Option>, /// Microsoft Entra ID access token, sent as `Authorization: Bearer` (and /// no `api-key`). Use an [`Resolvable::AsyncFn`] to get a fresh token on @@ -108,9 +112,6 @@ pub struct AzureOpenAIProviderSettings { /// `{base}/deployments/{deployment}{path}?api-version=` instead of /// `{base}/v1{path}`. pub use_deployment_based_urls: bool, - /// Rewrites every JSON request body once, after it is serialized and - /// before it is sent. - pub transform_request_body: Option, } impl std::fmt::Debug for AzureOpenAIProviderSettings { @@ -119,6 +120,7 @@ impl std::fmt::Debug for AzureOpenAIProviderSettings { f.debug_struct("AzureOpenAIProviderSettings") .field("resource_name", &self.resource_name) .field("base_url", &self.base_url) + .field("speech_base_url", &self.speech_base_url) .field("api_key", &self.api_key) .field("token_provider", &self.token_provider) .field( @@ -128,10 +130,6 @@ impl std::fmt::Debug for AzureOpenAIProviderSettings { .field("fetch", &self.fetch.is_some()) .field("api_version", &self.api_version) .field("use_deployment_based_urls", &self.use_deployment_based_urls) - .field( - "transform_request_body", - &self.transform_request_body.is_some(), - ) .finish() } } @@ -171,14 +169,19 @@ fn base_url_info(base_url: Option<&str>) -> Result { /// /// # Errors /// -/// Returns `AiMuxError::InvalidArgument` when both `api_key` and -/// `token_provider` are given or `base_url` is not an `http(s)` URL with a +/// Returns `AiMuxError::InvalidArgument` when a non-empty `api_key` and +/// `token_provider` are given or `base_url` / `speech_base_url` is not an `http(s)` URL with a /// host. Those are the only ways this fails: credentials and the resource /// name are resolved per request, not here. pub fn create_azure( settings: AzureOpenAIProviderSettings, ) -> Result { - if settings.api_key.is_some() && settings.token_provider.is_some() { + let has_api_key = match &settings.api_key { + None => false, + Some(Resolvable::Value(value)) => !value.is_empty(), + Some(_) => true, + }; + if has_api_key && settings.token_provider.is_some() { return Err(AiMuxError::InvalidArgument( "Both apiKey and tokenProvider were provided. Please use only one authentication \ method." @@ -190,10 +193,33 @@ pub fn create_azure( .as_deref() .map(validate_base_url) .transpose()?; + let speech_base_url = settings + .speech_base_url + .as_deref() + .map(validate_base_url) + .transpose()?; let info = base_url_info(base_url.as_deref())?; - let token_provider = settings.token_provider; - let headers = match &token_provider { - Some(_) => Resolvable::Value(combine_headers(&[&settings.headers.unwrap_or_default()])), + let fetch = settings + .token_provider + .clone() + .map(|token| { + Arc::new(BearerTokenFetch { + token, + fetch: settings.fetch.clone(), + }) as FetchFunction + }) + .or(settings.fetch); + let speech_headers = match &settings.token_provider { + Some(_) => Resolvable::Value(settings.headers.clone().unwrap_or_default()), + None => credential_headers( + Credential::explicit_or_env(settings.api_key.clone(), API_KEY_ENV_VAR, "Azure Speech"), + AuthScheme::Header("Ocp-Apim-Subscription-Key"), + Vec::new(), + settings.headers.clone(), + ), + }; + let headers = match settings.token_provider { + Some(_) => Resolvable::Value(settings.headers.unwrap_or_default()), None => credential_headers( Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Azure OpenAI"), AuthScheme::Header("api-key"), @@ -204,20 +230,56 @@ pub fn create_azure( Ok(AzureOpenAIProvider { resource_name: settings.resource_name, base_url, + speech_base_url, + speech_headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( + speech_headers, + options::NAMESPACE, + "4.0.84", + ), info, api_version: settings.api_version, use_deployment_based_urls: settings.use_deployment_based_urls, + fetch, headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( headers, options::NAMESPACE, "4.0.84", ), - fetch: settings.fetch, - token_provider, - transform_request_body: settings.transform_request_body, }) } +/// Resolve authentication after per-call headers have been merged. +struct BearerTokenFetch { + token: Resolvable, + fetch: Option, +} + +#[async_trait::async_trait] +impl Fetch for BearerTokenFetch { + async fn fetch(&self, mut request: FetchRequest) -> Result { + if !request.headers.contains_key("authorization") { + let token = self + .token + .resolve() + .await + .map_err(|error| FetchError::Other(error.to_string()))?; + request.headers.insert( + reqwest::header::AUTHORIZATION, + format!("Bearer {token}").parse().map_err( + |error: reqwest::header::InvalidHeaderValue| { + FetchError::Other(error.to_string()) + }, + )?, + ); + } + self.fetch + .clone() + .unwrap_or_else(default_fetch) + .fetch(request) + .await + } +} + /// The default provider: `create_azure` with default settings, created on /// first use. Creating it reads nothing from the environment and cannot fail; /// a missing key or resource name surfaces from the first request instead. @@ -235,13 +297,13 @@ pub fn azure() -> &'static AzureOpenAIProvider { pub struct AzureOpenAIProvider { resource_name: Option, base_url: Option, + speech_base_url: Option, + speech_headers: HeadersFn, info: BaseUrlInfo, api_version: Option, use_deployment_based_urls: bool, headers: HeadersFn, - token_provider: Option>, fetch: Option, - transform_request_body: Option, } /// Everything the URL of a request depends on, owned so a closure can hold it. @@ -349,10 +411,10 @@ impl AzureOpenAIProvider { provider: provider.to_string(), url: Arc::new(move |path| rules.url(path, &deployment)), headers: self.headers.clone(), - token_provider: self.token_provider.clone(), + token_provider: None, fetch: self.fetch.clone(), - supported_urls: aimux_core::language_model::SupportedUrls::default(), - transform_request_body: self.transform_request_body.clone(), + supported_urls: crate::openai::config::supported_urls(provider), + transform_request_body: None, responses: ResponsesProfile::default(), chat_options: options::CHAT_OPTIONS, } @@ -378,6 +440,7 @@ impl AzureOpenAIProvider { config.responses = ResponsesProfile { namespace: options::RESPONSES, file_id_prefixes: vec!["assistant-"], + explicit_message_item_type: self.info.is_foundry_project, }; OpenAIResponsesModel::from_config(deployment.to_string(), config) } @@ -402,13 +465,38 @@ impl AzureOpenAIProvider { ) } - /// A transcription model for a deployment; `provider()` is - /// `"azure.transcription"`. + /// A transcription model; `provider()` is `"azure.transcription"`. + /// The Azure Speech API is selected per call by `providerOptions.azure.api` + /// and defaults to Speech for the MAI transcription model, OpenAI otherwise. #[must_use] - pub fn transcription(&self, deployment: &str) -> OpenAITranscriptionModel { - OpenAITranscriptionModel::from_config( + pub fn transcription(&self, deployment: &str) -> AzureTranscriptionModel { + let mut speech = self.model_config("azure.transcription", deployment); + let rules = self.url_rules(); + let base_url = self.speech_base_url.clone(); + speech.url = Arc::new(move |_| { + let prefix = match &base_url { + Some(prefix) => without_trailing_slash(prefix), + None => { + // Reuse the resource-name validation and lazy environment lookup. + let mut rules = rules.clone(); + rules.base_url = None; + rules + .prefix()? + .replace(".openai.azure.com/openai", ".cognitiveservices.azure.com") + } + }; + Ok(format!( + "{prefix}/speechtotext/transcriptions:transcribe?api-version=2025-10-15" + )) + }); + speech.headers = self.speech_headers.clone(); + AzureTranscriptionModel::new( deployment.to_string(), - self.model_config("azure.transcription", deployment), + OpenAITranscriptionModel::from_config( + deployment.to_string(), + self.model_config("azure.transcription", deployment), + ), + speech, ) } diff --git a/aimux-providers/src/azure/options.rs b/aimux-providers/src/azure/options.rs index 7b713d90c..c6e823a0c 100644 --- a/aimux-providers/src/azure/options.rs +++ b/aimux-providers/src/azure/options.rs @@ -3,9 +3,9 @@ //! The AI SDK's Azure package reuses the OpenAI models. Its Responses model //! reads `providerOptions.azure` (and, when a call gave no `azure` options, //! `providerOptions.openai`, the functional fallback this port keeps) and -//! writes response metadata under `azure`; the other models read `openai` -//! only, like the OpenAI package itself. This module is the one place that -//! spells the Azure key. +//! writes response metadata under `azure`. Transcription reads `azure` for API +//! routing and Speech options and `openai` for OpenAI transcription options; +//! the other models read `openai` only. use crate::openai::responses::ResponsesNamespace; diff --git a/aimux-providers/src/azure/transcription.rs b/aimux-providers/src/azure/transcription.rs new file mode 100644 index 000000000..f6a13f991 --- /dev/null +++ b/aimux-providers/src/azure/transcription.rs @@ -0,0 +1,374 @@ +//! Per-request Azure transcription routing and the Azure Speech REST protocol. + +use std::collections::HashSet; + +use aimux_core::error::AiMuxError; +use aimux_core::shared::{SharedProviderOptions, Warning, provider_namespace}; +use aimux_core::transcription_model::{ + AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionResponse, + TranscriptionResult, TranscriptionSegment, TranscriptionStreamOptions, + TranscriptionStreamResult, +}; +use aimux_provider_utils::{HttpBody, MultipartForm}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; + +use super::options::NAMESPACE; +use crate::openai::{OpenAITranscriptionModel, config::OpenAIModelConfig}; + +#[derive(Default, Deserialize, Serialize)] +#[serde(deny_unknown_fields, rename_all = "camelCase")] +struct AzureOptions { + #[serde(skip_serializing_if = "Option::is_none")] + api: Option, + #[serde(skip_serializing_if = "Option::is_none")] + timestamps: Option, + #[serde(skip_serializing_if = "Option::is_none")] + transcribe_style: Option, + #[serde(skip_serializing_if = "Option::is_none")] + locales: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + diarization: Option, + #[serde(skip_serializing_if = "Option::is_none")] + phrase_list: Option, +} + +#[derive(Clone, Copy, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +enum Api { + Openai, + Speech, +} +#[derive(Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +enum Timestamps { + Word, + Segment, + None, +} +#[derive(Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +enum TranscribeStyle { + Verbatim, + Clean, +} +#[derive(Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +struct Diarization { + enabled: bool, +} +#[derive(Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +struct PhraseList { + phrases: Vec, +} + +fn is_default_speech(model_id: &str) -> bool { + model_id.eq_ignore_ascii_case("mai-transcribe-2") +} + +fn parse_options(options: Option<&SharedProviderOptions>) -> Result { + let Some(azure) = options.and_then(|options| options.get(NAMESPACE)) else { + return Ok(AzureOptions::default()); + }; + // Optional in the upstream schema means absent, never explicitly null. + if azure.values().any(Value::is_null) { + return Err(AiMuxError::InvalidArgument( + "Invalid Azure transcription options: null field".into(), + )); + } + let parsed: AzureOptions = + serde_json::from_value(Value::Object(azure.clone())).map_err(|error| { + AiMuxError::InvalidArgument(format!("Invalid Azure transcription options: {error}")) + })?; + if parsed + .locales + .as_ref() + .is_some_and(|locales| locales.len() != 1) + { + return Err(AiMuxError::InvalidArgument( + "Azure transcription locales must contain one language".into(), + )); + } + Ok(parsed) +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct SpeechResponse { + combined_phrases: Vec, + duration_milliseconds: Option, + phrases: Option>, +} +#[derive(Deserialize)] +struct CombinedPhrase { + text: String, +} +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct Phrase { + text: String, + offset_milliseconds: Option, + duration_milliseconds: Option, + locale: Option, + #[serde(rename = "speaker")] + _speaker: Option, + #[serde(rename = "confidence")] + _confidence: Option, + #[serde(rename = "words")] + _words: Option>, +} +#[derive(Deserialize)] +struct Word { + #[serde(rename = "text")] + _text: String, + #[serde(rename = "offsetMilliseconds")] + _offset_milliseconds: Option, + #[serde(rename = "durationMilliseconds")] + _duration_milliseconds: Option, +} + +/// Azure's transcription model, resolving the API from provider options on every call. +pub struct AzureTranscriptionModel { + model_id: String, + openai: OpenAITranscriptionModel, + speech: OpenAIModelConfig, +} + +impl AzureTranscriptionModel { + pub(crate) fn new( + model_id: String, + openai: OpenAITranscriptionModel, + speech: OpenAIModelConfig, + ) -> Self { + Self { + model_id, + openai, + speech, + } + } + + fn uses_speech(&self, options: &AzureOptions) -> bool { + matches!( + options.api.unwrap_or(if is_default_speech(&self.model_id) { + Api::Speech + } else { + Api::Openai + }), + Api::Speech + ) + } +} + +#[async_trait::async_trait] +impl TranscriptionModel for AzureTranscriptionModel { + fn provider(&self) -> &str { + "azure.transcription" + } + fn model_id(&self) -> &str { + &self.model_id + } + + async fn do_generate( + &self, + options: &TranscriptionCallOptions, + ) -> Result { + let mut azure = parse_options(options.provider_options.as_ref())?; + if !self.uses_speech(&azure) { + let mut result = self.openai.do_generate(options).await?; + azure.api = None; + let speech_options = serde_json::to_value(azure).expect("serializable options"); + for key in [ + "timestamps", + "transcribeStyle", + "locales", + "diarization", + "phraseList", + ] { + if speech_options.get(key).is_some() { + result.warnings.push(Warning::Unsupported { + feature: format!("providerOptions.azure.{key}"), + details: Some("This option requires the Azure Speech API.".into()), + }); + } + } + return Ok(result); + } + let timestamp = chrono::Utc::now().to_rfc3339(); + let audio = match &options.audio { + AudioInput::Binary(audio) => audio.clone(), + AudioInput::Base64(audio) => { + let audio: String = audio + .chars() + .filter(|character| !matches!(character, '\t' | '\n' | '\x0c' | '\r' | ' ')) + .map(|character| match character { + '-' => '+', + '_' => '/', + other => other, + }) + .collect(); + if audio.ends_with('=') && !audio.len().is_multiple_of(4) { + return Err(AiMuxError::InvalidArgument("invalid base64 padding".into())); + } + let config = base64::engine::general_purpose::GeneralPurposeConfig::new() + .with_decode_padding_mode(base64::engine::DecodePaddingMode::Indifferent) + .with_decode_allow_trailing_bits(true); + base64::Engine::decode( + &base64::engine::general_purpose::GeneralPurpose::new( + &base64::alphabet::STANDARD, + config, + ), + audio, + ) + .map_err(|error| AiMuxError::InvalidArgument(format!("invalid base64: {error}")))? + } + }; + let mut definition = json!({"enhancedMode": { + "enabled": true, + "model": if is_default_speech(&self.model_id) { "MAI-Transcribe-2" } else { &self.model_id }, + "modelOptions": {"timestamps": azure.timestamps.take().unwrap_or(Timestamps::Segment)}, + }}); + if let Some(style) = azure.transcribe_style.take() { + definition["enhancedMode"]["modelOptions"]["transcribeStyle"] = json!(style); + } + azure.api = None; + definition + .as_object_mut() + .expect("definition object") + .extend( + serde_json::to_value(azure) + .expect("serializable options") + .as_object() + .expect("options object") + .clone(), + ); + let media_type = options.media_type.to_lowercase(); + let subtype = media_type.split('/').nth(1).unwrap_or_default(); + let extension = match subtype { + "mpeg" => "mp3", + "x-wav" => "wav", + "opus" => "ogg", + "mp4" | "x-m4a" => "m4a", + other => other, + }; + let mut form = MultipartForm::new(); + form.file( + "audio", + &format!("audio.{extension}"), + &options.media_type, + &audio, + )?; + form.text("definition", &definition.to_string())?; + let (body, content_type) = form.finish(); + let response = aimux_provider_utils::post_to_api( + self.speech.http_request( + self.speech.url("")?, + self.speech + .request_headers(options.headers.as_ref()) + .await?, + options, + ), + HttpBody::Bytes(body, content_type), + aimux_provider_utils::create_json_response_handler::(), + aimux_provider_utils::create_standard_json_error_response_handler(), + ) + .await?; + let parsed = response.value; + let phrases = parsed.phrases.unwrap_or_default(); + let languages: HashSet<_> = phrases + .iter() + .filter_map(|phrase| phrase.locale.as_ref()) + .map(|locale| locale.split('-').next().unwrap_or_default().to_lowercase()) + .collect(); + let language = if languages.len() == 1 { + languages.into_iter().next().filter(|language| { + language.len() == 2 && language.bytes().all(|byte| byte.is_ascii_lowercase()) + }) + } else { + None + }; + let segments = phrases + .iter() + .filter_map(|phrase| { + Some(TranscriptionSegment { + text: phrase.text.clone(), + start_second: phrase.offset_milliseconds? / 1000.0, + end_second: (phrase.offset_milliseconds? + phrase.duration_milliseconds?) + / 1000.0, + }) + }) + .collect(); + let raw_body = response.raw_value.unwrap_or(Value::Null); + // Schema parsing strips unknown keys but retains absent versus null fields. + let mut metadata_phrases = raw_body + .get("phrases") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + for phrase in &mut metadata_phrases { + if let Some(phrase) = phrase.as_object_mut() { + phrase.retain(|key, _| { + matches!( + key.as_str(), + "text" + | "offsetMilliseconds" + | "durationMilliseconds" + | "locale" + | "speaker" + | "confidence" + | "words" + ) + }); + if let Some(words) = phrase.get_mut("words").and_then(Value::as_array_mut) { + for word in words { + if let Some(word) = word.as_object_mut() { + word.retain(|key, _| { + matches!( + key.as_str(), + "text" | "offsetMilliseconds" | "durationMilliseconds" + ) + }); + } + } + } + } + } + Ok(TranscriptionResult { + text: parsed + .combined_phrases + .into_iter() + .map(|phrase| phrase.text) + .collect::>() + .join(" "), + segments, + language, + duration_in_seconds: parsed + .duration_milliseconds + .map(|duration| duration / 1000.0), + warnings: Vec::new(), + request: None, + provider_metadata: Some(provider_namespace( + NAMESPACE, + json!({"phrases": metadata_phrases}), + )?), + response: TranscriptionResponse { + timestamp: Some(timestamp), + model_id: Some(self.model_id.clone()), + headers: Some(response.response_headers), + body: Some(raw_body), + }, + }) + } + + async fn do_stream( + &self, + options: TranscriptionStreamOptions, + ) -> Result { + if self.uses_speech(&parse_options(options.provider_options.as_ref())?) { + return Err(AiMuxError::UnsupportedFunctionality( + "streaming transcription with the Azure Speech API".into(), + )); + } + self.openai.do_stream(options).await + } +} diff --git a/aimux-providers/src/bedrock/convert.rs b/aimux-providers/src/bedrock/convert.rs index 8d0b4b20e..f1f150749 100644 --- a/aimux-providers/src/bedrock/convert.rs +++ b/aimux-providers/src/bedrock/convert.rs @@ -17,12 +17,13 @@ use super::options; use aimux_core::language_model_message::{ AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, - ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, UserPart, + ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, }; -use aimux_core::options::{CallOptions, ToolChoice}; +use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified}; +use aimux_provider_utils::is_full_media_type; use base64::Engine; use serde_json::{Value, json}; @@ -40,11 +41,18 @@ use serde_json::{Value, json}; /// /// # Errors /// -/// Returns an unsupported-functionality error for unsupported tool result content. +/// Returns an error for prompt content the Converse API cannot express. pub fn convert_prompt_to_bedrock( prompt: &LanguageModelPrompt, -) -> Result<(Vec, Vec), aimux_core::error::AiMuxError> { - validate_tool_result_content(prompt)?; +) -> Result<(Vec, Vec), aimux_core::AiMuxError> { + convert_prompt(prompt, false) +} + +fn convert_prompt( + prompt: &LanguageModelPrompt, + is_mistral: bool, +) -> Result<(Vec, Vec), aimux_core::AiMuxError> { + validate_prompt(prompt)?; #[derive(Clone, Copy, PartialEq, Eq)] enum Blk { System, @@ -85,7 +93,7 @@ pub fn convert_prompt_to_bedrock( unreachable!() }; system.push(json!({ "text": content })); - if let Some(cp) = cache_point(provider_options) { + if let Some(cp) = cache_point(provider_options.as_ref()) { system.push(cp); } } @@ -93,56 +101,70 @@ pub fn convert_prompt_to_bedrock( Blk::User => { let mut content: Vec = Vec::new(); for &i in idxs { - match &prompt[i] { - LanguageModelMessage::User { content: parts, .. } => { + let provider_options = match &prompt[i] { + LanguageModelMessage::User { + content: parts, + provider_options, + } => { for part in parts { - let provider_options = match part { + let po = match part { UserPart::Text(part) => { - content.push(json!({ "text": part.text })); + push_user_text(part, &mut content); &part.provider_options } UserPart::File(file) => { - if !push_file_part( - file, - &mut content, - &mut document_counter, - ) { + if matches!(file.data, FileData::Reference { .. }) { continue; } + push_file_part(file, &mut content, &mut document_counter); &file.provider_options } }; - if let Some(cp) = cache_point(provider_options) { + if let Some(cp) = cache_point(po.as_ref()) { content.push(cp); } } + provider_options } - LanguageModelMessage::Tool { content: parts, .. } => { + LanguageModelMessage::Tool { + content: parts, + provider_options, + } => { for part in parts { let ToolPart::ToolResult(part) = part else { continue; }; - let result_content = - resolve_tool_result_output(&part.output, &mut document_counter); - content.push(json!({ - "toolResult": { - "toolUseId": part.tool_call_id, - "content": result_content, - } - })); + push_tool_result( + part, + &mut content, + &mut document_counter, + is_mistral, + ); + if let Some(cp) = cache_point(part.provider_options.as_ref()) { + content.push(cp); + } } + provider_options } _ => unreachable!(), + }; + if let Some(cp) = cache_point(provider_options.as_ref()) { + content.push(cp); } } append_user_message(&mut messages, content); } Blk::Assistant => { let mut content: Vec = Vec::new(); + let mut results: Vec = Vec::new(); let num_msgs = idxs.len(); for (mj, &i) in idxs.iter().enumerate() { let is_last_message = mj == num_msgs - 1; - let LanguageModelMessage::Assistant { content: parts, .. } = &prompt[i] else { + let LanguageModelMessage::Assistant { + content: parts, + provider_options, + } = &prompt[i] + else { unreachable!() }; let has_reasoning = parts @@ -151,6 +173,9 @@ pub fn convert_prompt_to_bedrock( let num_parts = parts.len(); for (kj, p) in parts.iter().enumerate() { let is_last_content_part = kj == num_parts - 1; + if !matches!(p, AssistantPart::ToolResult(_)) && !results.is_empty() { + append_user_message(&mut messages, std::mem::take(&mut results)); + } match p { AssistantPart::Text(TextPart { text, .. }) => { // Skip empty text unless the message has reasoning. @@ -169,45 +194,32 @@ pub fn convert_prompt_to_bedrock( text, provider_options, }) => { - let options = options::read(provider_options.as_ref()); - let Some(sig) = options - .and_then(|options| options.get("signature")) + let metadata = options::read(provider_options.as_ref()); + if let Some(sig) = metadata + .and_then(|v| v.get("signature")) .and_then(Value::as_str) - else { - if let Some(redacted) = - options.and_then(|options| options.get("redactedContent")) - { - content.push(json!({ "reasoningContent": { "redactedContent": redacted } })); - } else if let Some(redacted) = - options.and_then(|options| options.get("redactedData")) - { - content.push(json!({ "reasoningContent": { "redactedReasoning": { "data": redacted } } })); - } - continue; - }; - // Only signed reasoning is replayed; unsigned - // reasoning is intentionally omitted. - content.push(json!({ - "reasoningContent": { - "reasoningText": { - "text": text, - "signature": sig - } - } - })); + { + content.push(json!({"reasoningContent":{"reasoningText":{"text":text,"signature":sig}}})); + } else if let Some(redacted) = + metadata.and_then(|v| v.get("redactedContent")) + { + content.push( + json!({"reasoningContent":{"redactedContent":redacted}}), + ); + } else if let Some(data) = + metadata.and_then(|v| v.get("redactedData")) + { + content.push(json!({"reasoningContent":{"redactedReasoning":{"data":data}}})); + } } AssistantPart::ToolResult(part) => { - if !content.is_empty() { - messages.push(json!({ "role": "assistant", "content": std::mem::take(&mut content) })); - } - let mut result = vec![ - json!({ "toolResult": { "toolUseId": part.tool_call_id, "content": resolve_tool_result_output(&part.output, &mut document_counter) } }), - ]; - if let Some(cp) = cache_point(&part.provider_options) { - result.push(cp); - } - append_user_message(&mut messages, result); - continue; + flush_assistant(&mut messages, &mut content); + push_tool_result( + part, + &mut results, + &mut document_counter, + is_mistral, + ); } AssistantPart::ToolCall(ToolCallPart { tool_call_id, @@ -222,7 +234,7 @@ pub fn convert_prompt_to_bedrock( }; content.push(json!({ "toolUse": { - "toolUseId": tool_call_id, + "toolUseId": normalize_tool_call_id(tool_call_id, is_mistral), "name": sanitize_tool_name(tool_name), "input": input_val, } @@ -230,26 +242,26 @@ pub fn convert_prompt_to_bedrock( } _ => {} } - let provider_options = match p { - AssistantPart::Text(part) => &part.provider_options, - AssistantPart::File(part) - if matches!( - part.data, - FileData::Data { .. } | FileData::Text { .. } - ) => - { - &part.provider_options + if let Some(cp) = assistant_cache_point(p) { + if matches!(p, AssistantPart::ToolResult(_)) { + results.push(cp); + } else { + content.push(cp); } - _ => continue, - }; - if let Some(cp) = cache_point(provider_options) { + } + } + if let Some(cp) = cache_point(provider_options.as_ref()) { + if matches!(parts.last(), Some(AssistantPart::ToolResult(_))) { + results.push(cp); + } else { content.push(cp); } } } - if !content.is_empty() { - messages.push(json!({ "role": "assistant", "content": content })); + if !results.is_empty() { + append_user_message(&mut messages, results); } + flush_assistant(&mut messages, &mut content); } } } @@ -257,118 +269,39 @@ pub fn convert_prompt_to_bedrock( Ok((system, messages)) } -fn append_user_message(messages: &mut Vec, content: Vec) { - if let Some(last) = messages.last_mut() - && last.get("role").and_then(Value::as_str) == Some("user") - && let Some(parts) = last.get_mut("content").and_then(Value::as_array_mut) +fn push_user_text(part: &TextPart, content: &mut Vec) { + let opts = options::read(part.provider_options.as_ref()); + if opts + .and_then(|v| v.get("guardContent")) + .and_then(Value::as_bool) + == Some(true) { - parts.extend(content); + let mut block = json!({"text":part.text}); + if let Some(q) = opts.and_then(|v| v.get("guardContentQualifiers")) { + block["qualifiers"] = q.clone(); + } + content.push(json!({"guardContent":{"text":block}})); } else { - messages.push(json!({ "role": "user", "content": content })); + content.push(json!({ "text": part.text })); } } -pub(crate) fn validate_tool_result_content( - prompt: &LanguageModelPrompt, -) -> Result<(), aimux_core::error::AiMuxError> { - for message in prompt { - let results: Vec<&ToolResultOutput> = match message { - LanguageModelMessage::Tool { content, .. } => content - .iter() - .filter_map(|part| { - if let ToolPart::ToolResult(part) = part { - Some(&part.output) - } else { - None - } - }) - .collect(), - LanguageModelMessage::Assistant { content, .. } => content - .iter() - .filter_map(|part| { - if let AssistantPart::ToolResult(part) = part { - Some(&part.output) - } else { - None - } - }) - .collect(), - _ => continue, - }; - for output in results { - if let ToolResultOutput::Content { value } = output { - for part in value { - let unsupported = match part { - ToolResultContent::Text(_) => false, - ToolResultContent::File(file) => { - let allowed_data = matches!(file.data, FileData::Data { .. }) - || matches!(&file.data, FileData::Url { url, .. } if url.starts_with("s3:")); - if !allowed_data { - return Err( - aimux_core::error::AiMuxError::UnsupportedFunctionality( - "tool result content part".to_string(), - ), - ); - } - let media_type = crate::google::convert::tool_file_media_type(file)?; - let supported = matches!( - media_type.as_str(), - "image/jpeg" - | "image/png" - | "image/gif" - | "image/webp" - | "video/x-matroska" - | "video/quicktime" - | "video/mp4" - | "video/webm" - | "video/x-flv" - | "video/mpeg" - | "video/mpg" - | "video/wmv" - | "video/x-ms-wmv" - | "video/3gpp" - | "application/pdf" - | "text/csv" - | "application/msword" - | "application/vnd.openxmlformats-officedocument.wordprocessingml.document" - | "application/vnd.ms-excel" - | "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" - | "text/html" - | "text/plain" - | "text/markdown" - ); - if !supported { - let kind = match media_type.split('/').next() { - Some("image") => "image", - Some("video") => "video", - _ => "file", - }; - return Err( - aimux_core::error::AiMuxError::UnsupportedFunctionality( - format!("{kind} mime type: {media_type}"), - ), - ); - } - matches!(file.data, FileData::Url { .. }) - && !(media_type.starts_with("image/") - || media_type.starts_with("video/")) - } - ToolResultContent::Custom { .. } => true, - }; - if unsupported { - return Err(aimux_core::error::AiMuxError::UnsupportedFunctionality( - "tool result content part".to_string(), - )); - } - } - } +fn push_tool_result( + part: &ToolResultPart, + content: &mut Vec, + doc_counter: &mut u32, + is_mistral: bool, +) { + let result_content = resolve_tool_result_output(&part.output, doc_counter); + content.push(json!({ + "toolResult": { + "toolUseId": normalize_tool_call_id(&part.tool_call_id, is_mistral), + "content": result_content, } - } - Ok(()) + })); } -/// Push a supported file into a Bedrock content array. -fn push_file_part(file: &FilePart, content: &mut Vec, doc_counter: &mut u32) -> bool { +fn push_file_part(file: &FilePart, content: &mut Vec, doc_counter: &mut u32) { let b64 = match &file.data { FileData::Data { data: FileBytes::Binary(bytes), @@ -377,40 +310,39 @@ fn push_file_part(file: &FilePart, content: &mut Vec, doc_counter: &mut u data: FileBytes::Base64(data), } => data.clone(), FileData::Text { text } => base64::engine::general_purpose::STANDARD.encode(text), - FileData::Url { url, .. } - if url.starts_with("s3:") && file.media_type.starts_with("image/") => - { - content.push(json!({ - "image": { - "format": mime_to_image_format(&file.media_type), - "source": { "s3Location": { "uri": url } } - } - })); - return true; - } - FileData::Url { .. } | FileData::Reference { .. } => return false, - }; - let is_text = matches!(file.data, FileData::Text { .. }); - let media_type = if is_text - && !file - .media_type - .split_once('/') - .is_some_and(|(_, subtype)| !subtype.is_empty() && subtype != "*") - { - "text/plain" - } else { - &file.media_type + FileData::Url { .. } => String::new(), + FileData::Reference { .. } => return, }; + let resolved_media_type = aimux_provider_utils::resolve_full_media_type(file).ok(); + let media_type = + if matches!(file.data, FileData::Text { .. }) && !is_full_media_type(&file.media_type) { + "text/plain" + } else { + resolved_media_type.as_deref().unwrap_or(&file.media_type) + }; push_file_block( &b64, media_type, file.filename.as_deref(), &file.provider_options, - is_text, content, doc_counter, ); - true + if let FileData::Url { url, .. } = &file.data + && let Some(block) = content.last_mut() + { + let inner = if block.get("guardContent").is_some() { + &mut block["guardContent"] + } else { + block + }; + let key = if media_type.starts_with("video/") { + "video" + } else { + "image" + }; + inner[key]["source"] = json!({"s3Location":{"uri":url}}); + } } /// Build a Bedrock `image` or `document` block from an already-base64 `bytes` @@ -423,21 +355,31 @@ fn push_file_block( media_type: &str, filename: Option<&str>, provider_options: &Option, - is_text: bool, content: &mut Vec, doc_counter: &mut u32, ) { let top_level = media_type.split('/').next().unwrap_or(""); - if top_level == "image" && !is_text { + if top_level == "image" { let format = mime_to_image_format(media_type); - content.push(json!({ - "image": { "format": format, "source": { "bytes": b64 } } - })); + let block = json!({ "image": { "format": format, "source": { "bytes": b64 } } }); + if options::read(provider_options.as_ref()) + .and_then(|v| v.get("guardContent")) + .and_then(Value::as_bool) + == Some(true) + { + content.push(json!({"guardContent":block})); + } else { + content.push(block); + } + } else if top_level == "video" { + content.push( + json!({"video":{"format":mime_to_video_format(media_type),"source":{"bytes":b64}}}), + ); } else { let format = mime_to_document_format(media_type); let name = match filename { - Some(f) => strip_file_extension(f), - None => { + Some(f) if !sanitize_document_name(f).is_empty() => sanitize_document_name(f), + _ => { *doc_counter += 1; format!("document-{}", *doc_counter) } @@ -455,9 +397,18 @@ fn push_file_block( /// Extract a `{ cachePoint: {...} }` block from a part's `providerOptions` /// (`amazonBedrock.cachePoint`), if present. -fn cache_point(provider_options: &Option) -> Option { - let cp = options::read(provider_options.as_ref())?.get("cachePoint")?; - Some(json!({ "cachePoint": cp.clone() })) +fn assistant_cache_point(part: &AssistantPart) -> Option { + let po = match part { + AssistantPart::Text(part) => &part.provider_options, + AssistantPart::File(part) if !matches!(part.data, FileData::Reference { .. }) => { + &part.provider_options + } + AssistantPart::Reasoning(part) => &part.provider_options, + AssistantPart::ToolCall(part) => &part.provider_options, + AssistantPart::ToolResult(part) => &part.provider_options, + _ => return None, + }; + cache_point(po.as_ref()) } /// Whether `citations.enabled` is set on a part's `amazonBedrock` provider @@ -490,12 +441,13 @@ fn resolve_tool_result_output(output: &ToolResultOutput, doc_counter: &mut u32) ToolResultContent::Text(part) => content.push(json!({ "text": part.text })), ToolResultContent::File(part) => { let mut resolved = part.clone(); - if let Ok(media_type) = crate::google::convert::tool_file_media_type(part) { + if let Ok(media_type) = aimux_provider_utils::resolve_full_media_type(part) + { resolved.media_type = media_type; } let part = &resolved; if let FileData::Url { url, .. } = &part.data { - let format = part.media_type.split('/').nth(1).unwrap_or(""); + let format = mime_to_video_format(&part.media_type).unwrap_or(""); let source = json!({ "s3Location": { "uri": url } }); content.push(if part.media_type.starts_with("image/") { json!({ "image": { "format": mime_to_image_format(&part.media_type), "source": source } }) } else { json!({ "video": { "format": format, "source": source } }) }); } else if part.media_type.starts_with("video/") { @@ -506,7 +458,7 @@ fn resolve_tool_result_output(output: &ToolResultOutput, doc_counter: &mut u32) } FileBytes::Base64(data) => data.clone(), }; - content.push(json!({ "video": { "format": part.media_type.split('/').nth(1).unwrap_or(""), "source": { "bytes": bytes } } })); + content.push(json!({ "video": { "format": mime_to_video_format(&part.media_type).unwrap_or(""), "source": { "bytes": bytes } } })); } } else { push_file_part(part, &mut content, doc_counter); @@ -523,7 +475,7 @@ fn resolve_tool_result_output(output: &ToolResultOutput, doc_counter: &mut u32) fn sanitize_tool_name(name: &str) -> String { let sanitized: String = name .chars() - .filter(|c| c.is_alphanumeric() || *c == '_' || *c == '-') + .filter(|c| c.is_ascii_alphanumeric() || *c == '_' || *c == '-') .collect(); if sanitized.is_empty() { "_".to_string() @@ -534,7 +486,7 @@ fn sanitize_tool_name(name: &str) -> String { fn strip_file_extension(name: &str) -> String { match name.rfind('.') { - Some(i) if i > 0 => name[..i].to_string(), + Some(i) => name[..i].to_string(), _ => name.to_string(), } } @@ -564,6 +516,341 @@ fn mime_to_document_format(media_type: &str) -> &'static str { } } +fn cache_point(provider_options: Option<&SharedProviderOptions>) -> Option { + Some(json!({"cachePoint": options::read(provider_options)?.get("cachePoint")?})) +} + +fn append_user_message(messages: &mut Vec, content: Vec) { + if let Some(last) = messages.last_mut() + && last["role"] == "user" + && last["content"] + .as_array() + .is_some_and(|blocks| blocks.iter().any(|b| b.get("toolResult").is_some())) + { + last["content"].as_array_mut().unwrap().extend(content); + } else { + messages.push(json!({"role":"user","content":content})); + } +} + +fn flush_assistant(messages: &mut Vec, content: &mut Vec) { + if content.iter().any(|v| v.get("cachePoint").is_none()) { + messages.push(json!({"role":"assistant","content":std::mem::take(content)})); + } else { + content.clear(); + } +} + +#[must_use] +pub fn normalize_tool_call_id(id: &str, is_mistral: bool) -> String { + if !is_mistral || (id.len() == 9 && id.bytes().all(|c| c.is_ascii_alphanumeric())) { + return id.to_owned(); + } + let mut hash = 14_695_981_039_346_656_037_u64; + for unit in id.encode_utf16() { + hash = (hash ^ u64::from(unit)).wrapping_mul(1_099_511_628_211); + } + let mut value = hash % 62_u64.pow(9); + let mut bytes = [b'0'; 9]; + let alphabet = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; + for byte in bytes.iter_mut().rev() { + *byte = alphabet[(value % 62) as usize]; + value /= 62; + } + String::from_utf8(bytes.to_vec()).unwrap() +} + +fn sanitize_document_name(name: &str) -> String { + strip_file_extension(name) + .split_whitespace() + .collect::>() + .join(" ") + .chars() + .filter(|c| c.is_ascii_alphanumeric() || " ()[]-".contains(*c)) + .collect::() + .trim() + .chars() + .take(200) + .collect::() + .trim() + .to_owned() +} + +fn mime_to_video_format(media_type: &str) -> Option<&'static str> { + match media_type { + "video/x-matroska" => Some("mkv"), + "video/quicktime" => Some("mov"), + "video/mp4" => Some("mp4"), + "video/webm" => Some("webm"), + "video/x-flv" => Some("flv"), + "video/mpeg" => Some("mpeg"), + "video/mpg" => Some("mpg"), + "video/wmv" | "video/x-ms-wmv" => Some("wmv"), + "video/3gpp" => Some("three_gp"), + _ => None, + } +} + +fn is_strict_schema_compatible(schema: &Value) -> bool { + let Some(object) = schema.as_object() else { + return true; + }; + let is_object = schema["type"] == "object" + || schema["type"] + .as_array() + .is_some_and(|a| a.iter().any(|v| v == "object")); + if is_object && schema["additionalProperties"] != false { + return false; + } + for key in [ + "properties", + "patternProperties", + "definitions", + "$defs", + "dependencies", + ] { + if let Some(map) = object.get(key).and_then(Value::as_object) + && map.values().any(|v| !is_strict_schema_compatible(v)) + { + return false; + } + } + for key in [ + "propertyNames", + "contains", + "not", + "if", + "then", + "else", + "items", + "anyOf", + "allOf", + "oneOf", + ] { + if let Some(value) = object.get(key) { + if let Some(array) = value.as_array() { + if array.iter().any(|v| !is_strict_schema_compatible(v)) { + return false; + } + } else if !is_strict_schema_compatible(value) { + return false; + } + } + } + true +} + +fn validate_media_type( + media_type: &str, + is_url: bool, + url: Option<&str>, +) -> Result<(), aimux_core::AiMuxError> { + let supported = match media_type.split('/').next() { + Some("image") => matches!( + media_type, + "image/png" | "image/jpeg" | "image/gif" | "image/webp" + ), + Some("video") => mime_to_video_format(media_type).is_some(), + _ => matches!( + media_type, + "application/pdf" + | "text/csv" + | "application/msword" + | "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + | "application/vnd.ms-excel" + | "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" + | "text/html" + | "text/plain" + | "text/markdown" + ), + }; + if !supported { + let kind = match media_type.split('/').next() { + Some("image") => "image", + Some("video") => "video", + _ => "file", + }; + return Err(aimux_core::AiMuxError::UnsupportedFunctionality(format!( + "{kind} mime type: {media_type}" + ))); + } + if is_url + && (!url.is_some_and(|u| u.starts_with("s3:")) + || !(media_type.starts_with("image/") || media_type.starts_with("video/"))) + { + return Err(aimux_core::AiMuxError::UnsupportedFunctionality(format!( + "file media type or data: {media_type}" + ))); + } + Ok(()) +} + +/// Reject prompts the Converse API cannot express. +/// +/// # Errors +/// +/// An unsupported message arrangement. +pub fn validate_prompt(prompt: &LanguageModelPrompt) -> Result<(), aimux_core::AiMuxError> { + let mut saw_non_system = false; + for message in prompt { + match message { + LanguageModelMessage::System { .. } => { + if saw_non_system { + return Err(aimux_core::AiMuxError::UnsupportedFunctionality( + "Multiple system messages that are separated by user/assistant messages" + .into(), + )); + } + } + LanguageModelMessage::User { content, .. } => { + saw_non_system = true; + for part in content { + match part { + UserPart::Text(part) => { + validate_part_options(part.provider_options.as_ref(), false)? + } + UserPart::File(file) => validate_file(file)?, + } + } + } + LanguageModelMessage::Assistant { content, .. } => { + saw_non_system = true; + for part in content { + match part { + AssistantPart::Text(part) => { + validate_part_options(part.provider_options.as_ref(), false)? + } + AssistantPart::File(file) => validate_file(file)?, + AssistantPart::Reasoning(part) => { + validate_part_options(part.provider_options.as_ref(), true)? + } + AssistantPart::ToolResult(part) => validate_tool_result(&part.output)?, + AssistantPart::ToolCall(_) + | AssistantPart::ReasoningFile(_) + | AssistantPart::Custom(_) => {} + } + } + } + LanguageModelMessage::Tool { content, .. } => { + saw_non_system = true; + for part in content { + if let ToolPart::ToolResult(part) = part { + validate_tool_result(&part.output)?; + } + } + } + } + } + Ok(()) +} + +fn validate_part_options( + provider_options: Option<&SharedProviderOptions>, + reasoning: bool, +) -> Result<(), aimux_core::AiMuxError> { + if let Some(options) = options::read(provider_options) { + if let Some(value) = options.get("guardContent") + && !value.is_boolean() + { + return Err(aimux_core::AiMuxError::InvalidArgument( + "Invalid Bedrock guardContent".into(), + )); + } + if let Some(value) = options.get("guardContentQualifiers") + && !value.as_array().is_some_and(|a| { + a.iter().all(|v| { + matches!( + v.as_str(), + Some("grounding_source" | "query" | "guard_content") + ) + }) + }) + { + return Err(aimux_core::AiMuxError::InvalidArgument( + "Invalid Bedrock guardContentQualifiers".into(), + )); + } + if let Some(value) = options.get("citations") + && (!value.is_object() || value.get("enabled").is_some_and(|v| !v.is_boolean())) + { + return Err(aimux_core::AiMuxError::InvalidArgument( + "Invalid Bedrock citations".into(), + )); + } + if reasoning { + for key in ["signature", "redactedContent", "redactedData"] { + if let Some(value) = options.get(key) + && !value.is_string() + && !value.is_null() + { + return Err(aimux_core::AiMuxError::InvalidArgument(format!( + "Invalid Bedrock {key}" + ))); + } + } + } + } + Ok(()) +} + +fn validate_file(file: &FilePart) -> Result<(), aimux_core::AiMuxError> { + if !matches!(file.data, FileData::Reference { .. }) { + validate_part_options(file.provider_options.as_ref(), false)?; + } + match &file.data { + FileData::Reference { .. } => Err(aimux_core::AiMuxError::UnsupportedFunctionality( + "File reference data".into(), + )), + FileData::Url { url, .. } => validate_media_type(&file.media_type, true, Some(url)), + FileData::Text { .. } => { + let media_type = if is_full_media_type(&file.media_type) { + &file.media_type + } else { + "text/plain" + }; + if media_type.starts_with("image/") || media_type.starts_with("video/") { + return Err(aimux_core::AiMuxError::UnsupportedFunctionality(format!( + "file media type or data: {media_type}" + ))); + } + validate_media_type(media_type, false, None) + } + FileData::Data { .. } => { + let media_type = aimux_provider_utils::resolve_full_media_type(file)?; + validate_media_type(&media_type, false, None) + } + } +} + +fn validate_tool_result(result: &ToolResultOutput) -> Result<(), aimux_core::AiMuxError> { + if let ToolResultOutput::Content { value } = result { + for part in value { + match part { + ToolResultContent::Text(_) => {} + ToolResultContent::File(file) => { + let (is_url, url) = match &file.data { + FileData::Data { .. } => (false, None), + FileData::Url { url, .. } => (true, Some(url.as_str())), + _ => { + return Err(aimux_core::AiMuxError::UnsupportedFunctionality( + "tool result file data".into(), + )); + } + }; + let media_type = aimux_provider_utils::resolve_full_media_type(file)?; + validate_media_type(&media_type, is_url, url)?; + } + ToolResultContent::Custom { .. } => { + return Err(aimux_core::AiMuxError::UnsupportedFunctionality( + "unsupported tool content part".into(), + )); + } + } + } + } + Ok(()) +} + // ── Tools ─────────────────────────────────────────────────────────────────── /// Whether a Bedrock model id supports the `strict` tool schema field. @@ -572,14 +859,28 @@ fn mime_to_document_format(media_type: &str) -> &'static str { /// `strict` (and `output_config.format`), so the field is omitted for them. #[must_use] pub fn supports_strict_tools(model_id: &str) -> bool { - const REJECTING: &[&str] = &[ - "claude-opus-4-7", - "claude-opus-4-8", - "claude-opus-5", - "claude-fable-5", - "claude-sonnet-5", - ]; - !REJECTING.iter().any(|m| model_id.contains(m)) + !MODELS_WITHOUT_STRICT_TOOL_SUPPORT + .iter() + .any(|m| model_id.contains(m)) +} + +const MODELS_WITHOUT_STRICT_TOOL_SUPPORT: &[&str] = &[ + "claude-opus-4-7", + "claude-opus-4-8", + "claude-opus-5", + "claude-fable-5", + "claude-sonnet-5", +]; + +fn supports_native_structured_output(model_id: &str) -> bool { + let prefix = MODELS_WITHOUT_STRICT_TOOL_SUPPORT[0] + .split_once('-') + .unwrap() + .0; + supports_strict_tools(model_id) + && !["sonnet-4-6", "haiku-4-5"] + .iter() + .any(|suffix| model_id.contains(&format!("{prefix}-{suffix}"))) } /// Prepare `FunctionTool`s into the Bedrock `toolConfig` JSON shape. @@ -592,9 +893,6 @@ pub fn supports_strict_tools(model_id: &str) -> bool { /// - `description` is omitted when empty/whitespace /// - `strict` is passed through only for models that `supports_strict_tools` /// -/// Provider-defined tools (web_search, anthropic provider tools) and the -/// `additionalTools`/`betas`/`toolWarnings` they produce are not modelled in -/// the Rust `FunctionTool` and are intentionally not handled here. #[must_use] pub fn prepare_tools( tools: &Option>, @@ -632,6 +930,7 @@ pub fn prepare_tools( } if let Some(strict) = t.strict && supports_strict + && (!strict || is_strict_schema_compatible(&t.input_schema)) { spec.insert("strict".to_string(), json!(strict)); } @@ -657,113 +956,687 @@ pub fn prepare_tools( // ── Request body ──────────────────────────────────────────────────────────── -/// Build the Bedrock Converse request body. -/// -/// When `providerOptions.bedrock.reasoningConfig = { type: 'enabled', -/// budgetTokens: N }` is present, it is translated into -/// `additionalModelRequestFields.thinking = { type: 'enabled', budget_tokens: N }` -/// and `inferenceConfig.maxTokens` is bumped by `budgetTokens`. The -/// `reasoningConfig` key never appears in the request body. User-supplied -/// `additionalModelRequestFields` are merged with the derived `thinking` field. +// The request capabilities come from Anthropic's getModelCapabilities, not +// Bedrock's separate strict-tool support list. Unknown newer model IDs assume +// adaptive thinking and reject sampling; legacy IDs retain their defaults. +fn reasoning_capabilities(model_id: &str) -> (f64, bool, bool, bool) { + let Some((_, model)) = model_id.split_once("claude-") else { + return (4096.0, false, false, false); + }; + if ["opus-5", "fable-5", "sonnet-5", "opus-4-8", "opus-4-7"] + .iter() + .any(|id| model.starts_with(id)) + { + (128000.0, true, true, true) + } else if ["sonnet-4-6", "opus-4-6"] + .iter() + .any(|id| model.starts_with(id)) + { + (128000.0, true, false, true) + } else if ["sonnet-4-5", "opus-4-5", "haiku-4-5"] + .iter() + .any(|id| model.starts_with(id)) + { + (64000.0, false, false, true) + } else if model.starts_with("opus-4-1") { + (32000.0, false, false, true) + } else if model.starts_with("sonnet-4-") || model.starts_with("sonnet-4@") { + (64000.0, false, false, false) + } else if model.starts_with("opus-4-") || model.starts_with("opus-4@") { + (32000.0, false, false, false) + } else { + let two = model.strip_prefix('v').unwrap_or(model); + let legacy = model == "instant" + || model.starts_with("instant-") + || two == "2" + || two.starts_with("2-") + || two.starts_with("2.") + || two.starts_with("2:") + || model == "3" + || model.starts_with("3-") + || model.starts_with("3."); + if legacy { + (4096.0, false, false, false) + } else { + (128000.0, true, true, true) + } + } +} + +/// Build a checked Converse request, preserving upstream warnings. /// /// # Errors /// -/// Returns an unsupported-functionality error for unsupported tool result content. -pub fn build_request_body( +/// A prompt or tool definition the Converse API cannot express. +pub fn build_request_body_checked( model_id: &str, - options: &CallOptions, -) -> Result { - let (system, messages) = convert_prompt_to_bedrock(&options.prompt)?; - - // Extract Bedrock-specific provider options. - let bedrock_opts = options::read(options.provider_options.as_ref()); - let reasoning_config = bedrock_opts.and_then(|bo| bo.get("reasoningConfig")); - let user_amrf = bedrock_opts.and_then(|bo| bo.get("additionalModelRequestFields")); - - // Derive the thinking config and budget_tokens from reasoningConfig. - let mut thinking: Option = None; - let mut budget_tokens: Option = None; - if let Some(rc) = reasoning_config - && rc.get("type").and_then(|v| v.as_str()) == Some("enabled") - && let Some(bt) = rc.get("budgetTokens").and_then(serde_json::Value::as_u64) + call: &CallOptions, + model_family: Option<&str>, +) -> Result<(Value, Vec, bool, bool), aimux_core::AiMuxError> { + use aimux_core::types::{ReasoningEffort, Warning}; + let mut warnings = Vec::new(); + let mut compatibility = Vec::new(); + let mut warn = |feature: &str, details: Option<&str>| { + warnings.push(Warning::Unsupported { + feature: feature.into(), + details: details.map(str::to_owned), + }) + }; + let bedrock = options::read(call.provider_options.as_ref()); + validate_options(bedrock)?; + let mut reasoning = bedrock + .and_then(|b| b.get("reasoningConfig")) + .cloned() + .unwrap_or(json!({})); + let anthropic = model_family == Some("anthropic") + || model_id.contains("anthropic") + || (model_id.contains(":application-inference-profile/") + && reasoning.get("budgetTokens").is_some()); + let openai = model_id + .split('.') + .any(|p| p == options::OPENAI_MODEL_FAMILY); + let oss = openai && model_id.contains("-oss-"); + let nova = model_id.contains("amazon.nova-2-lite-v1:0"); + let (max_reasoning_tokens, adaptive, rejects_sampling, supports_structured_output) = + reasoning_capabilities(model_id); + if let Some(effort) = &call.reasoning + && effort.is_custom() { - thinking = Some(json!({ "type": "enabled", "budget_tokens": bt })); - budget_tokens = Some(bt); + if anthropic && matches!(effort, ReasoningEffort::None) { + reasoning = json!({"type":"disabled"}); + } else if !matches!(effort, ReasoningEffort::None) + && (anthropic + || openai + || nova + || bedrock.is_some_and(|b| b.contains_key("reasoningConfig"))) + { + let level = match effort { + ReasoningEffort::Minimal | ReasoningEffort::Low => "low", + ReasoningEffort::Medium => "medium", + ReasoningEffort::High => "high", + ReasoningEffort::Xhigh => "max", + _ => "low", + }; + if (!anthropic || adaptive) + && matches!(effort, ReasoningEffort::Minimal | ReasoningEffort::Xhigh) + { + compatibility.push(Warning::Compatibility { feature: "reasoning".into(), details: Some(format!("reasoning \"{effort}\" is not directly supported by this model. mapped to effort \"{level}\".")) }); + } + let mut derived = if anthropic && adaptive { + json!({"type":"adaptive","maxReasoningEffort":level}) + } else if anthropic { + let percentage = match effort { + ReasoningEffort::Minimal => 0.02, + ReasoningEffort::Low => 0.1, + ReasoningEffort::Medium => 0.3, + ReasoningEffort::High => 0.6, + _ => 0.9, + }; + json!({"type":"enabled","budgetTokens":(max_reasoning_tokens * percentage).round().clamp(1024.0, max_reasoning_tokens) as u64}) + } else { + json!({"maxReasoningEffort":level}) + }; + if nova { + derived["type"] = json!("enabled"); + } + if let Some(explicit) = reasoning.as_object() { + derived.as_object_mut().unwrap().extend(explicit.clone()); + } + reasoning = derived; + } else if !matches!(effort, ReasoningEffort::None) { + warn( + "reasoning", + Some( + "Portable reasoning is not supported for this model and will be ignored. If the model supports a provider-specific reasoning configuration, use providerOptions.amazonBedrock.reasoningConfig.", + ), + ); + } + if reasoning["type"] == "disabled" { + reasoning.as_object_mut().unwrap().remove("budgetTokens"); + reasoning + .as_object_mut() + .unwrap() + .remove("maxReasoningEffort"); + } } - - let mut inference_config = serde_json::Map::new(); - if let Some(max) = options.max_output_tokens { - // When thinking is enabled, maxTokens is bumped by budgetTokens so the - // model has room for both reasoning and the visible response. - let max_tokens = budget_tokens.map(|bt| max + bt as u32).unwrap_or(max); - inference_config.insert("maxTokens".to_string(), json!(max_tokens)); + let thinking = anthropic && matches!(reasoning["type"].as_str(), Some("enabled" | "adaptive")); + let rejects_forced = anthropic + && (model_id.contains("sonnet-5-5") + || model_id.contains("opus-5-5") + || model_id.contains("fable-5-1")); + let mut inference = serde_json::Map::new(); + if let Some(max) = call.max_output_tokens { + inference.insert("maxTokens".into(), json!(max)); } - if let Some(temp) = options.temperature { - inference_config.insert("temperature".to_string(), json!(temp)); + if let Some(temp) = call.temperature { + if rejects_sampling || thinking || (openai && !oss) { + let details = if rejects_sampling { + format!("temperature is not supported by {model_id} and will be ignored") + } else if thinking { + "temperature is not supported when thinking is enabled".into() + } else { + "temperature is not supported by this OpenAI model on the Converse API".into() + }; + warn("temperature", Some(&details)); + } else { + let temp = if !openai || oss { + if temp > 1.0 { + warn( + "temperature", + Some(&format!( + "{temp} exceeds bedrock maximum of 1.0. clamped to 1.0" + )), + ); + 1.0 + } else if temp < 0.0 { + warn( + "temperature", + Some(&format!( + "{temp} is below bedrock minimum of 0. clamped to 0" + )), + ); + 0.0 + } else { + temp + } + } else { + temp + }; + inference.insert("temperature".into(), json!(temp)); + } } - if let Some(top_p) = options.top_p { - inference_config.insert("topP".to_string(), json!(top_p)); + if let Some(p) = call.top_p { + if rejects_sampling || thinking || (openai && !oss) { + let details = if rejects_sampling { + format!("topP is not supported by {model_id} and will be ignored") + } else if thinking { + "topP is not supported when thinking is enabled".into() + } else { + "topP is not supported by this OpenAI model on the Converse API".into() + }; + warn("topP", Some(&details)); + } else { + inference.insert("topP".into(), json!(p)); + } } - if let Some(top_k) = options.top_k { - inference_config.insert("topK".to_string(), json!(top_k)); + if let Some(k) = call.top_k { + if rejects_sampling || thinking { + let details = if rejects_sampling { + format!("topK is not supported by {model_id} and will be ignored") + } else { + "topK is not supported when thinking is enabled".into() + }; + warn("topK", Some(&details)); + } else { + inference.insert("topK".into(), json!(k)); + } } - if let Some(ref stop) = options.stop_sequences { - inference_config.insert("stopSequences".to_string(), json!(stop)); + if let Some(stop) = &call.stop_sequences { + if openai { + warn( + "stopSequences", + Some("stopSequences is not supported by this OpenAI model on the Converse API"), + ); + } else { + inference.insert("stopSequences".into(), json!(stop)); + } } - - // Build additionalModelRequestFields: merge user-supplied fields with the - // derived thinking config (thinking takes precedence on conflict). - let mut amrf = serde_json::Map::new(); - if let Some(user_amrf) = user_amrf - && let Some(obj) = user_amrf.as_object() - { - for (k, v) in obj { - amrf.insert(k.clone(), v.clone()); + for (feature, present) in [ + ("frequencyPenalty", call.frequency_penalty.is_some()), + ("presencePenalty", call.presence_penalty.is_some()), + ("seed", call.seed.is_some()), + ] { + if present { + warn(feature, None); } } - if let Some(thinking) = thinking { - amrf.insert("thinking".to_string(), thinking); + let mut fields = bedrock + .and_then(|b| b.get("additionalModelRequestFields")) + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + if thinking && reasoning["type"] == "enabled" { + if let Some(budget) = reasoning["budgetTokens"].as_u64() { + let max = inference + .get("maxTokens") + .and_then(Value::as_u64) + .unwrap_or(4096); + inference.insert("maxTokens".into(), json!(max.saturating_add(budget))); + fields.insert( + "thinking".into(), + json!({"type":"enabled","budget_tokens":budget}), + ); + } + } else if thinking && reasoning["type"] == "adaptive" { + let mut config = json!({"type":"adaptive"}); + if let Some(display) = reasoning.get("display") { + config["display"] = display.clone(); + } + fields.insert("thinking".into(), config); + } else if !anthropic { + if reasoning.get("budgetTokens").is_some() { + warn( + "budgetTokens", + Some( + "budgetTokens applies only to Anthropic models on Bedrock and will be ignored for this model.", + ), + ); + } + if reasoning["type"] == "adaptive" { + warn( + "adaptive thinking", + Some("adaptive thinking type applies only to Anthropic models on Bedrock."), + ); + } } - - let mut body = serde_json::Map::new(); - body.insert("messages".to_string(), Value::Array(messages)); - - if !system.is_empty() { - body.insert("system".to_string(), Value::Array(system)); + if let Some(effort) = reasoning.get("maxReasoningEffort") { + let key = if anthropic { + "output_config" + } else if openai && !oss { + "reasoning" + } else { + "reasoningConfig" + }; + if oss { + fields.insert("reasoning_effort".into(), effort.clone()); + } else { + let mut config = fields + .get(key) + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + config.insert( + if anthropic || openai { + "effort" + } else { + "maxReasoningEffort" + } + .into(), + effort.clone(), + ); + if !anthropic && !openai { + if let Some(kind) = reasoning.get("type").filter(|t| *t != "adaptive") { + config.insert("type".into(), kind.clone()); + } + if let Some(budget) = reasoning + .get("budgetTokens") + .filter(|_| reasoning["type"] == "enabled") + { + config.insert("budgetTokens".into(), budget.clone()); + } + } + fields.insert(key.into(), Value::Object(config)); + } + if nova + && reasoning["type"] == "enabled" + && matches!(effort.as_str(), Some("high" | "xhigh" | "max")) + && inference.remove("maxTokens").is_some() + { + warn( + "maxOutputTokens", + Some( + "maxOutputTokens is not supported when high reasoning is enabled and will be ignored", + ), + ); + } } - if !inference_config.is_empty() { - body.insert( - "inferenceConfig".to_string(), - Value::Object(inference_config), - ); + if let Some(betas) = bedrock.and_then(|b| b.get("anthropicBeta")) { + fields.insert("anthropic_beta".into(), betas.clone()); } - if !amrf.is_empty() { - body.insert( - "additionalModelRequestFields".to_string(), - Value::Object(amrf), - ); + let anthropic_options = call + .provider_options + .as_ref() + .and_then(|p| p.get("anthropic")); + if let Some(value) = anthropic_options.and_then(|p| p.get("structuredOutputMode")) + && !value + .as_str() + .is_some_and(|v| ["outputFormat", "jsonTool", "auto"].contains(&v)) + { + return Err(aimux_core::AiMuxError::InvalidArgument( + "Invalid Anthropic structuredOutputMode".into(), + )); } - - // Tools - let function_tools: Option> = options.tools.as_ref().map(|tools| { - tools + if anthropic_options + .and_then(|p| p.get("disableParallelToolUse")) + .is_some_and(|v| !v.is_boolean()) + { + return Err(aimux_core::AiMuxError::InvalidArgument( + "Invalid Anthropic disableParallelToolUse".into(), + )); + } + let structured_mode = bedrock + .and_then(|b| b.get("structuredOutputMode")) + .or_else(|| anthropic_options.and_then(|p| p.get("structuredOutputMode"))) + .and_then(Value::as_str) + .unwrap_or("auto"); + if structured_mode == "jsonTool" + && let Some(output_config) = fields + .get_mut("output_config") + .and_then(Value::as_object_mut) + { + output_config.remove("format"); + if output_config.is_empty() { + fields.remove("output_config"); + } + } + let schema = match &call.response_format { + Some(ResponseFormat::Json { schema, .. }) => schema.as_ref(), + _ => None, + }; + let supports_native = supports_native_structured_output(model_id); + let uses_native = anthropic + && schema.is_some() + && (structured_mode == "outputFormat" + || (structured_mode == "auto" + && supports_native + && (supports_structured_output || thinking || model_family == Some("anthropic")))); + let uses_json_instruction = !uses_native + && anthropic + && schema.is_some() + && (rejects_forced + || (structured_mode != "jsonTool" + && !supports_strict_tools(model_id) + && call.tools.as_ref().is_some_and(|tools| !tools.is_empty()))); + let uses_json_tool = schema.is_some() && !uses_native && !uses_json_instruction; + if uses_native { + let output_config = fields.entry("output_config").or_insert_with(|| json!({})); + if !output_config.is_object() { + *output_config = json!({}); + } + output_config["format"] = json!({"type":"json_schema", "schema":crate::anthropic::sanitize_json_schema::sanitize_json_schema(schema.unwrap())}); + } + let mut all_tools = call.tools.clone().unwrap_or_default(); + if uses_json_tool { + all_tools.push(Tool::Function( + FunctionTool::new("json", schema.unwrap().clone()) + .with_description("Respond with a JSON object."), + )); + } + let tool_choice = if uses_json_tool { + Some(ToolChoice::Required) + } else { + call.tool_choice.clone() + }; + let disable_parallel = anthropic_options + .and_then(|p| p.get("disableParallelToolUse")) + .and_then(Value::as_bool) + == Some(true); + let mut provider_tools = super::tools::prepare_provider_tools( + Some(&all_tools), + tool_choice.as_ref(), + anthropic, + disable_parallel, + rejects_forced, + )?; + for warning in provider_tools.warnings.drain(..) { + match warning { + Warning::Unsupported { feature, details } => warn(&feature, details.as_deref()), + other => compatibility.push(other), + } + } + let tools: Option> = Some( + all_tools .iter() - .filter_map(|t| match t { - Tool::Function(ft) => Some(ft.clone()), - Tool::Provider(_) => None, + .filter_map(|tool| match tool { + Tool::Function(f) => Some(f.clone()), + _ => None, }) - .collect() - }); - let tool_config = prepare_tools(&function_tools, options.tool_choice.as_ref(), model_id); - if tool_config - .as_object() - .map(|o| !o.is_empty()) - .unwrap_or(false) + .collect(), + ); + let function_choice = + if provider_tools.using_anthropic_tools && matches!(tool_choice, Some(ToolChoice::None)) { + Some(ToolChoice::Auto) + } else { + tool_choice.clone() + }; + let mut tool_config = prepare_tools(&tools, function_choice.as_ref(), model_id); + if !provider_tools.tools.is_empty() { + let mut combined = provider_tools.tools; + if let Some(functions) = tool_config.get("tools").and_then(Value::as_array) { + combined.extend(functions.clone()); + } + tool_config["tools"] = json!(combined); + } + let has_additional_tools = provider_tools.tool_choice.is_some(); + if provider_tools.using_anthropic_tools { + tool_config.as_object_mut().unwrap().remove("toolChoice"); + if let Some(choice) = provider_tools.tool_choice { + fields.insert("tool_choice".into(), choice); + } + } + if !provider_tools.betas.is_empty() { + let mut betas = bedrock + .and_then(|b| b.get("anthropicBeta")) + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + betas.extend(provider_tools.betas.into_iter().map(Value::String)); + fields.insert("anthropic_beta".into(), json!(betas)); + } + if !provider_tools.using_anthropic_tools + && rejects_forced + && matches!(tool_choice, Some(ToolChoice::Required | ToolChoice::Tool { .. })) + && all_tools.iter().any(|tool| !matches!(tool, Tool::Provider(provider) if matches!(provider.id.as_str(), "anthropic.web_search_20250305" | "anthropic.web_search_20260318" | "anthropic.web_fetch_20260318"))) { - body.insert("toolConfig".to_string(), tool_config); + if tool_config.get("tools").is_some() { tool_config["toolChoice"] = json!({"auto":{}}); } + let details = match &tool_choice { + Some(ToolChoice::Tool {tool_name}) => format!("toolChoice 'tool' is not supported by this model because it rejects forced tool use. Only the '{tool_name}' tool is sent with 'auto' tool choice. Instruct the model to use the tool in the prompt and verify that a tool call was made."), + _ => "toolChoice 'required' is not supported by this model because it rejects forced tool use. Using 'auto' instead. Instruct the model to use a tool in the prompt and verify that a tool call was made.".into(), + }; + warn("toolChoice", Some(&details)); } + if let Some(tools) = &tools { + for tool in tools.iter().filter(|tool| !matches!(&function_choice, Some(ToolChoice::Tool { tool_name }) if tool_name != &tool.name)) { + if let Some(strict) = tool.strict + && !supports_strict_tools(model_id) + { + warn( + "strict", + Some(&format!( + "Tool '{}' has strict: {}, but strict mode is not supported by this model on Amazon Bedrock. The strict property will be ignored.", + tool.name, strict + )), + ); + } else if tool.strict == Some(true) && !is_strict_schema_compatible(&tool.input_schema) + { + warn( + "strict", + Some(&format!( + "Tool '{}' has strict: true, but Amazon Bedrock requires every object in a strict tool schema to set additionalProperties: false. The strict property will be ignored.", + tool.name + )), + ); + } + } + } + if !provider_tools.using_anthropic_tools + && anthropic + && disable_parallel + && tool_config.get("tools").is_some() + { + let choice = if rejects_forced { + json!({"type":"auto","disable_parallel_tool_use":true}) + } else { + match &tool_choice { + Some(ToolChoice::Required) => { + json!({"type":"any","disable_parallel_tool_use":true}) + } + Some(ToolChoice::Tool { tool_name }) => { + json!({"type":"tool","name":tool_name,"disable_parallel_tool_use":true}) + } + _ => json!({"type":"auto","disable_parallel_tool_use":true}), + } + }; + fields.insert("tool_choice".into(), choice); + tool_config.as_object_mut().unwrap().remove("toolChoice"); + } + let mut prompt = call.prompt.clone(); + if tool_config.get("tools").is_none() && !has_additional_tools { + let has_tools = prompt.iter().any(|message| match message { + LanguageModelMessage::Assistant { content, .. } => content.iter().any(|part| { + matches!( + part, + AssistantPart::ToolCall(_) | AssistantPart::ToolResult(_) + ) + }), + LanguageModelMessage::Tool { content, .. } => !content.is_empty(), + _ => false, + }); + if has_tools { + for message in &mut prompt { + match message { + LanguageModelMessage::Assistant { content, .. } => content.retain(|part| { + !matches!( + part, + AssistantPart::ToolCall(_) | AssistantPart::ToolResult(_) + ) + }), + LanguageModelMessage::Tool { content, .. } => content.clear(), + _ => {} + } + } + prompt.retain(|message| match message { + LanguageModelMessage::System { .. } => true, + LanguageModelMessage::User { content, .. } => !content.is_empty(), + LanguageModelMessage::Assistant { content, .. } => !content.is_empty(), + LanguageModelMessage::Tool { content, .. } => !content.is_empty(), + }); + warn( + "toolContent", + Some( + "Tool calls and results removed from conversation because Bedrock does not support tool content without active tools.", + ), + ); + } + } + if uses_json_instruction { + let instruction = format!( + "JSON schema:\n{}\nYou MUST answer with only a JSON object that matches the JSON schema above. Do not wrap it in markdown fences or include any other text.", + schema.unwrap() + ); + if let Some(LanguageModelMessage::System { content, .. }) = prompt.first_mut() { + if !content.is_empty() { + content.push_str("\n\n"); + } + content.push_str(&instruction); + } else { + prompt.insert( + 0, + LanguageModelMessage::System { + content: instruction, + provider_options: None, + }, + ); + } + } + let (system, messages) = convert_prompt(&prompt, model_id.contains("mistral."))?; + let mut body = json!({"messages":messages}); + if !system.is_empty() { + body["system"] = json!(system); + } + if !inference.is_empty() { + body["inferenceConfig"] = Value::Object(inference); + } + if !fields.is_empty() { + body["additionalModelRequestFields"] = Value::Object(fields); + } + if anthropic { + body["additionalModelResponseFieldPaths"] = json!(["/delta/stop_sequence"]); + } + if let Some(tier) = bedrock.and_then(|b| b.get("serviceTier")) { + body["serviceTier"] = json!({"type":tier}); + } + if let Some(bedrock) = bedrock { + for (key, value) in bedrock { + if !matches!( + key.as_str(), + "reasoningConfig" + | "additionalModelRequestFields" + | "serviceTier" + | "structuredOutputMode" + ) { + body[key] = value.clone(); + } + } + } + if tool_config.get("tools").is_some() { + body["toolConfig"] = tool_config; + } + warnings.extend(compatibility); + Ok((body, warnings, uses_json_instruction, uses_json_tool)) +} - Ok(Value::Object(body)) +fn validate_options( + options: Option<&aimux_core::shared::JsonObject>, +) -> Result<(), aimux_core::AiMuxError> { + let Some(options) = options else { + return Ok(()); + }; + for (key, choices) in [ + ( + "structuredOutputMode", + &["outputFormat", "jsonTool", "auto"][..], + ), + ( + "serviceTier", + &["reserved", "priority", "default", "flex"][..], + ), + ] { + if let Some(value) = options.get(key) + && !value.as_str().is_some_and(|v| choices.contains(&v)) + { + return Err(aimux_core::AiMuxError::InvalidArgument(format!( + "Invalid Bedrock {key}" + ))); + } + } + if let Some(value) = options.get("additionalModelRequestFields") + && !value.is_object() + { + return Err(aimux_core::AiMuxError::InvalidArgument( + "Invalid Bedrock additionalModelRequestFields".into(), + )); + } + if let Some(value) = options.get("anthropicBeta") + && !value + .as_array() + .is_some_and(|a| a.iter().all(Value::is_string)) + { + return Err(aimux_core::AiMuxError::InvalidArgument( + "Invalid Bedrock anthropicBeta".into(), + )); + } + if let Some(reasoning) = options.get("reasoningConfig") { + if !reasoning.is_object() { + return Err(aimux_core::AiMuxError::InvalidArgument( + "Invalid Bedrock reasoningConfig".into(), + )); + } + for (key, choices) in [ + ("type", &["enabled", "disabled", "adaptive"][..]), + ("display", &["omitted", "summarized"][..]), + ( + "maxReasoningEffort", + &["low", "medium", "high", "xhigh", "max"][..], + ), + ] { + if let Some(value) = reasoning.get(key) + && !value.as_str().is_some_and(|v| choices.contains(&v)) + { + return Err(aimux_core::AiMuxError::InvalidArgument(format!( + "Invalid Bedrock reasoningConfig.{key}" + ))); + } + } + if let Some(budget) = reasoning.get("budgetTokens") + && !budget.is_number() + { + return Err(aimux_core::AiMuxError::InvalidArgument( + "Invalid Bedrock reasoningConfig.budgetTokens".into(), + )); + } + } + Ok(()) } /// Map a Bedrock `stopReason` to the unified `FinishReason`. @@ -788,9 +1661,8 @@ pub fn map_finish_reason(reason: &str) -> FinishReason { /// /// Mirrors the TS `convertAmazonBedrockUsage(usage: AmazonBedrockUsage | /// undefined | null)`: a `None`/null/undefined usage yields an all-`None` -/// `Usage` (the TS `undefined` fields). The TS `raw` echo is not modelled on -/// the Rust `Usage` type (see `convert-usage` tests for the skipped `raw` -/// cases); `outputTokens.text` mirrors the TS `outputTokens.text` field. +/// `Usage` (the TS `undefined` fields). The original provider payload is +/// retained in `raw`; `outputTokens.text` mirrors the TS field. #[must_use] pub fn convert_usage(usage: Option<&BedrockUsage>) -> aimux_core::types::Usage { use aimux_core::types::Usage; diff --git a/aimux-providers/src/bedrock/embedding.rs b/aimux-providers/src/bedrock/embedding.rs index 1694e5c70..1f64f6bec 100644 --- a/aimux-providers/src/bedrock/embedding.rs +++ b/aimux-providers/src/bedrock/embedding.rs @@ -20,7 +20,7 @@ use aimux_core::shared::SharedProviderOptions; use aimux_provider_utils::HttpBody; -use super::options; +use super::{AmazonBedrockEmbeddingModelFamily, options}; use crate::shared::EndpointConfig; /// An Amazon Bedrock embedding model (e.g. `"amazon.titan-embed-text-v2:0"`). @@ -30,11 +30,35 @@ use crate::shared::EndpointConfig; pub struct BedrockEmbeddingModel { model_id: String, config: EndpointConfig, + model_family: Option, } impl BedrockEmbeddingModel { pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { - Self { model_id, config } + Self { + model_id, + config, + model_family: None, + } + } + + pub(crate) fn with_model_family( + mut self, + model_family: Option, + ) -> Self { + self.model_family = model_family; + self + } + + fn model_family(&self) -> &str { + match self.model_family { + Some(AmazonBedrockEmbeddingModelFamily::Titan) => "titan", + Some(AmazonBedrockEmbeddingModelFamily::Cohere) => options::COHERE_MODEL_FAMILY, + Some(AmazonBedrockEmbeddingModelFamily::Nova) => "nova", + None if is_nova_embedding_model(&self.model_id) => "nova", + None if is_cohere_embedding_model(&self.model_id) => options::COHERE_MODEL_FAMILY, + None => "titan", + } } } @@ -47,7 +71,7 @@ fn is_cohere_embedding_model(model_id: &str) -> bool { /// Returns `true` if the model ID is an Amazon Nova embedding model. fn is_nova_embedding_model(model_id: &str) -> bool { - model_id.starts_with("amazon.nova-") && model_id.contains("embed") + model_id.contains("amazon.nova-") && model_id.contains("embed") } #[async_trait] @@ -61,7 +85,7 @@ impl EmbeddingModel for BedrockEmbeddingModel { } fn max_embeddings_per_call(&self) -> Option { - if is_cohere_embedding_model(&self.model_id) { + if self.model_family() == options::COHERE_MODEL_FAMILY { Some(96) } else { Some(1) @@ -76,10 +100,17 @@ impl EmbeddingModel for BedrockEmbeddingModel { &self, options: &EmbeddingCallOptions, ) -> Result { - let bedrock_options = parse_bedrock_provider_options(options.provider_options.as_ref()); + if options.values.len() > self.max_embeddings_per_call().unwrap_or(1) as usize { + return Err(AiMuxError::InvalidArgument(format!( + "Too many embedding values for amazon-bedrock model {}: maximum {} per call", + self.model_id, + self.max_embeddings_per_call().unwrap_or(1) + ))); + } + let bedrock_options = parse_bedrock_provider_options(options.provider_options.as_ref())?; - let is_nova = is_nova_embedding_model(&self.model_id); - let is_cohere = is_cohere_embedding_model(&self.model_id); + let is_nova = self.model_family() == "nova"; + let is_cohere = self.model_family() == options::COHERE_MODEL_FAMILY; // Build request body based on model family. let body = if is_nova { @@ -152,11 +183,10 @@ impl EmbeddingModel for BedrockEmbeddingModel { let exchange = self.config.exchange(options.headers.as_ref()).await?; let body = exchange.transform_body(body); let body_str = serde_json::to_string(&body).unwrap_or_default(); - // Bedrock model IDs contain dots and colons (e.g. - // `amazon.titan-embed-text-v2:0`). The TS reference URL-encodes the - // model ID, but these characters are valid in URL paths and AWS accepts - // them unencoded — matching the LanguageModel implementation. - let url = exchange.url(&format!("/model/{}/invoke", self.model_id)); + let url = exchange.url(&format!( + "/model/{}/invoke", + super::encode_model_id(&self.model_id) + )); let resp = aimux_provider_utils::post_to_api( exchange.request(url, options), @@ -187,8 +217,11 @@ impl EmbeddingModel for BedrockEmbeddingModel { ] } else if let Some(embeddings_arr) = raw_value.get("embeddings").and_then(|e| e.as_array()) { - let first = &embeddings_arr[0]; - if first.get("embeddingType").is_some() { + if embeddings_arr + .first() + .is_some_and(|first| first.get("embeddingType").is_some()) + { + let first = &embeddings_arr[0]; // Nova response: { embeddings: [{ embeddingType, embedding }] } vec![ first @@ -289,20 +322,73 @@ struct BedrockEmbeddingProviderOptions { /// Reads the `amazonBedrock` key only. fn parse_bedrock_provider_options( options: Option<&SharedProviderOptions>, -) -> BedrockEmbeddingProviderOptions { +) -> Result { let provider_opts = options::read(options); + if let Some(opts) = provider_opts { + let valid_dimensions = [ + ("dimensions", &[1024, 512, 256][..]), + ("embeddingDimension", &[256, 384, 1024, 3072][..]), + ("outputDimension", &[256, 512, 1024, 1536][..]), + ] + .iter() + .all(|(key, allowed)| { + opts.get(*key).is_none_or(|v| { + v.as_f64() + .is_some_and(|n| allowed.iter().any(|allowed| n == f64::from(*allowed))) + }) + }); + let valid_enums = [ + ( + "embeddingPurpose", + &[ + "GENERIC_INDEX", + "TEXT_RETRIEVAL", + "IMAGE_RETRIEVAL", + "VIDEO_RETRIEVAL", + "DOCUMENT_RETRIEVAL", + "AUDIO_RETRIEVAL", + "GENERIC_RETRIEVAL", + "CLASSIFICATION", + "CLUSTERING", + ][..], + ), + ( + "inputType", + &[ + "search_document", + "search_query", + "classification", + "clustering", + ][..], + ), + ("truncate", &["NONE", "START", "END"][..]), + ] + .iter() + .all(|(key, allowed)| { + opts.get(*key) + .is_none_or(|v| v.as_str().is_some_and(|s| allowed.contains(&s))) + }); + if !valid_dimensions + || !valid_enums + || opts.get("normalize").is_some_and(|v| !v.is_boolean()) + { + return Err(AiMuxError::InvalidArgument( + "invalid amazonBedrock provider options".into(), + )); + } + } - BedrockEmbeddingProviderOptions { + Ok(BedrockEmbeddingProviderOptions { dimensions: provider_opts .and_then(|o| o.get("dimensions")) - .and_then(serde_json::Value::as_u64) + .and_then(serde_json::Value::as_f64) .map(|d| d as u32), normalize: provider_opts .and_then(|o| o.get("normalize")) .and_then(serde_json::Value::as_bool), embedding_dimension: provider_opts .and_then(|o| o.get("embeddingDimension")) - .and_then(serde_json::Value::as_u64) + .and_then(serde_json::Value::as_f64) .map(|d| d as u32), embedding_purpose: provider_opts .and_then(|o| o.get("embeddingPurpose")) @@ -318,7 +404,7 @@ fn parse_bedrock_provider_options( .map(std::string::ToString::to_string), output_dimension: provider_opts .and_then(|o| o.get("outputDimension")) - .and_then(serde_json::Value::as_u64) + .and_then(serde_json::Value::as_f64) .map(|d| d as u32), - } + }) } diff --git a/aimux-providers/src/bedrock/event_stream.rs b/aimux-providers/src/bedrock/event_stream.rs index b9492d5ae..80b83fac9 100644 --- a/aimux-providers/src/bedrock/event_stream.rs +++ b/aimux-providers/src/bedrock/event_stream.rs @@ -176,7 +176,7 @@ pub fn decode_messages(data: &[u8]) -> Vec { let payload_start = header_end; let payload_end = offset + total_length - 4; // last 4 bytes are CRC - if header_end > data.len() || payload_end > data.len() { + if header_end > payload_end || payload_end > data.len() { break; } @@ -215,7 +215,7 @@ pub fn decode_messages(data: &[u8]) -> Vec { h_offset += val_len; match name.as_str() { ":message-type" => message_type = val, - ":event-type" => event_type = val, + ":event-type" | ":exception-type" => event_type = val, _ => {} } } @@ -259,96 +259,37 @@ pub fn decode_messages(data: &[u8]) -> Vec { messages } -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn round_trip_single_message() { - let payload = r#"{"contentBlockIndex":0,"delta":{"text":"Hello"}}"#; - let encoded = encode_message("event", "contentBlockDelta", payload); - let decoded = decode_messages(&encoded); - assert_eq!(decoded.len(), 1); - assert_eq!(decoded[0].message_type, "event"); - assert_eq!(decoded[0].event_type, "contentBlockDelta"); - assert_eq!(decoded[0].data, payload); - } - - #[test] - fn round_trip_multiple_messages() { - let messages = vec![ - ("event", "messageStart", r#"{"role":"assistant""#), - ( - "event", - "contentBlockDelta", - r#"{"contentBlockIndex":0,"delta":{"text":"Hi"}}"#, - ), - ("event", "messageStop", r#"{"stopReason":"end_turn"}"#), - ]; - let encoded = encode_messages(&messages); - let decoded = decode_messages(&encoded); - assert_eq!(decoded.len(), 3); - assert_eq!(decoded[0].event_type, "messageStart"); - assert_eq!(decoded[1].event_type, "contentBlockDelta"); - assert_eq!(decoded[2].event_type, "messageStop"); - } - - #[test] - fn corrupt_prelude_crc_is_rejected() { - let payload = r#"{"delta":{"text":"Hi"}}"#; - let mut encoded = encode_message("event", "contentBlockDelta", payload); - // Flip a byte in the prelude CRC (bytes 8..12). - encoded[8] ^= 0xFF; - let decoded = decode_messages(&encoded); - assert_eq!( - decoded.len(), - 0, - "frame with bad prelude CRC must be dropped" - ); - } - - #[test] - fn corrupt_message_crc_is_rejected() { - let payload = r#"{"delta":{"text":"Hi"}}"#; - let mut encoded = encode_message("event", "contentBlockDelta", payload); - // Flip the last byte (part of the trailing message CRC). - let last = encoded.len() - 1; - encoded[last] ^= 0xFF; - let decoded = decode_messages(&encoded); - assert_eq!( - decoded.len(), - 0, - "frame with bad message CRC must be dropped" - ); - } - - #[test] - fn corrupt_payload_is_rejected_by_crc() { - let payload = r#"{"delta":{"text":"Hi"}}"#; - let mut encoded = encode_message("event", "contentBlockDelta", payload); - // Flip a byte in the payload region (middle of the frame). - let mid = encoded.len() / 2; - encoded[mid] ^= 0xFF; - let decoded = decode_messages(&encoded); - assert_eq!( - decoded.len(), - 0, - "frame with corrupted payload must be caught by message CRC" - ); - } - - #[test] - fn valid_frame_before_corrupt_one_is_kept() { - // Two frames: first valid, second with a flipped payload byte. - let good = encode_message("event", "messageStart", r#"{"role":"assistant"}"#); - let mut bad = encode_message("event", "contentBlockDelta", r#"{"delta":{"text":"x"}}"#); - let bad_mid = bad.len() / 2; - bad[bad_mid] ^= 0xFF; - let mut buf = good.clone(); - buf.extend_from_slice(&bad); - let decoded = decode_messages(&buf); - // First frame decodes; decoding stops (break) at the corrupt second frame. - assert_eq!(decoded.len(), 1); - assert_eq!(decoded[0].event_type, "messageStart"); - } +/// Decode frames incrementally, surfacing transport, framing, and checksum errors. +#[must_use] +pub fn decode_stream( + mut input: futures::stream::BoxStream<'static, Result>, +) -> futures::stream::BoxStream<'static, Result> { + use futures::StreamExt; + Box::pin(async_stream::stream! { + let mut buffer = Vec::new(); + while let Some(chunk) = input.next().await { + match chunk { + Ok(chunk) => buffer.extend_from_slice(&chunk), + Err(error) => { yield Err(error); return; } + } + while buffer.len() >= 4 { + let length = u32::from_be_bytes(buffer[..4].try_into().expect("four bytes")) as usize; + if length < 16 { + yield Err(aimux_core::AiMuxError::InvalidResponseData("Invalid Amazon Bedrock event-stream frame length".into())); + return; + } + if buffer.len() < length { break; } + let decoded = decode_messages(&buffer[..length]); + let Some(message) = decoded.into_iter().next() else { + yield Err(aimux_core::AiMuxError::InvalidResponseData("Invalid Amazon Bedrock event-stream frame".into())); + return; + }; + buffer.drain(..length); + yield Ok(message); + } + } + if !buffer.is_empty() { + yield Err(aimux_core::AiMuxError::InvalidResponseData(format!("Incomplete Amazon Bedrock event-stream frame: {} buffered bytes remain at end of stream.", buffer.len()))); + } + }) } diff --git a/aimux-providers/src/bedrock/image.rs b/aimux-providers/src/bedrock/image.rs index 0e7e32c0a..04f0cbda0 100644 --- a/aimux-providers/src/bedrock/image.rs +++ b/aimux-providers/src/bedrock/image.rs @@ -43,6 +43,18 @@ impl BedrockImageModel { Self { model_id, config } } + /// Whether this model advertises image file inputs. + #[must_use] + pub fn supports_file_inputs(&self) -> Option { + (self.model_id == "amazon.nova-canvas-v1:0").then_some(true) + } + + /// Whether this model advertises mask inputs. + #[must_use] + pub fn supports_mask_inputs(&self) -> Option { + self.supports_file_inputs() + } + fn get_base64_data(file: &ImageFile) -> Result { match file { ImageFile::Url { .. } => Err(AiMuxError::InvalidArgument( @@ -82,6 +94,59 @@ impl ImageModel for BedrockImageModel { // Parse provider options (`amazonBedrock`) let ab_opts = options::read(Some(&options.provider_options)); + if let Some(opts) = ab_opts { + for (key, allowed) in [ + ("quality", &["standard", "premium"][..]), + ( + "style", + &[ + "3D_ANIMATED_FAMILY_FILM", + "DESIGN_SKETCH", + "FLAT_VECTOR_ILLUSTRATION", + "GRAPHIC_NOVEL_ILLUSTRATION", + "MAXIMALISM", + "MIDCENTURY_RETRO", + "PHOTOREALISM", + "SOFT_DIGITAL_PAINTING", + ][..], + ), + ( + "taskType", + &[ + "TEXT_IMAGE", + "IMAGE_VARIATION", + "INPAINTING", + "OUTPAINTING", + "BACKGROUND_REMOVAL", + ][..], + ), + ("outPaintingMode", &["DEFAULT", "PRECISE"][..]), + ] { + if opts + .get(key) + .is_some_and(|v| !v.as_str().is_some_and(|s| allowed.contains(&s))) + { + return Err(AiMuxError::InvalidArgument( + "invalid amazonBedrock provider options".into(), + )); + } + } + for key in ["negativeText", "maskPrompt"] { + if opts.get(key).is_some_and(|v| !v.is_string()) { + return Err(AiMuxError::InvalidArgument( + "invalid amazonBedrock provider options".into(), + )); + } + } + for key in ["cfgScale", "similarityStrength"] { + if opts.get(key).is_some_and(|v| !v.is_number()) { + return Err(AiMuxError::InvalidArgument( + "invalid amazonBedrock provider options".into(), + )); + } + } + } + // Build image generation config let mut image_gen_config = Map::new(); if let Some(w) = width { @@ -90,16 +155,22 @@ impl ImageModel for BedrockImageModel { if let Some(h) = height { image_gen_config.insert("height".into(), json!(h)); } - if let Some(seed) = options.seed { + if let Some(seed) = options.seed.filter(|v| *v != 0) { image_gen_config.insert("seed".into(), json!(seed)); } if options.n > 0 { image_gen_config.insert("numberOfImages".into(), json!(options.n)); } - if let Some(v) = ab_opts.and_then(|o| o.get("quality")) { + if let Some(v) = ab_opts + .and_then(|o| o.get("quality")) + .filter(|v| v.as_str().is_some_and(|s| !s.is_empty())) + { image_gen_config.insert("quality".into(), v.clone()); } - if let Some(v) = ab_opts.and_then(|o| o.get("cfgScale")) { + if let Some(v) = ab_opts + .and_then(|o| o.get("cfgScale")) + .filter(|v| v.as_f64().is_some_and(|n| n != 0.0)) + { image_gen_config.insert("cfgScale".into(), v.clone()); } @@ -124,10 +195,13 @@ impl ImageModel for BedrockImageModel { "INPAINTING" => { let mut params = Map::new(); params.insert("image".into(), json!(source_image_b64)); - if let Some(ref p) = options.prompt { + if let Some(p) = options.prompt.as_ref().filter(|p| !p.is_empty()) { params.insert("text".into(), json!(p)); } - if let Some(v) = ab_opts.and_then(|o| o.get("negativeText")) { + if let Some(v) = ab_opts + .and_then(|o| o.get("negativeText")) + .filter(|v| v.as_str().is_some_and(|s| !s.is_empty())) + { params.insert("negativeText".into(), v.clone()); } if has_mask { @@ -147,13 +221,19 @@ impl ImageModel for BedrockImageModel { "OUTPAINTING" => { let mut params = Map::new(); params.insert("image".into(), json!(source_image_b64)); - if let Some(ref p) = options.prompt { + if let Some(p) = options.prompt.as_ref().filter(|p| !p.is_empty()) { params.insert("text".into(), json!(p)); } - if let Some(v) = ab_opts.and_then(|o| o.get("negativeText")) { + if let Some(v) = ab_opts + .and_then(|o| o.get("negativeText")) + .filter(|v| v.as_str().is_some_and(|s| !s.is_empty())) + { params.insert("negativeText".into(), v.clone()); } - if let Some(v) = ab_opts.and_then(|o| o.get("outPaintingMode")) { + if let Some(v) = ab_opts + .and_then(|o| o.get("outPaintingMode")) + .filter(|v| v.as_str().is_some_and(|s| !s.is_empty())) + { params.insert("outPaintingMode".into(), v.clone()); } if has_mask { @@ -184,10 +264,13 @@ impl ImageModel for BedrockImageModel { .collect::>()?; let mut params = Map::new(); params.insert("images".into(), json!(images)); - if let Some(ref p) = options.prompt { + if let Some(p) = options.prompt.as_ref().filter(|p| !p.is_empty()) { params.insert("text".into(), json!(p)); } - if let Some(v) = ab_opts.and_then(|o| o.get("negativeText")) { + if let Some(v) = ab_opts + .and_then(|o| o.get("negativeText")) + .filter(|v| v.as_str().is_some_and(|s| !s.is_empty())) + { params.insert("negativeText".into(), v.clone()); } if let Some(v) = ab_opts.and_then(|o| o.get("similarityStrength")) { @@ -212,10 +295,16 @@ impl ImageModel for BedrockImageModel { if let Some(ref p) = options.prompt { params.insert("text".into(), json!(p)); } - if let Some(v) = ab_opts.and_then(|o| o.get("negativeText")) { + if let Some(v) = ab_opts + .and_then(|o| o.get("negativeText")) + .filter(|v| v.as_str().is_some_and(|s| !s.is_empty())) + { params.insert("negativeText".into(), v.clone()); } - if let Some(v) = ab_opts.and_then(|o| o.get("style")) { + if let Some(v) = ab_opts + .and_then(|o| o.get("style")) + .filter(|v| v.as_str().is_some_and(|s| !s.is_empty())) + { params.insert("style".into(), v.clone()); } args.insert("taskType".into(), json!("TEXT_IMAGE")); @@ -235,11 +324,15 @@ impl ImageModel for BedrockImageModel { }); } + let current_date = chrono::Utc::now().to_rfc3339(); let exchange = self.config.exchange(options.headers.as_ref()).await?; let body = exchange.transform_body(Value::Object(args)); let body_str = serde_json::to_string(&body).map_err(|e| AiMuxError::JsonParse(e.to_string()))?; - let url = exchange.url(&format!("/model/{}/invoke", self.model_id)); + let url = exchange.url(&format!( + "/model/{}/invoke", + super::encode_model_id(&self.model_id) + )); let resp = aimux_provider_utils::post_to_api( exchange.request(url.clone(), options), @@ -300,7 +393,7 @@ impl ImageModel for BedrockImageModel { warnings, provider_metadata: None, response: ImageResponse { - timestamp: Some(chrono::Utc::now().to_rfc3339()), + timestamp: Some(current_date), model_id: Some(self.model_id.clone()), headers: Some(rh), }, diff --git a/aimux-providers/src/bedrock/mod.rs b/aimux-providers/src/bedrock/mod.rs index 221ced35a..3e838a57d 100644 --- a/aimux-providers/src/bedrock/mod.rs +++ b/aimux-providers/src/bedrock/mod.rs @@ -34,6 +34,7 @@ pub mod image; mod model; pub(crate) mod options; pub mod reranking; +pub(crate) mod tools; mod types; use std::sync::{Arc, OnceLock}; @@ -76,6 +77,54 @@ const SESSION_TOKEN_ENV_VAR: &str = "AWS_SESSION_TOKEN"; const RUNTIME_ENDPOINT_ENV_VAR: &str = "AWS_ENDPOINT_URL_BEDROCK_RUNTIME"; const AGENT_RUNTIME_ENDPOINT_ENV_VAR: &str = "AWS_ENDPOINT_URL_BEDROCK_AGENT_RUNTIME"; +pub(crate) fn encode_model_id(model_id: &str) -> String { + const SAFE: &percent_encoding::AsciiSet = &percent_encoding::NON_ALPHANUMERIC + .remove(b'-') + .remove(b'_') + .remove(b'.') + .remove(b'!') + .remove(b'~') + .remove(b'*') + .remove(b'\'') + .remove(b'(') + .remove(b')'); + percent_encoding::utf8_percent_encode(model_id, SAFE).to_string() +} + +/// The chat family override supported by the upstream provider. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum AmazonBedrockChatModelFamily { + Anthropic, +} + +impl AmazonBedrockChatModelFamily { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Anthropic => "anthropic", + } + } +} + +/// The embedding families supported by the upstream provider. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum AmazonBedrockEmbeddingModelFamily { + Titan, + Cohere, + Nova, +} + +/// Optional family override for opaque inference-profile IDs. +#[derive(Clone, Debug, Default)] +pub struct AmazonBedrockChatModelSettings { + pub model_family: Option, +} + +/// Optional embedding family override for opaque inference-profile IDs. +#[derive(Clone, Debug, Default)] +pub struct AmazonBedrockEmbeddingModelSettings { + pub model_family: Option, +} + pub(crate) fn bedrock_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -83,11 +132,17 @@ pub(crate) fn bedrock_failed_response_handler() -> aimux_provider_utils::Respons // as `{ "error": { ... } }` by compatible gateways. let error = data.get("error").unwrap_or(data); aimux_provider_utils::ProviderErrorParts { - message: error - .get("message") - .and_then(Value::as_str) - .unwrap_or_default() - .to_owned(), + message: { + let message = error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(); + match error.get("type").and_then(Value::as_str) { + Some(kind) => format!("{kind}: {message}"), + None => message, + } + }, provider_code: error .get("type") .or_else(|| error.get("__type")) @@ -128,8 +183,8 @@ pub struct AmazonBedrockProviderSettings { pub session_token: Option, /// Base URL for the Bedrock Runtime calls (and the Agent Runtime /// `rerank` calls). Default `https://bedrock-runtime.{region}.amazonaws.com`, - /// or the service endpoint environment variable, then `AWS_ENDPOINT_URL`; - /// a trailing slash is removed. + /// or `AWS_ENDPOINT_URL_BEDROCK_RUNTIME` when that is set; a trailing slash + /// is removed. pub base_url: Option, /// Extra headers on every request. A `None` value removes the header. /// Per-call headers win over these. @@ -289,9 +344,21 @@ impl Auth { if let Some(url) = load_optional_setting(None, endpoint_env_var) .or_else(|| load_optional_setting(None, "AWS_ENDPOINT_URL")) { - return Ok(without_trailing_slash(&url)); + return validate_base_url(&url); } - Ok(format!("https://{service}.{}.amazonaws.com", region()?)) + let region = region()?; + let suffix = [ + ("cn-", "amazonaws.com.cn"), + ("us-iso-", "c2s.ic.gov"), + ("us-isob-", "sc2s.sgov.gov"), + ("eu-isoe-", "cloud.adc-e.uk"), + ("us-isof-", "csp.hci.ic.gov"), + ("eusc-", "amazonaws.eu"), + ] + .iter() + .find(|(prefix, _)| region.starts_with(prefix)) + .map_or("amazonaws.com", |(_, suffix)| *suffix); + Ok(format!("https://{service}.{region}.{suffix}")) } } @@ -341,7 +408,13 @@ pub fn create_amazon_bedrock( let header_auth = auth.clone(); let headers: HeadersFn = Resolvable::from_async_fn(move || { let auth = header_auth.clone(); - let user = user_headers.clone().unwrap_or_default(); + let mut user = combine_headers(&[&user_headers.clone().unwrap_or_default()]); + let suffix = concat!("ai-sdk-amazon-bedrock/", env!("CARGO_PKG_VERSION")); + let agent = user + .get("user-agent") + .and_then(|value| value.as_deref()) + .map_or_else(|| suffix.to_string(), |value| format!("{value} {suffix}")); + user.insert("user-agent".into(), Some(agent)); async move { match auth.bearer().await? { Some(key) => { @@ -455,6 +528,19 @@ impl AmazonBedrockProvider { BedrockModel::from_config(model_id.to_string(), self.runtime_config()) } + #[must_use] + pub fn chat_with_settings( + &self, + model_id: &str, + settings: &AmazonBedrockChatModelSettings, + ) -> BedrockModel { + self.chat(model_id).with_model_family( + settings + .model_family + .map(|family| family.as_str().to_owned()), + ) + } + /// An embedding model (e.g. `"amazon.titan-embed-text-v2:0"`); /// `provider()` is `"amazon-bedrock"`. #[must_use] @@ -462,6 +548,25 @@ impl AmazonBedrockProvider { BedrockEmbeddingModel::from_config(model_id.to_string(), self.runtime_config()) } + #[must_use] + pub fn embedding_with_settings( + &self, + model_id: &str, + settings: &AmazonBedrockEmbeddingModelSettings, + ) -> BedrockEmbeddingModel { + self.embedding(model_id) + .with_model_family(settings.model_family) + } + + #[must_use] + pub fn call_with_settings( + &self, + model_id: &str, + settings: &AmazonBedrockChatModelSettings, + ) -> Arc { + Arc::new(self.chat_with_settings(model_id, settings)) + } + /// An image generation model (e.g. `"amazon.titan-image-generator-v1"` or /// `"amazon.nova-canvas-v1:0"`); `provider()` is `"amazon-bedrock"`. #[must_use] @@ -536,12 +641,14 @@ impl ProviderDiscovery for AmazonBedrockProvider { let exchange = config.exchange(None).await?; let mut headers = exchange.headers(); headers.push(("Accept".to_string(), "application/json".to_string())); + let mut request = exchange.with_transport(aimux_provider_utils::HttpRequest { + url, + headers, + ..Default::default() + }); + request.credentialed_origin = Some(request.url.clone()); let resp = aimux_provider_utils::get_from_api( - exchange.with_transport(aimux_provider_utils::HttpRequest { - url, - headers, - ..Default::default() - }), + request, aimux_provider_utils::create_json_response_handler(), bedrock_failed_response_handler(), ) diff --git a/aimux-providers/src/bedrock/model.rs b/aimux-providers/src/bedrock/model.rs index 6d3546ccc..3e40090b5 100644 --- a/aimux-providers/src/bedrock/model.rs +++ b/aimux-providers/src/bedrock/model.rs @@ -6,14 +6,14 @@ //! etc.). use aimux_core::tool::RawToolCall; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use async_trait::async_trait; use bytes::Bytes; use futures::{StreamExt, stream::BoxStream}; use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; use aimux_core::stream_part::StreamPart; @@ -25,7 +25,9 @@ use serde_json::json; use aimux_provider_utils::HttpBody; -use super::convert::{build_request_body, convert_usage, map_finish_reason}; +use super::convert::{ + build_request_body_checked, convert_usage, map_finish_reason, normalize_tool_call_id, +}; use super::options; use super::types::{BedrockContentBlock, BedrockConverseResponse}; use crate::shared::EndpointConfig; @@ -88,11 +90,21 @@ fn bedrock_event_stream_response_handler() pub struct BedrockModel { model_id: String, config: EndpointConfig, + model_family: Option, } impl BedrockModel { pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { - Self { model_id, config } + Self { + model_id, + config, + model_family: None, + } + } + + pub(crate) fn with_model_family(mut self, model_family: Option) -> Self { + self.model_family = model_family; + self } /// `/model/{model-id}/converse` or `/converse-stream`. @@ -102,11 +114,11 @@ impl BedrockModel { } else { "converse" }; - // Bedrock model IDs contain dots and colons (e.g. - // `anthropic.claude-3-5-sonnet-20240620-v1:0`). These characters are - // valid in URL paths and are sent unencoded — matching the AWS CLI - // and SDK behaviour. - format!("/model/{}/{}", self.model_id, suffix) + format!( + "/model/{}/{}", + super::encode_model_id(&self.model_id), + suffix + ) } } @@ -120,8 +132,21 @@ impl LanguageModel for BedrockModel { &self.model_id } + fn supported_urls(&self) -> SupportedUrls { + let s3 = regex::Regex::new(r"^s3://").expect("static pattern"); + SupportedUrls( + [ + ("image/*".into(), vec![s3.clone()]), + ("video/*".into(), vec![s3]), + ] + .into_iter() + .collect(), + ) + } + async fn do_generate(&self, options: &CallOptions) -> Result { - let body = build_request_body(&self.model_id, options)?; + let (body, warnings, uses_json_instruction, uses_json_tool) = + build_request_body_checked(&self.model_id, options, self.model_family.as_deref())?; let exchange = self.config.exchange(options.headers.as_ref()).await?; let body = exchange.transform_body(body); let body_str = serde_json::to_string(&body).unwrap_or_default(); @@ -134,22 +159,44 @@ impl LanguageModel for BedrockModel { ) .await?; - let response_body = resp.raw_value; let response_headers = resp.response_headers; + let raw_data = resp.raw_value; let data: BedrockConverseResponse = resp.value; // Extract content from response.output.message.content let mut content = Vec::new(); + let mut extractor = uses_json_instruction.then(JsonObjectTextExtractor::default); + let mut is_json_response_from_tool = false; if let Some(output) = &data.output && let Some(message) = &output.message { for block in &message.content { - extract_content(block, &mut content); + let before = content.len(); + extract_content(block, &mut content, &self.model_id); + for part in &mut content[before..] { + if let GenerateContent::ToolCall(tool) = part + && uses_json_tool + && tool.tool_name == "json" + { + is_json_response_from_tool = true; + *part = GenerateContent::Text { + text: tool.input.clone(), + provider_metadata: None, + }; + } + } + if let Some(extractor) = &mut extractor { + for part in &mut content[before..] { + if let GenerateContent::Text { text, .. } = part { + *text = extractor.process(text); + } + } + } } } - let finish_reason = data + let mut finish_reason = data .stop_reason .as_deref() .map(map_finish_reason) @@ -158,7 +205,27 @@ impl LanguageModel for BedrockModel { raw: None, }); - let usage = convert_usage(data.usage.as_ref()); + if is_json_response_from_tool && finish_reason.unified == FinishReasonUnified::ToolCalls { + finish_reason.unified = FinishReasonUnified::Stop; + } + let mut provider_metadata = response_provider_metadata(&data); + if is_json_response_from_tool { + let mut payload = provider_metadata + .as_ref() + .and_then(|meta| meta.get(options::AMAZON_BEDROCK)) + .cloned() + .unwrap_or_default(); + payload.insert("isJsonResponseFromTool".into(), json!(true)); + payload + .entry("stopSequence") + .or_insert(serde_json::Value::Null); + provider_metadata = Some(options::metadata(json!(payload))); + } + let mut usage = convert_usage(data.usage.as_ref()); + usage.raw = raw_data + .as_ref() + .and_then(|raw| raw.get("usage")) + .and_then(|value| value.as_object().cloned()); let request_id = response_headers .get("x-amzn-requestid") @@ -169,25 +236,28 @@ impl LanguageModel for BedrockModel { content, finish_reason, usage, - warnings: Vec::new(), - provider_metadata: None, + warnings, + provider_metadata, response: Some(aimux_core::shared::ResponseInfo { id: request_id, timestamp: response_headers.get("date").cloned(), model_id: Some(self.model_id.clone()), headers: Some(response_headers), - body: response_body, + body: None, }), request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), }) } async fn do_stream(&self, options: &CallOptions) -> Result { - let body = build_request_body(&self.model_id, options)?; + let (body, warnings, uses_json_instruction, uses_json_tool) = + build_request_body_checked(&self.model_id, options, self.model_family.as_deref())?; let exchange = self.config.exchange(options.headers.as_ref()).await?; let body = exchange.transform_body(body); let body_str = serde_json::to_string(&body).unwrap_or_default(); let url = exchange.url(&self.path(true)); + let error_url = url.clone(); + let error_body = body.clone(); let resp = aimux_provider_utils::post_to_api( exchange.request(url, options), HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), @@ -198,20 +268,8 @@ impl LanguageModel for BedrockModel { let response_headers = resp.response_headers; - // Bedrock converse-stream returns binary AWS event stream format. - // We read the full body and decode it, then emit stream parts. - // (For true streaming we'd decode incrementally, but the Bedrock event - // stream codec requires buffering whole frames anyway.) - let mut buf: Vec = Vec::new(); - let mut body_stream = resp.value; - while let Some(chunk) = body_stream.next().await { - match chunk { - Ok(bytes) => buf.extend_from_slice(&bytes), - Err(e) => return Err(e), - } - } - let response_bytes = buf; - + let body_stream = resp.value; + let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); let request_id = response_headers .get("x-amzn-requestid") .cloned() @@ -223,7 +281,7 @@ impl LanguageModel for BedrockModel { let model_id = self.model_id.clone(); let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings: vec![] }); + yield Ok(StreamPart::StreamStart { warnings }); yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { id: request_id, @@ -231,14 +289,12 @@ impl LanguageModel for BedrockModel { model_id: Some(model_id.clone()), })); - let messages = super::event_stream::decode_messages(&response_bytes); + let mut messages = super::event_stream::decode_stream(body_stream); - let mut text_id: Option = None; - let mut reasoning_id: Option = None; - // Reasoning signatures arrive in their own reasoningContent delta - // (usually with no text); accumulated here and attached to the - // concluding ReasoningEnd. - let mut reasoning_signature: Option = None; + let mut extractor = uses_json_instruction.then(JsonObjectTextExtractor::default); + let mut is_json_response_from_tool = false; + let mut text_blocks = HashSet::new(); + let mut reasoning_blocks: HashMap> = HashMap::new(); let mut block_counter = 0usize; // Tool call state: block_index → (id, name, accumulated_json) let mut tool_blocks: HashMap = HashMap::new(); @@ -252,16 +308,41 @@ impl LanguageModel for BedrockModel { serde_json::Map::new(); let mut stop_sequence: Option = None; - for msg in &messages { - if msg.message_type != "event" { - continue; - } - + while let Some(message) = messages.next().await { + let msg = match message { + Ok(msg) => msg, + Err(error) => { yield Err(error); return; } + }; let payload: serde_json::Value = match serde_json::from_str(&msg.data) { Ok(v) => v, - Err(_) => continue, + Err(error) => { + final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); + yield Ok(StreamPart::Error { error: AiMuxError::from(error) }); + continue; + } }; - + if include_raw_chunks { + yield Ok(StreamPart::Raw { raw_value: json!({ msg.event_type.clone(): payload.clone() }) }); + } + if msg.message_type != "event" || msg.event_type.ends_with("Exception") { + final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); + let (status_code, is_retryable) = match msg.event_type.as_str() { + "internalServerException" | "InternalServerException" => (Some(500), true), + "modelStreamErrorException" | "ModelStreamErrorException" => (Some(424), true), + "serviceUnavailableException" | "ServiceUnavailableException" => (Some(503), true), + "throttlingException" | "ThrottlingException" => (Some(429), true), + "validationException" | "ValidationException" => (Some(400), false), + _ => (None, false), + }; + let message = payload.get("message").and_then(|v| v.as_str()).map(str::to_string) + .unwrap_or_else(|| format!("Amazon Bedrock stream failed with {}", msg.event_type)); + yield Ok(StreamPart::Error { error: AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + status_code, is_retryable, provider_code: Some(msg.event_type.clone()), + data: Some(json!({msg.event_type.clone(): payload.clone()})), + ..aimux_core::ApiCallError::new(message, error_url.clone(), error_body.clone()) + })) }); + continue; + } match msg.event_type.as_str() { "messageStart" => { // Nothing to emit; response metadata already sent. @@ -285,11 +366,12 @@ impl LanguageModel for BedrockModel { .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); - let id = if id.is_empty() { - format!("call-{idx}") + let id = normalize_tool_call_id(&if id.is_empty() { + aimux_provider_utils::generate_id() } else { id - }; + }, model_id.contains("mistral.")); + if !(uses_json_tool && name == "json") { yield Ok(StreamPart::ToolInputStart { id: id.clone(), tool_name: name.clone(), @@ -298,18 +380,19 @@ impl LanguageModel for BedrockModel { title: None, provider_metadata: None, }); + } tool_blocks.insert(idx, (id, name, String::new())); } else { // Text block. block_counter = idx + 1; let id = idx.to_string(); - text_id = Some(id.clone()); + text_blocks.insert(idx); yield Ok(StreamPart::TextStart { id, provider_metadata: None}); } } else { // Default: text block. let id = idx.to_string(); - text_id = Some(id.clone()); + text_blocks.insert(idx); yield Ok(StreamPart::TextStart { id, provider_metadata: None}); } } @@ -323,65 +406,47 @@ impl LanguageModel for BedrockModel { // Text delta if let Some(text) = delta.get("text").and_then(|v| v.as_str()) && !text.is_empty() { - if text_id.is_none() { + if text_blocks.insert(idx) { let id = idx.to_string(); - text_id = Some(id.clone()); + text_blocks.insert(idx); yield Ok(StreamPart::TextStart { id, provider_metadata: None}); } - if let Some(id) = &text_id { - yield Ok(StreamPart::TextDelta { - id: id.clone(), - delta: text.to_string(), - provider_metadata: None, - }); + let delta = extractor.as_mut().map(|extractor| extractor.process(text)).unwrap_or_else(|| text.to_string()); + if !delta.is_empty() { + yield Ok(StreamPart::TextDelta { id: idx.to_string(), delta, provider_metadata: None }); } } // Tool use input delta if let Some(partial) = delta.get("toolUse").and_then(|t| t.get("input")) && let Some(partial_str) = partial.as_str() - && let Some((id, _name, acc)) = tool_blocks.get_mut(&idx) + && let Some((id, name, acc)) = tool_blocks.get_mut(&idx) && !partial_str.is_empty() { acc.push_str(partial_str); + if !(uses_json_tool && name == "json") { let id = id.clone(); yield Ok(StreamPart::ToolInputDelta { id, delta: partial_str.to_string(), provider_metadata: None, }); + } } - // Reasoning delta — `reasoningContent.text` carries - // incremental reasoning text; `reasoningContent.signature` - // carries the reasoning signature (typically in a - // final, text-less delta). The signature is attached - // to the concluding ReasoningEnd via provider_metadata - // so extended-thinking multi-turn can echo it back — - // same shape as the non-streaming path. Empty text - // deltas are skipped. if let Some(rc) = delta.get("reasoningContent") { - if let Some(text) = rc.get("text").and_then(|v| v.as_str()) - && !text.is_empty() - { - let id = idx.to_string(); - if reasoning_id.as_deref() != Some(id.as_str()) { - reasoning_id = Some(id.clone()); - yield Ok(StreamPart::ReasoningStart { - id, - provider_metadata: None, - }); - } - yield Ok(StreamPart::ReasoningDelta { - id: idx.to_string(), - delta: text.to_string(), - provider_metadata: None, - }); + if !["text", "signature", "data", "redactedContent"].iter().any(|key| rc.get(key).and_then(|v| v.as_str()).is_some_and(|s| !s.is_empty())) { continue; } + let id = idx.to_string(); + if let std::collections::hash_map::Entry::Vacant(entry) = reasoning_blocks.entry(idx) { + entry.insert(None); + yield Ok(StreamPart::ReasoningStart { id: id.clone(), provider_metadata: None }); } - if let Some(sig) = - rc.get("signature").and_then(|v| v.as_str()) - { - reasoning_signature - .get_or_insert_with(String::new) - .push_str(sig); + if let Some(text) = rc.get("text").and_then(|v| v.as_str()) { + yield Ok(StreamPart::ReasoningDelta { id: id.clone(), delta: text.to_string(), provider_metadata: None }); + } else if let Some(sig) = rc.get("signature").and_then(|v| v.as_str()).filter(|s| !s.is_empty()) { + yield Ok(StreamPart::ReasoningDelta { id: id.clone(), delta: String::new(), provider_metadata: Some(options::metadata(json!({"signature": sig}))) }); + } else if let Some(data) = rc.get("data").and_then(|v| v.as_str()).filter(|s| !s.is_empty()) { + yield Ok(StreamPart::ReasoningDelta { id: id.clone(), delta: String::new(), provider_metadata: Some(options::metadata(json!({"redactedData": data}))) }); + } else if let Some(data) = rc.get("redactedContent").and_then(|v| v.as_str()) { + reasoning_blocks.entry(idx).or_default().get_or_insert_with(String::new).push_str(data); } } } @@ -393,6 +458,13 @@ impl LanguageModel for BedrockModel { .unwrap_or(0) as usize; if let Some((id, name, acc)) = tool_blocks.remove(&idx) { + if uses_json_tool && name == "json" { + is_json_response_from_tool = true; + let id = idx.to_string(); + yield Ok(StreamPart::TextStart { id: id.clone(), provider_metadata: None }); + yield Ok(StreamPart::TextDelta { id: id.clone(), delta: acc, provider_metadata: None }); + yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); + } else { yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: None}); // Empty input normalizes to "{}" per the upstream // provider. @@ -402,49 +474,27 @@ impl LanguageModel for BedrockModel { acc }; yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, + tool_call_id: normalize_tool_call_id(&id, model_id.contains("mistral.")), tool_name: name, input, provider_executed: None, dynamic: None, provider_metadata: None, })); - } else if reasoning_id.is_some() { - let id = idx.to_string(); - if reasoning_id.as_deref() == Some(id.as_str()) { - let provider_metadata = - reasoning_signature_meta(reasoning_signature.take()); - yield Ok(StreamPart::ReasoningEnd { - id, - provider_metadata, - }); - reasoning_id = None; - } - } else if text_id.is_some() { - // Only end if this is the current text block. - let id = idx.to_string(); - if text_id.as_deref() == Some(id.as_str()) { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); - text_id = None; } + } else if let Some(redacted) = reasoning_blocks.remove(&idx) { + yield Ok(StreamPart::ReasoningEnd { id: idx.to_string(), provider_metadata: reasoning_redacted_meta(redacted) }); + } else if text_blocks.remove(&idx) { + yield Ok(StreamPart::TextEnd { id: idx.to_string(), provider_metadata: None }); } } "messageStop" => { if let Some(reason) = payload.get("stopReason").and_then(|v| v.as_str()) { - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); - } - if let Some(id) = reasoning_id.take() { - let provider_metadata = - reasoning_signature_meta(reasoning_signature.take()); - yield Ok(StreamPart::ReasoningEnd { - id, - provider_metadata, - }); - } - final_finish_reason = Some(map_finish_reason(reason)); + let mut finish_reason = map_finish_reason(reason); + if is_json_response_from_tool && finish_reason.unified == FinishReasonUnified::ToolCalls { finish_reason.unified = FinishReasonUnified::Stop; } + final_finish_reason = Some(finish_reason); } // #26: surface which stop sequence sentinel fired // (additionalModelResponseFields.delta.stop_sequence), if any. @@ -460,8 +510,17 @@ impl LanguageModel for BedrockModel { "metadata" => { if let Some(usage) = payload.get("usage") { let bedrock_usage: super::types::BedrockUsage = - serde_json::from_value(usage.clone()).unwrap_or_default(); + match serde_json::from_value(usage.clone()) { + Ok(usage) => usage, + Err(error) => { + final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); + yield Ok(StreamPart::Error { error: AiMuxError::InvalidResponseData(format!("{error}; usage: {usage}")) }); + continue; + } + }; final_usage = convert_usage(Some(&bedrock_usage)); + final_usage.raw = usage.as_object().cloned(); + if let Some(cache_usage) = cache_usage_metadata(Some(&bedrock_usage)) { finish_meta.insert("usage".into(), cache_usage); } } // #27: surface guardrails trace, performanceConfig, and // serviceTier into the Finish chunk's provider_metadata @@ -481,23 +540,13 @@ impl LanguageModel for BedrockModel { } } - // Close any remaining text block. - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); - } - // Close any remaining reasoning block (truncated stream without - // contentBlockStop): still attach the accumulated signature. - if let Some(id) = reasoning_id.take() { - let provider_metadata = - reasoning_signature_meta(reasoning_signature.take()); - yield Ok(StreamPart::ReasoningEnd { id, - provider_metadata, - }); - } - // #26: merge the stop sentinel into the metadata payload, then build // the Finish provider_metadata under `amazonBedrock` (mirrors the // TS `doStream` flush handler). + if is_json_response_from_tool { + finish_meta.insert("isJsonResponseFromTool".into(), json!(true)); + finish_meta.insert("stopSequence".into(), serde_json::Value::Null); + } if let Some(seq) = stop_sequence { finish_meta.insert( "stopSequence".to_string(), @@ -513,7 +562,7 @@ impl LanguageModel for BedrockModel { yield Ok(StreamPart::Finish { finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Stop, + unified: FinishReasonUnified::Other, raw: None, }), usage: final_usage, @@ -531,10 +580,9 @@ impl LanguageModel for BedrockModel { } } -/// Wrap the accumulated reasoning signature as provider_metadata in the same -/// shape the non-streaming path emits (under `amazonBedrock`). -fn reasoning_signature_meta(sig: Option) -> Option { - sig.map(|s| options::metadata(json!({ "signature": s }))) +/// Attach accumulated encrypted reasoning once at the end of the block. +fn reasoning_redacted_meta(sig: Option) -> Option { + sig.map(|s| options::metadata(json!({ "redactedContent": s }))) } /// Extract `GenerateContent` items from a non-streaming content block. @@ -547,18 +595,40 @@ fn reasoning_signature_meta(sig: Option) -> Option { /// `None` when no signature is present). /// `reasoningContent.redactedReasoning` yields a `Reasoning` item with empty /// text and `redactedData` under the same key. -fn extract_content(block: &BedrockContentBlock, content: &mut Vec) { +fn extract_content( + block: &BedrockContentBlock, + content: &mut Vec, + model_id: &str, +) { if let Some(text) = &block.text { content.push(GenerateContent::Text { text: text.clone(), provider_metadata: None, }); + } else if let Some(citation) = &block.citations_content + && let Some(parts) = citation.get("content").and_then(|v| v.as_array()) + { + for part in parts { + if let Some(text) = part.get("text").and_then(|v| v.as_str()) { + content.push(GenerateContent::Text { + text: text.into(), + provider_metadata: None, + }); + } + } } if let Some(tool_use) = &block.tool_use { content.push(GenerateContent::ToolCall(RawToolCall { - tool_call_id: tool_use.tool_use_id.clone(), + tool_call_id: normalize_tool_call_id( + &tool_use.tool_use_id, + model_id.contains("mistral."), + ), tool_name: tool_use.name.clone(), - input: tool_use.input.to_string(), + input: if tool_use.input.is_null() { + "{}".into() + } else { + tool_use.input.to_string() + }, provider_executed: None, dynamic: None, provider_metadata: None, @@ -574,11 +644,17 @@ fn extract_content(block: &BedrockContentBlock, content: &mut Vec) -> Option { + let usage = usage?; + let mut payload = serde_json::Map::new(); + if let Some(tokens) = usage.cache_write_input_tokens { + payload.insert("cacheWriteInputTokens".into(), json!(tokens)); + } + if let Some(details) = &usage.cache_details { + payload.insert( + "cacheDetails".into(), + serde_json::to_value(details).unwrap_or_default(), + ); + } + (!payload.is_empty()).then_some(serde_json::Value::Object(payload)) +} + +fn response_provider_metadata(data: &BedrockConverseResponse) -> Option { + if data.trace.is_none() + && data.usage.is_none() + && data.performance_config.is_none() + && data.service_tier.is_none() + && data + .additional_model_response_fields + .as_ref() + .and_then(|v| v.get("delta")) + .and_then(|v| v.get("stop_sequence")) + .is_none() + { + return None; + } + let mut payload = serde_json::Map::new(); + for (key, value) in [ + ("trace", &data.trace), + ("performanceConfig", &data.performance_config), + ("serviceTier", &data.service_tier), + ] { + if let Some(value) = value { + if key == "trace" && !value.is_object() { + continue; + } + payload.insert(key.into(), value.clone()); + } + } + if let Some(usage) = cache_usage_metadata(data.usage.as_ref()) { + payload.insert("usage".into(), usage); + } + payload.insert( + "stopSequence".into(), + data.additional_model_response_fields + .as_ref() + .and_then(|v| v.get("delta")) + .and_then(|v| v.get("stop_sequence")) + .cloned() + .unwrap_or(serde_json::Value::Null), + ); + Some(options::metadata(serde_json::Value::Object(payload))) +} + +#[derive(Default)] +struct JsonObjectTextExtractor { + started: bool, + completed: bool, + depth: usize, + in_string: bool, + escaped: bool, +} + +impl JsonObjectTextExtractor { + fn process(&mut self, text: &str) -> String { + let mut result = String::new(); + for character in text.chars() { + if self.completed { + break; + } + if !self.started { + if character != '{' { + continue; + } + self.started = true; + self.depth = 1; + result.push(character); + continue; + } + result.push(character); + if self.escaped { + self.escaped = false; + continue; + } + if character == '\\' && self.in_string { + self.escaped = true; + continue; + } + if character == '"' { + self.in_string = !self.in_string; + continue; + } + if self.in_string { + continue; + } + if character == '{' { + self.depth += 1; + } else if character == '}' { + self.depth -= 1; + self.completed = self.depth == 0; + } + } + result + } +} diff --git a/aimux-providers/src/bedrock/options.rs b/aimux-providers/src/bedrock/options.rs index 39fb7132b..cbfb21bff 100644 --- a/aimux-providers/src/bedrock/options.rs +++ b/aimux-providers/src/bedrock/options.rs @@ -1,7 +1,4 @@ -//! The providerOptions and providerMetadata namespaces of Amazon Bedrock. -//! -//! Options use `amazonBedrock`, with `bedrock` as the upstream fallback alias. -//! Response metadata uses the canonical `amazonBedrock` key. +//! Amazon Bedrock canonical and legacy option namespaces. use aimux_core::shared::{JsonObject, SharedProviderOptions, provider_namespace}; use aimux_core::types::ProviderMetadata; @@ -10,6 +7,10 @@ use serde_json::Value; /// The namespace: read and written. pub(crate) const AMAZON_BEDROCK: &str = "amazonBedrock"; +/// Model family names used by Bedrock's family-specific request formats. +pub(crate) const COHERE_MODEL_FAMILY: &str = "cohere"; +pub(crate) const OPENAI_MODEL_FAMILY: &str = "openai"; + /// The Bedrock options in a providerOptions container (`amazonBedrock`). pub(crate) fn read(provider_options: Option<&SharedProviderOptions>) -> Option<&JsonObject> { let options = provider_options?; @@ -20,5 +21,8 @@ pub(crate) fn read(provider_options: Option<&SharedProviderOptions>) -> Option<& /// Wrap `payload` as response metadata under the namespace key. pub(crate) fn metadata(payload: Value) -> ProviderMetadata { - provider_namespace(AMAZON_BEDROCK, payload).expect("provider metadata must be an object") + let mut metadata = + provider_namespace(AMAZON_BEDROCK, payload.clone()).expect("metadata payload is an object"); + metadata.extend(provider_namespace("bedrock", payload).expect("metadata payload is an object")); + metadata } diff --git a/aimux-providers/src/bedrock/reranking.rs b/aimux-providers/src/bedrock/reranking.rs index 3570a9f53..6bc9d2b95 100644 --- a/aimux-providers/src/bedrock/reranking.rs +++ b/aimux-providers/src/bedrock/reranking.rs @@ -31,11 +31,20 @@ struct BedrockRerankingOptions { fn parse_bedrock_reranking_options( provider_options: Option<&SharedProviderOptions>, -) -> BedrockRerankingOptions { +) -> Result { let mut opts = BedrockRerankingOptions::default(); if let Some(po) = provider_options { let bedrock = options::read(Some(po)); if let Some(bedrock) = bedrock { + if bedrock.get("nextToken").is_some_and(|v| !v.is_string()) + || bedrock + .get("additionalModelRequestFields") + .is_some_and(|v| !v.is_object()) + { + return Err(AiMuxError::InvalidArgument( + "invalid amazonBedrock provider options".into(), + )); + } if let Some(token) = bedrock.get("nextToken").and_then(|v| v.as_str()) { opts.next_token = Some(token.to_string()); } @@ -44,7 +53,7 @@ fn parse_bedrock_reranking_options( } } } - opts + Ok(opts) } /// The response from the Bedrock `/rerank` endpoint. @@ -94,7 +103,7 @@ impl RerankingModel for BedrockRerankingModel { &self, options: &RerankingCallOptions, ) -> Result { - let bedrock_options = parse_bedrock_reranking_options(options.provider_options.as_ref()); + let bedrock_options = parse_bedrock_reranking_options(options.provider_options.as_ref())?; let region = (self.region)()?; let model_arn = format!( @@ -142,13 +151,17 @@ impl RerankingModel for BedrockRerankingModel { "modelConfiguration": { "modelArn": model_arn }, - "numberOfResults": options.top_n }, "type": "BEDROCK_RERANKING_MODEL" }, "sources": sources }); + if let Some(top_n) = options.top_n { + body["rerankingConfiguration"]["bedrockRerankingConfiguration"]["numberOfResults"] = + json!(top_n); + } + if let Some(ref token) = bedrock_options.next_token { body["nextToken"] = json!(token); } diff --git a/aimux-providers/src/bedrock/tools.rs b/aimux-providers/src/bedrock/tools.rs new file mode 100644 index 000000000..feb877229 --- /dev/null +++ b/aimux-providers/src/bedrock/tools.rs @@ -0,0 +1,282 @@ +//! Provider-defined tool input schemas from the pinned upstream factories. + +use serde_json::Value; + +fn input_schema(id: &str) -> Option { + let schema = match id { + "anthropic.advisor_20260301" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{},"additionalProperties":false}"# + } + "anthropic.bash_20241022" | "anthropic.bash_20250124" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"command":{"type":"string"},"restart":{"type":"boolean"}},"required":["command"],"additionalProperties":false}"# + } + "anthropic.code_execution_20250522" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"code":{"type":"string"}},"required":["code"],"additionalProperties":false}"# + } + "anthropic.code_execution_20250825" | "anthropic.code_execution_20260120" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","oneOf":[{"type":"object","properties":{"type":{"type":"string","const":"programmatic-tool-call"},"code":{"type":"string"}},"required":["type","code"],"additionalProperties":false},{"type":"object","properties":{"type":{"type":"string","const":"bash_code_execution"},"command":{"type":"string"}},"required":["type","command"],"additionalProperties":false},{"oneOf":[{"type":"object","properties":{"type":{"type":"string","const":"text_editor_code_execution"},"command":{"type":"string","const":"view"},"path":{"type":"string"}},"required":["type","command","path"],"additionalProperties":false},{"type":"object","properties":{"type":{"type":"string","const":"text_editor_code_execution"},"command":{"type":"string","const":"create"},"path":{"type":"string"},"file_text":{"type":["string","null"]}},"required":["type","command","path"],"additionalProperties":false},{"type":"object","properties":{"type":{"type":"string","const":"text_editor_code_execution"},"command":{"type":"string","const":"str_replace"},"path":{"type":"string"},"old_str":{"type":"string"},"new_str":{"type":"string"}},"required":["type","command","path","old_str","new_str"],"additionalProperties":false}]}]}"# + } + "anthropic.computer_20241022" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"action":{"type":"string","enum":["key","type","mouse_move","left_click","left_click_drag","right_click","middle_click","double_click","screenshot","cursor_position"]},"coordinate":{"type":"array","items":{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991}},"text":{"type":"string"}},"required":["action"],"additionalProperties":false}"# + } + "anthropic.computer_20250124" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"action":{"type":"string","enum":["key","hold_key","type","cursor_position","mouse_move","left_mouse_down","left_mouse_up","left_click","left_click_drag","right_click","middle_click","double_click","triple_click","scroll","wait","screenshot"]},"coordinate":{"type":"array","items":[{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991}],"additionalItems":false,"minItems":2,"maxItems":2},"duration":{"type":"number"},"scroll_amount":{"type":"number"},"scroll_direction":{"type":"string","enum":["up","down","left","right"]},"start_coordinate":{"type":"array","items":[{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991}],"additionalItems":false,"minItems":2,"maxItems":2},"text":{"type":"string"}},"required":["action"],"additionalProperties":false}"# + } + "anthropic.computer_20251124" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"action":{"type":"string","enum":["key","hold_key","type","cursor_position","mouse_move","left_mouse_down","left_mouse_up","left_click","left_click_drag","right_click","middle_click","double_click","triple_click","scroll","wait","screenshot","zoom"]},"coordinate":{"type":"array","items":[{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991}],"additionalItems":false,"minItems":2,"maxItems":2},"duration":{"type":"number"},"region":{"type":"array","items":[{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991}],"additionalItems":false,"minItems":4,"maxItems":4},"scroll_amount":{"type":"number"},"scroll_direction":{"type":"string","enum":["up","down","left","right"]},"start_coordinate":{"type":"array","items":[{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991}],"additionalItems":false,"minItems":2,"maxItems":2},"text":{"type":"string"}},"required":["action"],"additionalProperties":false}"# + } + "anthropic.computer_toolset_20260801" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"action":{"type":"string","enum":["screenshot","zoom","left_click","right_click","middle_click","double_click","triple_click","left_click_drag","mouse_move","left_mouse_down","left_mouse_up","cursor_position","scroll","type","key","hold_key","wait"]},"coordinate":{"type":"array","items":[{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991}],"additionalItems":false,"minItems":2,"maxItems":2},"duration":{"type":"number"},"region":{"type":"array","items":[{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991}],"additionalItems":false,"minItems":4,"maxItems":4},"repeat":{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},"scroll_amount":{"type":"number"},"scroll_direction":{"type":"string","enum":["up","down","left","right"]},"start_coordinate":{"type":"array","items":[{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991}],"additionalItems":false,"minItems":2,"maxItems":2},"text":{"type":"string"}},"required":["action"],"additionalProperties":false}"# + } + "anthropic.memory_20250818" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","oneOf":[{"type":"object","properties":{"command":{"type":"string","const":"view"},"path":{"type":"string"},"view_range":{"type":"array","items":[{"type":"number"},{"type":"number"}],"additionalItems":false,"minItems":2,"maxItems":2}},"required":["command","path"],"additionalProperties":false},{"type":"object","properties":{"command":{"type":"string","const":"create"},"path":{"type":"string"},"file_text":{"type":"string"}},"required":["command","path","file_text"],"additionalProperties":false},{"type":"object","properties":{"command":{"type":"string","const":"str_replace"},"path":{"type":"string"},"old_str":{"type":"string"},"new_str":{"type":"string"}},"required":["command","path","old_str","new_str"],"additionalProperties":false},{"type":"object","properties":{"command":{"type":"string","const":"insert"},"path":{"type":"string"},"insert_line":{"type":"number"},"insert_text":{"type":"string"}},"required":["command","path","insert_line","insert_text"],"additionalProperties":false},{"type":"object","properties":{"command":{"type":"string","const":"delete"},"path":{"type":"string"}},"required":["command","path"],"additionalProperties":false},{"type":"object","properties":{"command":{"type":"string","const":"rename"},"old_path":{"type":"string"},"new_path":{"type":"string"}},"required":["command","old_path","new_path"],"additionalProperties":false}]}"# + } + "anthropic.text_editor_20241022" | "anthropic.text_editor_20250124" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"command":{"type":"string","enum":["view","create","str_replace","insert","undo_edit"]},"path":{"type":"string"},"file_text":{"type":"string"},"insert_line":{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},"new_str":{"type":"string"},"insert_text":{"type":"string"},"old_str":{"type":"string"},"view_range":{"type":"array","items":{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991}}},"required":["command","path"],"additionalProperties":false}"# + } + "anthropic.text_editor_20250429" | "anthropic.text_editor_20250728" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"command":{"type":"string","enum":["view","create","str_replace","insert"]},"path":{"type":"string"},"file_text":{"type":"string"},"insert_line":{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991},"new_str":{"type":"string"},"insert_text":{"type":"string"},"old_str":{"type":"string"},"view_range":{"type":"array","items":{"type":"integer","minimum":-9007199254740991,"maximum":9007199254740991}}},"required":["command","path"],"additionalProperties":false}"# + } + "anthropic.web_fetch_20250910" + | "anthropic.web_fetch_20260209" + | "anthropic.web_fetch_20260318" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"url":{"type":"string"}},"required":["url"],"additionalProperties":false}"# + } + "anthropic.web_search_20250305" + | "anthropic.web_search_20260209" + | "anthropic.web_search_20260318" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"query":{"type":"string"}},"required":["query"],"additionalProperties":false}"# + } + "anthropic.tool_search_regex_20251119" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"pattern":{"type":"string"},"limit":{"type":"number"}},"required":["pattern"],"additionalProperties":false}"# + } + "anthropic.tool_search_bm25_20251119" => { + r#"{"$schema":"http://json-schema.org/draft-07/schema#","type":"object","properties":{"query":{"type":"string"},"limit":{"type":"number"}},"required":["query"],"additionalProperties":false}"# + } + _ => return None, + }; + Some(serde_json::from_str(schema).expect("upstream tool schema")) +} + +use aimux_core::AiMuxError; + +use aimux_core::tool::{Tool, ToolChoice}; +use aimux_core::types::Warning; +use serde_json::json; + +use crate::anthropic::prepare_tools::{AnthropicTool, prepare_tools_with_provider}; + +pub(crate) struct PreparedProviderTools { + pub tools: Vec, + pub tool_choice: Option, + pub betas: Vec, + pub warnings: Vec, + pub using_anthropic_tools: bool, +} + +pub(crate) fn prepare_provider_tools( + tools: Option<&[Tool]>, + choice: Option<&ToolChoice>, + anthropic: bool, + disable_parallel: bool, + rejects_forced: bool, +) -> Result { + let mut result = PreparedProviderTools { + tools: Vec::new(), + tool_choice: None, + betas: Vec::new(), + warnings: Vec::new(), + using_anthropic_tools: false, + }; + let providers: Vec<_> = tools + .unwrap_or_default() + .iter() + .filter_map(|tool| { + let Tool::Provider(tool) = tool else { + return None; + }; + if matches!( + tool.id.as_str(), + "anthropic.web_search_20250305" + | "anthropic.web_search_20260318" + | "anthropic.web_fetch_20260318" + ) { + let kind = tool.id.strip_prefix("anthropic.").unwrap(); + result.warnings.push(Warning::Unsupported { + feature: format!("{kind} tool"), + details: Some(format!( + "The {kind} tool is not supported on Amazon Bedrock." + )), + }); + None + } else { + Some(tool) + } + }) + .collect(); + let selected: Vec<_> = providers + .iter() + .copied() + .filter(|tool| { + !anthropic + || !rejects_forced + || !matches!(choice, Some(ToolChoice::Tool { tool_name }) if tool_name != &tool.name) + }) + .collect(); + if !anthropic || selected.is_empty() { + for tool in selected { + result.warnings.push(Warning::Unsupported { + feature: format!("tool {}", tool.id), + details: None, + }); + } + return Ok(result); + } + result.using_anthropic_tools = true; + let anthropic_tools: Vec<_> = providers + .iter() + .map(|tool| AnthropicTool::Provider { + id: tool.id.clone(), + name: tool.name.clone(), + args: tool.args.clone(), + }) + .collect(); + let fallback = + rejects_forced && matches!(choice, Some(ToolChoice::Required | ToolChoice::Tool { .. })); + let prepared = prepare_tools_with_provider( + Some(&anthropic_tools), + if fallback { + Some(&ToolChoice::Auto) + } else { + choice + }, + disable_parallel, + false, + false, + false, + ); + result.tool_choice = prepared.tool_choice; + for tool in &providers { + validate_args(&tool.id, &tool.args)?; + } + result.betas.extend(prepared.betas); + result.warnings.extend(prepared.tool_warnings); + if fallback { + let details = match choice { + Some(ToolChoice::Tool { tool_name }) => format!("toolChoice 'tool' is not supported by this model because it rejects forced tool use. Only the '{tool_name}' tool is sent with 'auto' tool choice. Instruct the model to use the tool in the prompt and verify that a tool call was made."), + _ => "toolChoice 'required' is not supported by this model because it rejects forced tool use. Using 'auto' instead. Instruct the model to use a tool in the prompt and verify that a tool call was made.".to_owned(), + }; + result.warnings.push(Warning::Unsupported { + feature: "toolChoice".into(), + details: Some(details), + }); + } + for tool in selected { + if let Some(schema) = input_schema(&tool.id) { + result + .tools + .push(json!({"toolSpec":{"name":tool.name,"inputSchema":{"json":schema}}})); + } else { + result.warnings.push(Warning::Unsupported { + feature: "tool ${tool.id}".into(), + details: None, + }); + } + } + Ok(result) +} + +// Upstream validates configuration only for these factories; older tools pass +// their configuration through unchecked. +fn validate_args(id: &str, args: &serde_json::Map) -> Result<(), AiMuxError> { + let web_fetch = matches!( + id, + "anthropic.web_fetch_20250910" | "anthropic.web_fetch_20260209" + ); + let web_search = id == "anthropic.web_search_20260209"; + if !web_fetch + && !web_search + && !matches!( + id, + "anthropic.advisor_20260301" + | "anthropic.computer_toolset_20260801" + | "anthropic.text_editor_20250728" + ) + { + return Ok(()); + } + let invalid = || { + AiMuxError::InvalidArgument(format!( + "Invalid configuration for provider-defined tool {id}" + )) + }; + let optional = |name: &str, valid: fn(&Value) -> bool| args.get(name).is_none_or(valid); + let valid = if web_fetch || web_search { + optional("maxUses", Value::is_number) + && optional("allowedDomains", |v| { + v.as_array().is_some_and(|a| a.iter().all(Value::is_string)) + }) + && optional("blockedDomains", |v| { + v.as_array().is_some_and(|a| a.iter().all(Value::is_string)) + }) + && if web_fetch { + optional("maxContentTokens", Value::is_number) + && optional("citations", |v| { + v.is_object() && v.get("enabled").is_some_and(Value::is_boolean) + }) + } else { + optional("userLocation", |v| { + v.is_object() + && v.get("type").and_then(Value::as_str) == Some("approximate") + && ["city", "region", "country", "timezone"] + .iter() + .all(|key| v.get(*key).is_none_or(Value::is_string)) + }) + } + } else if id == "anthropic.advisor_20260301" { + args.get("model").is_some_and(Value::is_string) + && optional("maxUses", Value::is_number) + && optional("maxTokens", |v| { + v.as_f64() + .is_some_and(|n| (1024.0..=9007199254740991.0).contains(&n) && n.fract() == 0.0) + }) + && optional("caching", |v| { + v.is_object() + && v.get("type").and_then(Value::as_str) == Some("ephemeral") + && matches!(v.get("ttl").and_then(Value::as_str), Some("5m" | "1h")) + }) + } else if id == "anthropic.text_editor_20250728" { + optional("maxCharacters", Value::is_number) + } else { + optional("configs", |v| { + v.as_object().is_some_and(|configs| { + configs.iter().all(|(member, config)| { + matches!( + member.as_str(), + "screenshot" + | "zoom" + | "left_click" + | "right_click" + | "middle_click" + | "double_click" + | "triple_click" + | "left_click_drag" + | "mouse_move" + | "left_mouse_down" + | "left_mouse_up" + | "cursor_position" + | "scroll" + | "type" + | "key" + | "hold_key" + | "wait" + ) && config.is_object() + && ["enabled", "deferLoading"] + .iter() + .all(|key| config.get(*key).is_none_or(Value::is_boolean)) + }) + }) + }) + }; + if valid { Ok(()) } else { Err(invalid()) } +} diff --git a/aimux-providers/src/bedrock/types.rs b/aimux-providers/src/bedrock/types.rs index 1a925d740..251a585f4 100644 --- a/aimux-providers/src/bedrock/types.rs +++ b/aimux-providers/src/bedrock/types.rs @@ -28,6 +28,12 @@ pub struct BedrockConverseResponse { pub additional_model_response_fields: Option, #[serde(default)] pub metrics: Option, + #[serde(default)] + pub trace: Option, + #[serde(default)] + pub performance_config: Option, + #[serde(default)] + pub service_tier: Option, } /// `output` block — wraps the generated `message`. @@ -55,6 +61,8 @@ pub struct BedrockContentBlock { pub tool_use: Option, #[serde(default, rename = "reasoningContent")] pub reasoning_content: Option, + #[serde(default, rename = "citationsContent")] + pub citations_content: Option, } /// `toolUse` block. @@ -73,16 +81,28 @@ pub struct BedrockToolUse { #[derive(Debug, Default, Deserialize, Serialize)] #[serde(rename_all = "camelCase")] pub struct BedrockUsage { - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub input_tokens: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub output_tokens: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub total_tokens: Option, - #[serde(default, rename = "cacheReadInputTokens")] + #[serde( + default, + rename = "cacheReadInputTokens", + skip_serializing_if = "Option::is_none" + )] pub cache_read_input_tokens: Option, - #[serde(default, rename = "cacheWriteInputTokens")] + #[serde( + default, + rename = "cacheWriteInputTokens", + skip_serializing_if = "Option::is_none" + )] pub cache_write_input_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cache_details: Option>, + #[serde(flatten)] + pub extra: serde_json::Map, } // ── Error response ─────────────────────────────────────────────────────────── @@ -95,3 +115,20 @@ pub struct BedrockError { #[serde(default, rename = "type")] pub error_type: Option, } + +#[derive(Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct BedrockCacheDetail { + #[serde(deserialize_with = "deserialize_cache_input_tokens")] + pub input_tokens: u32, + pub ttl: String, + #[serde(flatten)] + pub extra: serde_json::Map, +} + +fn deserialize_cache_input_tokens<'de, D: serde::Deserializer<'de>>( + deserializer: D, +) -> Result { + u32::deserialize(deserializer) + .map_err(|error| serde::de::Error::custom(format!("inputTokens: {error}"))) +} diff --git a/aimux-providers/src/cohere/convert.rs b/aimux-providers/src/cohere/convert.rs index 97c96443f..71ae58b65 100644 --- a/aimux-providers/src/cohere/convert.rs +++ b/aimux-providers/src/cohere/convert.rs @@ -3,7 +3,7 @@ //! Mirrors the TS `convert-to-cohere-chat-prompt.ts`, //! `cohere-prepare-tools.ts`, and `map-cohere-finish-reason.ts`. -use aimux_core::error::AiMuxError; +use aimux_core::AiMuxError; use aimux_core::language_model_message::{ AssistantPart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolCallPart, ToolPart, ToolResultPart, UserPart, @@ -66,11 +66,7 @@ pub fn prepare_tools(tools: &Option>, tool_choice: Option<&ToolChoice> } } } - if cohere_tools.is_empty() { - None - } else { - Some(cohere_tools) - } + Some(cohere_tools) } }; @@ -124,7 +120,7 @@ pub struct ConvertedPrompt { /// - Tool messages are flat `{role:"tool", content, tool_call_id}`. /// /// # Errors -/// Returns an error for image text data or an unresolved inline image media type. +/// Rejects provider references, non-image file URLs and invalid image options. pub fn convert_prompt_to_cohere( prompt: &LanguageModelPrompt, ) -> Result { @@ -148,10 +144,10 @@ pub fn convert_prompt_to_cohere( } } UserPart::File(file) => { - use base64::Engine; if get_top_level_media_type(&file.media_type) == "image" { let url = match &file.data { FileData::Data { data } => { + use base64::Engine; let b64 = match data { FileBytes::Binary(bytes) => { base64::engine::general_purpose::STANDARD @@ -168,7 +164,7 @@ pub fn convert_prompt_to_cohere( FileData::Url { url, .. } => url.clone(), FileData::Reference { .. } => { return Err(AiMuxError::UnsupportedFunctionality( - "image file parts with provider references".to_string(), + "image file parts with provider references".into(), )); } FileData::Text { .. } => { @@ -178,9 +174,7 @@ pub fn convert_prompt_to_cohere( } }; has_image = true; - parts.push( - json!({ "type": "image_url", "image_url": { "url": url } }), - ); + parts.push(image_part(url, &file.provider_options)?); } else { let text = match &file.data { FileData::Data { @@ -192,12 +186,12 @@ pub fn convert_prompt_to_cohere( FileData::Text { text } => text.clone(), FileData::Url { .. } => { return Err(AiMuxError::UnsupportedFunctionality( - "File URL data: URLs should be downloaded by the AI SDK and not reach this point. This indicates a configuration issue.".to_string(), + "File URL data: URLs should be downloaded by the AI SDK and not reach this point. This indicates a configuration issue.".into(), )); } FileData::Reference { .. } => { return Err(AiMuxError::UnsupportedFunctionality( - "file parts with provider references".to_string(), + "file parts with provider references".into(), )); } }; @@ -294,6 +288,24 @@ pub fn convert_prompt_to_cohere( }) } +fn image_part( + url: String, + provider_options: &Option, +) -> Result { + let mut image = json!({ "url": url }); + if let Some(detail) = super::options::cohere_options(provider_options.as_ref()) + .and_then(|options| options.get("detail")) + { + if !matches!(detail.as_str(), Some("auto" | "low" | "high")) { + return Err(AiMuxError::InvalidArgument( + "cohere.detail must be auto, low or high".into(), + )); + } + image["detail"] = detail.clone(); + } + Ok(json!({ "type": "image_url", "image_url": image })) +} + fn join_text_parts(content: &[AssistantPart]) -> String { content .iter() @@ -322,7 +334,7 @@ pub struct RequestBodyResult { /// and the `thinking` config resolved from `reasoning` / provider options. /// /// # Errors -/// Returns an error for image text data or an unresolved inline image media type. +/// Rejects unsupported prompt files and invalid provider options. pub fn build_request_body( model_id: &str, options: &CallOptions, @@ -374,10 +386,11 @@ pub fn build_request_body( match rf { ResponseFormat::Text => {} ResponseFormat::Json { schema, .. } => { - body["response_format"] = json!({ - "type": "json_object", - "json_schema": schema, - }); + let mut format = json!({ "type": "json_object" }); + if let Some(schema) = schema { + format["json_schema"] = schema.clone(); + } + body["response_format"] = format; } } } @@ -429,7 +442,7 @@ fn reasoning_budget_percentage(reasoning: ReasoningEffort) -> Option { /// - other reasoning levels → `{ type: "enabled", token_budget: }`. /// /// # Errors -/// Returns a parse error for invalid Cohere thinking options. +/// Rejects non-object thinking options, unknown types and nonnumeric budgets. pub fn resolve_cohere_thinking( reasoning: Option, provider_options: &Option, @@ -438,30 +451,29 @@ pub fn resolve_cohere_thinking( if let Some(cohere) = super::options::cohere_options(provider_options.as_ref()) && let Some(thinking) = cohere.get("thinking") { - let thinking = thinking.as_object().ok_or_else(|| { - AiMuxError::InvalidArgument("invalid cohere provider options".to_string()) - })?; - let thinking = crate::openai::convert::parse_option_fields( - thinking, - crate::cohere::options::NAMESPACE, - |key, value| { - let valid = match key { - "type" => value - .as_str() - .is_some_and(|value| matches!(value, "enabled" | "disabled")), - "tokenBudget" => value.is_number(), - _ => return None, - }; - Some(if valid { Ok(value.clone()) } else { Err(()) }) + if !thinking.is_object() { + return Err(AiMuxError::InvalidArgument( + "cohere.thinking must be an object".into(), + )); + } + let t_type = match thinking.get("type") { + None => "enabled", + Some(value) => match value.as_str() { + Some(value @ ("enabled" | "disabled")) => value, + _ => { + return Err(AiMuxError::InvalidArgument( + "cohere.thinking.type must be enabled or disabled".into(), + )); + } }, - )?; - let t_type = thinking - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or("enabled") - .to_string(); + }; let mut obj = json!({ "type": t_type }); if let Some(budget) = thinking.get("tokenBudget") { + if !budget.is_number() { + return Err(AiMuxError::InvalidArgument( + "cohere.thinking.tokenBudget must be a number".into(), + )); + } obj["token_budget"] = budget.clone(); } return Ok(Some(obj)); diff --git a/aimux-providers/src/cohere/embedding.rs b/aimux-providers/src/cohere/embedding.rs index 9ca536956..2e2eef584 100644 --- a/aimux-providers/src/cohere/embedding.rs +++ b/aimux-providers/src/cohere/embedding.rs @@ -6,6 +6,7 @@ //! Endpoint: `POST {base_url}/embed` use async_trait::async_trait; +use serde::Deserialize; use serde_json::{Map, Value, json}; use aimux_core::embedding_model::{ @@ -50,7 +51,14 @@ impl EmbeddingModel for CohereEmbeddingModel { &self, options: &EmbeddingCallOptions, ) -> Result { - let cohere_options = parse_cohere_provider_options(options.provider_options.as_ref()); + let cohere_options = parse_cohere_provider_options(options.provider_options.as_ref())?; + if options.values.len() > 96 { + return Err(AiMuxError::InvalidArgument(format!( + "Too many values for a single embedding call: {} supports at most 96 values, received {}", + self.provider(), + options.values.len() + ))); + } let mut body = Map::new(); body.insert("model".to_string(), json!(self.model_id)); @@ -77,43 +85,19 @@ impl EmbeddingModel for CohereEmbeddingModel { let resp = aimux_provider_utils::post_json_to_api( exchange.request(exchange.url("/embed"), options), exchange.transform_body(Value::Object(body)), - aimux_provider_utils::create_json_response_handler(), + aimux_provider_utils::create_json_response_handler::(), super::cohere_failed_response_handler(), ) .await?; let response_headers = resp.response_headers; - let raw_value: Value = resp.value; - - // Extract embeddings: response.embeddings.float - let embeddings: Vec> = raw_value - .get("embeddings") - .and_then(|e| e.get("float")) - .and_then(|f| f.as_array()) - .map(|arr| { - arr.iter() - .map(|row| { - row.as_array() - .unwrap_or(&vec![]) - .iter() - .filter_map(|v| v.as_f64().map(|f| f as f32)) - .collect() - }) - .collect() - }) - .unwrap_or_default(); - - // Extract usage: response.meta.billed_units.input_tokens - let usage = raw_value - .get("meta") - .and_then(|m| m.get("billed_units")) - .and_then(|b| b.get("input_tokens")) - .and_then(serde_json::Value::as_u64) - .map(|tokens| EmbeddingUsage { - tokens: tokens as u32, - }) - .unwrap_or_default(); + let raw_value = resp.raw_value; + let data = resp.value; + let embeddings = data.embeddings.float; + let usage = EmbeddingUsage { + tokens: data.meta.billed_units.input_tokens as u32, + }; Ok(EmbeddingResult { embeddings, @@ -121,7 +105,7 @@ impl EmbeddingModel for CohereEmbeddingModel { provider_metadata: None, response: Some(EmbeddingResponse { headers: Some(response_headers), - body: Some(raw_value), + body: raw_value, }), warnings: Vec::new(), }) @@ -138,20 +122,62 @@ struct CohereEmbeddingProviderOptions { fn parse_cohere_provider_options( options: Option<&SharedProviderOptions>, -) -> CohereEmbeddingProviderOptions { +) -> Result { let provider_opts = super::options::cohere_options(options); - CohereEmbeddingProviderOptions { - input_type: provider_opts - .and_then(|o| o.get("inputType")) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - truncate: provider_opts - .and_then(|o| o.get("truncate")) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - output_dimension: provider_opts - .and_then(|o| o.get("outputDimension")) - .and_then(serde_json::Value::as_u64) - .map(|d| d as u32), - } + let string_option = |key: &str, allowed: &[&str]| -> Result, AiMuxError> { + provider_opts + .and_then(|o| o.get(key)) + .map(|value| { + value + .as_str() + .filter(|value| allowed.contains(value)) + .map(str::to_owned) + .ok_or_else(|| AiMuxError::InvalidArgument(format!("Invalid cohere.{key}"))) + }) + .transpose() + }; + let output_dimension = provider_opts + .and_then(|o| o.get("outputDimension")) + .map(|value| { + value + .as_f64() + .filter(|value| [256.0, 512.0, 1024.0, 1536.0].contains(value)) + .map(|value| value as u32) + .ok_or_else(|| AiMuxError::InvalidArgument("Invalid cohere.outputDimension".into())) + }) + .transpose()?; + Ok(CohereEmbeddingProviderOptions { + input_type: string_option( + "inputType", + &[ + "search_document", + "search_query", + "classification", + "clustering", + ], + )?, + truncate: string_option("truncate", &["NONE", "START", "END"])?, + output_dimension, + }) +} + +#[derive(Deserialize)] +struct CohereEmbeddingResponse { + embeddings: CohereFloatEmbeddings, + meta: CohereEmbeddingMeta, +} + +#[derive(Deserialize)] +struct CohereFloatEmbeddings { + float: Vec>, +} + +#[derive(Deserialize)] +struct CohereEmbeddingMeta { + billed_units: CohereEmbeddingBilledUnits, +} + +#[derive(Deserialize)] +struct CohereEmbeddingBilledUnits { + input_tokens: f64, } diff --git a/aimux-providers/src/cohere/mod.rs b/aimux-providers/src/cohere/mod.rs index aec4b1d61..eadfd8302 100644 --- a/aimux-providers/src/cohere/mod.rs +++ b/aimux-providers/src/cohere/mod.rs @@ -1,7 +1,7 @@ //! Cohere provider. //! //! [`create_cohere`] is the Rust form of the AI SDK's `createCohere`: it takes -//! [`CohereProviderSettings`], validates the base URL, fixes the provider name +//! [`CohereProviderSettings`], normalizes the base URL, fixes the provider name //! and returns a [`CohereProvider`]. The API key is not read there; it is //! loaded in the request headers of every call, from the setting or from //! `COHERE_API_KEY`. [`cohere()`] is the default instance. @@ -33,7 +33,9 @@ use aimux_core::language_model::LanguageModel; use aimux_core::model_catalogue::RuntimeModel; use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_core::reranking_model::RerankingModel; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{ + FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, without_trailing_slash, +}; use crate::shared::{Credential, EndpointConfig, provider_headers}; @@ -73,11 +75,11 @@ pub struct CohereProviderSettings { /// Extra headers on every request. A `None` value removes the header, /// including `Authorization`. Per-call headers win over these. pub headers: Option, + /// Generates a unique identifier for each citation. + pub generate_id: Option String + Send + Sync>>, /// The transport: a mock, a signing decorator, a proxy-aware client. /// `None` uses the process default, resolved per request. pub fetch: Option, - /// Generates unique IDs for chat output. - pub generate_id: Option String + Send + Sync>>, } impl std::fmt::Debug for CohereProviderSettings { @@ -90,42 +92,44 @@ impl std::fmt::Debug for CohereProviderSettings { "headers", &self.headers.as_ref().map(std::collections::HashMap::len), ) - .field("fetch", &self.fetch.is_some()) .field("generate_id", &self.generate_id.is_some()) + .field("fetch", &self.fetch.is_some()) .finish() } } /// Create a Cohere provider. /// +/// Only a trailing slash is removed from `base_url`. The API key is loaded +/// per request, not here. +/// /// # Errors /// -/// Returns `AiMuxError::InvalidArgument` when `base_url` is not an `http(s)` -/// URL with a host. That is the only way this fails: the key is loaded per -/// request, not here. +/// This factory does not fail; request validation happens when a model is called. pub fn create_cohere(settings: CohereProviderSettings) -> Result { let base_url = match settings.base_url.as_deref() { - Some(url) => validate_base_url(url)?, + Some(url) => without_trailing_slash(url), None => DEFAULT_BASE_URL.to_string(), }; + let headers = provider_headers( + Credential::explicit_or_env( + settings.api_key.map(Resolvable::Value), + API_KEY_ENV_VAR, + "Cohere", + ), + Vec::new(), + settings.headers, + ); + let headers = aimux_provider_utils::headers::with_user_agent_suffix_fn( + headers, + options::NAMESPACE, + "4.0.52", + ); Ok(CohereProvider { - name: DEFAULT_NAME.to_string(), + generate_id: settings.generate_id, base_url, - headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( - provider_headers( - Credential::explicit_or_env( - settings.api_key.map(Resolvable::Value), - API_KEY_ENV_VAR, - "Cohere", - ), - Vec::new(), - settings.headers, - ), - options::NAMESPACE, - "4.0.52", - ), + headers, fetch: settings.fetch, - generate_id: settings.generate_id, }) } @@ -143,17 +147,16 @@ pub fn cohere() -> &'static CohereProvider { /// A Cohere provider (the AI SDK's `CohereProvider`). Cheap to clone the /// models out of; it holds no HTTP client. pub struct CohereProvider { - name: String, + generate_id: Option String + Send + Sync>>, base_url: String, headers: HeadersFn, fetch: Option, - generate_id: Option String + Send + Sync>>, } impl CohereProvider { fn model_config(&self, method: &str) -> EndpointConfig { EndpointConfig::fixed( - format!("{}.{method}", self.name), + format!("{DEFAULT_NAME}.{method}"), self.base_url.clone(), self.headers.clone(), self.fetch.clone(), @@ -161,7 +164,7 @@ impl CohereProvider { ) } - /// A chat model; `provider()` is `"{name}.chat"`. + /// A chat model; `provider()` is `"cohere.chat"`. #[must_use] pub fn chat(&self, model_id: &str) -> CohereModel { CohereModel::from_config(model_id.to_string(), self.model_config("chat")) @@ -169,14 +172,14 @@ impl CohereProvider { } /// An embedding model (e.g. `"embed-english-v3.0"`); `provider()` is - /// `"{name}.textEmbedding"`. + /// `"cohere.textEmbedding"`. #[must_use] pub fn embedding(&self, model_id: &str) -> CohereEmbeddingModel { CohereEmbeddingModel::from_config(model_id.to_string(), self.model_config("textEmbedding")) } /// A reranking model (e.g. `"rerank-english-v3.0"`); `provider()` is - /// `"{name}.reranking"`. + /// `"cohere.reranking"`. #[must_use] pub fn reranking(&self, model_id: &str) -> CohereRerankingModel { CohereRerankingModel::from_config(model_id.to_string(), self.model_config("reranking")) diff --git a/aimux-providers/src/cohere/model.rs b/aimux-providers/src/cohere/model.rs index 9c9a3685e..9c3a1e825 100644 --- a/aimux-providers/src/cohere/model.rs +++ b/aimux-providers/src/cohere/model.rs @@ -19,33 +19,33 @@ use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usa use crate::shared::EndpointConfig; use super::convert::{build_request_body, parse_finish_reason}; -use super::types::{ChatResponse, StreamEvent, TokenPair}; +use super::types::{ChatResponse, StreamEvent, UsageResponse}; +use std::sync::Arc; /// A Cohere language model. pub struct CohereModel { model_id: String, config: EndpointConfig, - generate_id: std::sync::Arc String + Send + Sync>, + generate_id: Arc String + Send + Sync>, } impl CohereModel { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { + Self { + model_id, + config, + generate_id: Arc::new(aimux_provider_utils::generate_id), + } + } pub(crate) fn with_generate_id( mut self, - generate_id: Option String + Send + Sync>>, + generate_id: Option String + Send + Sync>>, ) -> Self { if let Some(generate_id) = generate_id { self.generate_id = generate_id; } self } - - pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { - Self { - model_id, - config, - generate_id: std::sync::Arc::new(aimux_provider_utils::generate_id), - } - } } // ── Usage conversion ───────────────────────────────────────────────────────── @@ -55,7 +55,8 @@ impl CohereModel { /// Mirrors the TS `convertCohereUsage`: /// - `input.total = input_tokens`, `input.noCache = input_tokens` /// - `output.total = output_tokens` -fn convert_usage(tokens: &TokenPair) -> Usage { +fn convert_usage(usage: &UsageResponse) -> Usage { + let tokens = &usage.tokens; Usage { input_tokens: aimux_core::types::InputTokenUsage { total: Some(tokens.input_tokens), @@ -65,12 +66,11 @@ fn convert_usage(tokens: &TokenPair) -> Usage { }, output_tokens: aimux_core::types::OutputTokenUsage { total: Some(tokens.output_tokens), + text: Some(tokens.output_tokens), ..Default::default() }, // RFC-0015 P0-3: keep the raw provider usage payload. - raw: serde_json::to_value(tokens) - .ok() - .and_then(|value| value.as_object().cloned()), + raw: usage.raw.as_object().cloned(), } } @@ -106,7 +106,6 @@ impl LanguageModel for CohereModel { ) .await?; - let response_body = resp.raw_value; let response_headers = resp.response_headers; let data: ChatResponse = resp.value; @@ -155,7 +154,7 @@ impl LanguageModel for CohereModel { .and_then(|src| src.get("document")) .and_then(|d| d.get("title")) .and_then(|t| t.as_str()) - .filter(|s| !s.is_empty()) + .filter(|title| !title.is_empty()) .map(std::string::ToString::to_string) .unwrap_or_else(|| "Document".to_string()); @@ -212,7 +211,7 @@ impl LanguageModel for CohereModel { } let finish_reason = parse_finish_reason(&data.finish_reason); - let usage = convert_usage(&data.usage.tokens); + let usage = convert_usage(&data.usage); Ok(GenerateResult { content, @@ -225,7 +224,7 @@ impl LanguageModel for CohereModel { timestamp: None, model_id: None, headers: Some(response_headers), - body: response_body, + body: resp.raw_value, }), request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), }) @@ -239,301 +238,107 @@ impl LanguageModel for CohereModel { let resp = aimux_provider_utils::post_json_to_api( exchange.request(endpoint.clone(), options), body.clone(), - aimux_provider_utils::create_event_source_response_handler::(), + aimux_provider_utils::create_event_source_response_handler::(), super::cohere_failed_response_handler(), ) .await?; let response_headers = resp.response_headers; - let mut sse_stream = resp.value; - let first_event = match sse_stream.next().await { - Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), - first_event => first_event, - }; - match first_event.as_ref() { - Some(Ok(event)) if event.event_type == "error" => { - return Err(cohere_stream_error( - event, - &endpoint, - &body, - &response_headers, - )); - } - _ => {} - } - + let sse_stream = resp.value; let stream_warnings = body_result.warnings; - let stream_body = body.clone(); - let stream_response_headers = response_headers.clone(); + let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings: stream_warnings }); - let mut final_usage = Usage::default(); - let mut final_finish_reason: Option = None; + let mut finish_reason = FinishReason { unified: FinishReasonUnified::Other, raw: None }; let mut pending_tool_call: Option = None; let mut is_reasoning = false; - let mut stream_errored = false; - - let mut sse_iter = Box::pin( - futures::stream::iter(first_event.into_iter()).chain(sse_stream), - ); - + let mut sse_iter = sse_stream; while let Some(event) = sse_iter.next().await { - if stream_errored { - break; - } - - match event { - Ok(parsed) => { - match parsed.event_type.as_str() { - "message-start" => { - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: parsed.id.clone(), - timestamp: None, - model_id: None, - })); - } - - "content-start" => { - let idx = parsed.index.unwrap_or(0); - let content_type = parsed - .delta - .as_ref() - .and_then(|d| d.message.as_ref()) - .and_then(|m| m.content.as_ref()) - .and_then(|c| c.get("type")) - .and_then(|t| t.as_str()); - - if content_type == Some("thinking") { - is_reasoning = true; - yield Ok(StreamPart::ReasoningStart { - id: format!("reasoning-{idx}"), - provider_metadata: None, - }); - } else { - yield Ok(StreamPart::TextStart { - id: format!("{idx}"), - provider_metadata: None, - }); - } - } - - "content-delta" => { - let idx = parsed.index.unwrap_or(0); - let content = parsed - .delta - .as_ref() - .and_then(|d| d.message.as_ref()) - .and_then(|m| m.content.as_ref()); - - if let Some(content) = content { - // Thinking delta. - if let Some(thinking) = - content.get("thinking").and_then(|t| t.as_str()) - { - yield Ok(StreamPart::ReasoningDelta { - id: format!("reasoning-{idx}"), - delta: thinking.to_string(), - provider_metadata: None, - }); - } - // Text delta. - else if let Some(text) = - content.get("text").and_then(|t| t.as_str()) - { - yield Ok(StreamPart::TextDelta { - id: format!("{idx}"), - delta: text.to_string(), - provider_metadata: None, - }); - } - } - } - - "content-end" => { - let idx = parsed.index.unwrap_or(0); - if is_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: format!("reasoning-{idx}"), - provider_metadata: None, - }); - is_reasoning = false; - } else { - yield Ok(StreamPart::TextEnd { - id: format!("{idx}"), - provider_metadata: None, - }); - } - } - - "tool-plan-delta" => { - // Tool plan deltas are not emitted as stream parts - // (no corresponding variant). Silently consume. - } - - "tool-call-start" => { - let tc = parsed - .delta - .as_ref() - .and_then(|d| d.message.as_ref()) - .and_then(|m| m.tool_calls.as_ref()); - - if let Some(tc) = tc { - let id = tc - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = tc - .get("function") - .and_then(|f| f.get("name")) - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let initial_args = tc - .get("function") - .and_then(|f| f.get("arguments")) - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - - pending_tool_call = Some(PendingToolCall { - id: id.clone(), - name: name.clone(), - arguments: initial_args.clone(), - }); - - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: name, - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - - if !initial_args.is_empty() { - yield Ok(StreamPart::ToolInputDelta { - id, - delta: initial_args, - provider_metadata: None, - }); - } - } - } - - "tool-call-delta" => { - if let Some(ptc) = &mut pending_tool_call { - let args_delta = parsed - .delta - .as_ref() - .and_then(|d| d.message.as_ref()) - .and_then(|m| m.tool_calls.as_ref()) - .and_then(|tc| tc.get("function")) - .and_then(|f| f.get("arguments")) - .and_then(|v| v.as_str()) - .unwrap_or(""); - - if !args_delta.is_empty() { - ptc.arguments.push_str(args_delta); - yield Ok(StreamPart::ToolInputDelta { - id: ptc.id.clone(), - delta: args_delta.to_string(), - provider_metadata: None, - }); - } - } - } - - "tool-call-end" => { - if let Some(ptc) = pending_tool_call.take() { - yield Ok(StreamPart::ToolInputEnd { - id: ptc.id.clone(), - provider_metadata: None, - }); - - // Providers never parse tool input — Core - // owns JSON parsing, schema validation, - // and repair (aimux-core::parse_tool_call). - // Trim the accumulated text and default - // empty to "{}" to match the TS - // provider's convention; forward it - // verbatim otherwise, malformed JSON - // included, so a bad call surfaces as a - // retained `invalid: true` tool call - // instead of a terminal stream error. - let trimmed = ptc.arguments.trim(); - let text = if trimmed.is_empty() { "{}" } else { trimmed }; - let input = text.to_string(); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: ptc.id, - tool_name: ptc.name, - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - } - } - - "message-end" => { - if let Some(delta) = &parsed.delta { - if let Some(reason) = &delta.finish_reason { - final_finish_reason = Some(parse_finish_reason(reason)); - } - if let Some(usage) = &delta.usage { - final_usage = convert_usage(&usage.tokens); - } - } - } - - "error" => { - yield Ok(StreamPart::Error { - error: cohere_stream_error( - &parsed, - &endpoint, - &stream_body, - &stream_response_headers, - ), - }); - stream_errored = true; - break; - } - - // citation-start, citation-end, and any unknown - // event types are silently consumed. - _ => {} - } + let raw = match event { + Ok(raw) => raw, + Err(error) => { + if !error.is_recoverable_stream_error() { yield Err(error); return; } + if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: Value::Null }); } + finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; + yield Ok(StreamPart::Error { error }); + continue; } + }; + if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: raw.clone() }); } + let parsed = match StreamEvent::parse(raw) { + Ok(parsed) => parsed, Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; + finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; + yield Ok(StreamPart::Error { error }); + continue; + } + }; + let id = parsed.index.map(|index| index.to_string()).unwrap_or_default(); + match parsed.event_type.as_str() { + "message-start" => yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { id: parsed.id, timestamp: None, model_id: None })), + "content-start" => { + let content = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.content.as_ref()).unwrap(); + if content["type"] == "thinking" { + is_reasoning = true; + yield Ok(StreamPart::ReasoningStart { id, provider_metadata: None }); + } else { + yield Ok(StreamPart::TextStart { id, provider_metadata: None }); } } + "content-delta" => { + let content = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.content.as_ref()).unwrap(); + if let Some(text) = content.get("text").and_then(Value::as_str) { + yield Ok(StreamPart::TextDelta { id, delta: text.into(), provider_metadata: None }); + } else { + yield Ok(StreamPart::ReasoningDelta { id, delta: content["thinking"].as_str().unwrap().into(), provider_metadata: None }); + } + } + "content-end" => { + if is_reasoning { + yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); + is_reasoning = false; + } else { + yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); + } + } + "tool-call-start" => { + let tool = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.tool_calls.as_ref()).unwrap(); + let id = tool["id"].as_str().unwrap().to_string(); + let name = tool["function"]["name"].as_str().unwrap().to_string(); + let arguments = tool["function"]["arguments"].as_str().unwrap().to_string(); + pending_tool_call = Some(PendingToolCall { id: id.clone(), name: name.clone(), arguments: arguments.clone() }); + yield Ok(StreamPart::ToolInputStart { id: id.clone(), tool_name: name, provider_executed: None, dynamic: None, title: None, provider_metadata: None }); + if !arguments.is_empty() { yield Ok(StreamPart::ToolInputDelta { id, delta: arguments, provider_metadata: None }); } + } + "tool-call-delta" => { + if let Some(tool) = &mut pending_tool_call { + let delta = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.tool_calls.as_ref()).unwrap()["function"]["arguments"].as_str().unwrap().to_string(); + tool.arguments.push_str(&delta); + yield Ok(StreamPart::ToolInputDelta { id: tool.id.clone(), delta, provider_metadata: None }); + } + } + "tool-call-end" => { + if let Some(tool) = pending_tool_call.take() { + yield Ok(StreamPart::ToolInputEnd { id: tool.id.clone(), provider_metadata: None }); + let text = tool.arguments.trim(); + let input = match serde_json::from_str::(if text.is_empty() { "{}" } else { text }) { + Ok(value) if !contains_prototype_key(&value) => value.to_string(), + Ok(_) => { yield Err(AiMuxError::InvalidResponseData("Object contains forbidden prototype property".into())); return; } + Err(error) => { yield Err(AiMuxError::JsonParse(error.to_string())); return; } + }; + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: tool.id, tool_name: tool.name, input, provider_executed: None, dynamic: None, provider_metadata: None })); + } + } + "message-end" => { + let delta = parsed.delta.unwrap(); + finish_reason = parse_finish_reason(delta.finish_reason.as_ref().unwrap()); + final_usage = convert_usage(delta.usage.as_ref().unwrap()); + } + _ => {} } } - - yield Ok(StreamPart::Finish { - finish_reason: if stream_errored { - FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - } - } else { - final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }) - }, - usage: if stream_errored { - Usage::default() - } else { - final_usage - }, - provider_metadata: Some(super::options::cohere_metadata(serde_json::json!({}))), - }); + yield Ok(StreamPart::Finish { finish_reason, usage: final_usage, provider_metadata: None }); }; Ok(StreamResult { @@ -546,32 +351,14 @@ impl LanguageModel for CohereModel { } } -fn cohere_stream_error( - event: &StreamEvent, - url: &str, - request_body: &Value, - response_headers: &std::collections::HashMap, -) -> AiMuxError { - let status = event.status_code; - let provider_code = event.code.as_ref().and_then(|value| match value { - Value::String(value) => Some(value.clone()), - Value::Number(value) => Some(value.to_string()), - _ => None, - }); - AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { - status_code: status, - provider_code, - response_body: serde_json::to_string(event).ok(), - response_headers: Some(response_headers.clone()), - data: serde_json::to_value(event).ok(), - is_retryable: status.is_some_and(aimux_core::error::is_retryable_status), - ..aimux_core::ApiCallError::new( - event - .message - .clone() - .unwrap_or_else(|| "Cohere stream error".to_string()), - url, - aimux_provider_utils::logging::redact_error_context(request_body.clone()), - ) - })) +fn contains_prototype_key(value: &Value) -> bool { + match value { + Value::Object(object) => object.iter().any(|(key, value)| { + key == "__proto__" + || (key == "constructor" && value.is_object() && value.get("prototype").is_some()) + || contains_prototype_key(value) + }), + Value::Array(values) => values.iter().any(contains_prototype_key), + _ => false, + } } diff --git a/aimux-providers/src/cohere/options.rs b/aimux-providers/src/cohere/options.rs index d8d4741f3..be149bf06 100644 --- a/aimux-providers/src/cohere/options.rs +++ b/aimux-providers/src/cohere/options.rs @@ -20,5 +20,5 @@ pub(crate) fn cohere_options( /// `{ "cohere": payload }`. pub(crate) fn cohere_metadata(payload: Value) -> ProviderMetadata { - provider_namespace(NAMESPACE, payload).expect("provider metadata must be an object") + provider_namespace(NAMESPACE, payload).expect("metadata payload is an object") } diff --git a/aimux-providers/src/cohere/reranking.rs b/aimux-providers/src/cohere/reranking.rs index bd35ae9d2..ac25fe153 100644 --- a/aimux-providers/src/cohere/reranking.rs +++ b/aimux-providers/src/cohere/reranking.rs @@ -20,26 +20,27 @@ use crate::shared::EndpointConfig; /// Cohere provider-specific reranking options. #[derive(Debug, Clone, Default)] struct CohereRerankingOptions { - max_tokens_per_doc: Option, - priority: Option, + max_tokens_per_doc: Option, + priority: Option, } fn parse_cohere_reranking_options( provider_options: Option<&SharedProviderOptions>, -) -> CohereRerankingOptions { - let mut opts = CohereRerankingOptions::default(); - if let Some(cohere) = super::options::cohere_options(provider_options) { - if let Some(v) = cohere - .get("maxTokensPerDoc") - .and_then(serde_json::Value::as_u64) - { - opts.max_tokens_per_doc = Some(v); - } - if let Some(v) = cohere.get("priority").and_then(serde_json::Value::as_u64) { - opts.priority = Some(v); - } - } - opts +) -> Result { + let cohere = super::options::cohere_options(provider_options); + let number_option = |key: &str| -> Result, AiMuxError> { + cohere + .and_then(|options| options.get(key)) + .map(|value| match value { + Value::Number(number) => Ok(number.clone()), + _ => Err(AiMuxError::InvalidArgument(format!("Invalid cohere.{key}"))), + }) + .transpose() + }; + Ok(CohereRerankingOptions { + max_tokens_per_doc: number_option("maxTokensPerDoc")?, + priority: number_option("priority")?, + }) } /// The response from the Cohere `/rerank` endpoint. @@ -82,7 +83,7 @@ impl RerankingModel for CohereRerankingModel { &self, options: &RerankingCallOptions, ) -> Result { - let cohere_options = parse_cohere_reranking_options(options.provider_options.as_ref()); + let cohere_options = parse_cohere_reranking_options(options.provider_options.as_ref())?; let mut warnings = Vec::new(); @@ -105,9 +106,11 @@ impl RerankingModel for CohereRerankingModel { "model": self.model_id, "query": options.query, "documents": documents, - "top_n": options.top_n, }); + if let Some(top_n) = options.top_n { + body["top_n"] = json!(top_n); + } if let Some(max_tokens) = cohere_options.max_tokens_per_doc { body["max_tokens_per_doc"] = json!(max_tokens); } diff --git a/aimux-providers/src/cohere/types.rs b/aimux-providers/src/cohere/types.rs index 1ade29189..bbff669ee 100644 --- a/aimux-providers/src/cohere/types.rs +++ b/aimux-providers/src/cohere/types.rs @@ -54,11 +54,30 @@ pub struct FunctionCallResponse { pub arguments: String, } -#[derive(Debug, Deserialize)] +#[derive(Debug)] pub struct UsageResponse { - #[serde(default)] - pub billed_units: Option, pub tokens: TokenPair, + pub raw: Value, +} + +impl<'de> Deserialize<'de> for UsageResponse { + fn deserialize>(deserializer: D) -> Result { + let raw = Value::deserialize(deserializer)?; + #[derive(Deserialize)] + struct Validated { + tokens: TokenPair, + #[serde(default)] + billed_units: Option, + #[serde(default)] + cached_tokens: Option, + } + let validated: Validated = + serde_json::from_value(raw.clone()).map_err(serde::de::Error::custom)?; + Ok(Self { + tokens: validated.tokens, + raw, + }) + } } #[derive(Debug, Deserialize, Serialize)] @@ -73,25 +92,19 @@ pub struct TokenPair { // always has a `type` field matching the event name. We parse as a generic // `Value` and dispatch on the `type` field in the model code. -#[derive(Debug, Deserialize, Serialize)] +#[derive(Debug, Deserialize)] pub struct StreamEvent { #[serde(rename = "type")] pub event_type: String, #[serde(default)] - pub index: Option, + pub index: Option, #[serde(default)] pub id: Option, #[serde(default)] pub delta: Option, - #[serde(default)] - pub message: Option, - #[serde(default)] - pub status_code: Option, - #[serde(default)] - pub code: Option, } -#[derive(Debug, Deserialize, Serialize, Default)] +#[derive(Debug, Deserialize, Default)] pub struct StreamDelta { #[serde(default)] pub message: Option, @@ -101,7 +114,7 @@ pub struct StreamDelta { pub usage: Option, } -#[derive(Debug, Deserialize, Serialize, Default)] +#[derive(Debug, Deserialize, Default)] pub struct StreamMessage { /// Content can be {type:"text",text:""} or {type:"thinking",thinking:""}. #[serde(default)] @@ -112,11 +125,72 @@ pub struct StreamMessage { /// Tool plan string. #[serde(default)] pub tool_plan: Option, - #[serde(default)] - pub role: Option, } -#[derive(Debug, Deserialize, Serialize)] -pub struct StreamUsage { - pub tokens: TokenPair, +pub type StreamUsage = UsageResponse; + +impl StreamEvent { + pub fn parse(raw: Value) -> Result { + let valid = match raw["type"].as_str() { + Some("citation-start" | "citation-end" | "tool-call-end") => true, + Some("message-start") => raw + .get("id") + .is_none_or(|id| id.is_null() || id.is_string()), + Some("content-end") => raw["index"].is_number(), + Some("content-start" | "content-delta") => { + let content = &raw["delta"]["message"]["content"]; + raw["index"].is_number() + && if raw["type"] == "content-delta" { + content["text"].is_string() || content["thinking"].is_string() + } else { + match content["type"].as_str() { + Some("text") => content["text"].is_string(), + Some("thinking") => content["thinking"].is_string(), + _ => false, + } + } + } + Some("message-end") => { + raw["delta"]["finish_reason"].is_string() && raw["delta"]["usage"].is_object() + } + Some("tool-plan-delta") => raw["delta"]["message"]["tool_plan"].is_string(), + Some("tool-call-start" | "tool-call-delta") => { + let tool = &raw["delta"]["message"]["tool_calls"]; + tool["function"]["arguments"].is_string() + && (raw["type"] == "tool-call-delta" + || (tool["id"].is_string() + && tool["type"] == "function" + && tool["function"]["name"].is_string())) + } + _ => false, + }; + if valid { + // Strip fields outside the selected upstream event schema before deserialization. + let mut value = serde_json::json!({ "type": raw["type"] }); + match raw["type"].as_str().unwrap() { + "content-start" | "content-delta" => { + value["index"] = raw["index"].clone(); + value["delta"] = serde_json::json!({ "message": { "content": raw["delta"]["message"]["content"] } }); + } + "content-end" => value["index"] = raw["index"].clone(), + "message-start" => value["id"] = raw["id"].clone(), + "message-end" => { + value["delta"] = serde_json::json!({ "finish_reason": raw["delta"]["finish_reason"], "usage": raw["delta"]["usage"] }) + } + "tool-call-start" | "tool-call-delta" => { + value["delta"] = serde_json::json!({ "message": { "tool_calls": raw["delta"]["message"]["tool_calls"] } }) + } + "tool-plan-delta" => { + value["delta"] = serde_json::json!({ "message": { "tool_plan": raw["delta"]["message"]["tool_plan"] } }) + } + _ => {} + } + if let Ok(event) = serde_json::from_value(value) { + return Ok(event); + } + } + Err(aimux_core::error::AiMuxError::InvalidResponseData(format!( + "Invalid Cohere stream event: {raw}" + ))) + } } diff --git a/aimux-providers/src/deepseek/convert.rs b/aimux-providers/src/deepseek/convert.rs index 2ba10b45d..353bc875a 100644 --- a/aimux-providers/src/deepseek/convert.rs +++ b/aimux-providers/src/deepseek/convert.rs @@ -38,14 +38,19 @@ fn resolve_deepseek_image_media_type(part: &FilePart) -> Result Result { - let mut options = if matches!(part.data, FileData::Reference { .. }) { - Default::default() - } else { - parse_file_part_options(part.provider_options.as_ref(), provider_options_name)? - }; - if !allow_file_data { + if tool_result && let FileData::Reference { reference } = &part.data { + let file_id = reference.get("deepseek").ok_or_else(|| { + AiMuxError::InvalidArgument( + "No provider reference found for provider 'deepseek'.".into(), + ) + })?; + return Ok(json!({ "type": "file", "file_id": file_id })); + } + let mut options = + parse_file_part_options(part.provider_options.as_ref(), provider_options_name)?; + if tool_result { options.file_data = None; } let image_url = |url: &str| { @@ -80,15 +85,15 @@ fn convert_image_part( } FileData::Data { data } => { let media_type = resolve_deepseek_image_media_type(part)?; + let data = match data { + FileBytes::Binary(bytes) => base64::engine::general_purpose::STANDARD.encode(bytes), + FileBytes::Base64(data) => data.clone(), + }; let media_type = if media_type == "image/jpg" { "image/jpeg" } else { &media_type }; - let data = match data { - FileBytes::Binary(bytes) => base64::engine::general_purpose::STANDARD.encode(bytes), - FileBytes::Base64(data) => data.clone(), - }; let data_url = format!("data:{media_type};base64,{data}"); if options.file_data != Some(true) { return Ok(image_url(&data_url)); @@ -104,9 +109,7 @@ fn convert_image_part( } Ok(file) } - FileData::Text { .. } => Err(AiMuxError::UnsupportedFunctionality( - "text file parts".to_string(), - )), + FileData::Text { .. } => unreachable!("inline text is not an image input"), } } @@ -116,16 +119,6 @@ pub(crate) struct ConvertedMessages { pub warnings: Vec, } -fn join_text(content: &[UserPart]) -> String { - content - .iter() - .filter_map(|part| match part { - UserPart::Text(TextPart { text, .. }) => Some(text.as_str()), - _ => None, - }) - .collect() -} - /// The `messages` of the request body. /// /// # Errors @@ -209,23 +202,25 @@ pub(crate) fn convert_to_deepseek_chat_messages( } LanguageModelMessage::User { content: parts, .. } => { - let has_image_part = parts.iter().any(|part| { - matches!(part, - UserPart::File(file) if file.media_type.split('/').next() == Some("image") - && !matches!(file.data, FileData::Text { .. }) - ) - }); + let is_image = |file: &FilePart| { + file.media_type.split('/').next() == Some("image") + && !matches!(file.data, FileData::Text { .. }) + }; + let has_image_part = parts + .iter() + .any(|part| matches!(part, UserPart::File(file) if is_image(file))); let mut content = Vec::new(); + let mut text = String::new(); for part in parts { match part { - UserPart::Text(TextPart { text, .. }) => { - content.push(json!({ "type": "text", "text": text })); + UserPart::Text(TextPart { + text: part_text, .. + }) => { + text.push_str(part_text); + content.push(json!({ "type": "text", "text": part_text })); } - UserPart::File(file) - if file.media_type.split('/').next() == Some("image") - && !matches!(file.data, FileData::Text { .. }) => - { - content.push(convert_image_part(file, provider_options_name, true)?); + UserPart::File(file) if is_image(file) => { + content.push(convert_image_part(file, provider_options_name, false)?); } UserPart::File(_) => warnings.push(Warning::Unsupported { feature: "user message part type: file".to_string(), @@ -240,7 +235,7 @@ pub(crate) fn convert_to_deepseek_chat_messages( if has_image_part { Value::Array(content) } else { - json!(join_text(parts)) + json!(text) }, ); name(&mut wire); @@ -332,7 +327,7 @@ pub(crate) fn convert_to_deepseek_chat_messages( messages.push(json!({ "role": "tool", "tool_call_id": tool_call_id, - "content": convert_tool_result(output, provider_options_name, &mut warnings)?, + "content": convert_tool_result_content(output, provider_options_name, &mut warnings)?, })); } } @@ -342,7 +337,7 @@ pub(crate) fn convert_to_deepseek_chat_messages( Ok(ConvertedMessages { messages, warnings }) } -fn convert_tool_result( +fn convert_tool_result_content( output: &ToolResultOutput, provider_options_name: &str, warnings: &mut Vec, @@ -350,37 +345,35 @@ fn convert_tool_result( let ToolResultOutput::Content { value } = output else { return Ok(crate::openai::convert::tool_result_to_content(output)); }; - let is_image = |file: &FilePart| { - file.media_type.split('/').next() == Some("image") - && !matches!(file.data, FileData::Text { .. }) + let is_image = |part: &ToolResultContent| { + matches!(part, ToolResultContent::File(file) + if file.media_type.split('/').next() == Some("image") + && !matches!(file.data, FileData::Text { .. })) }; - if !value - .iter() - .any(|part| matches!(part, ToolResultContent::File(file) if is_image(file))) - { + if !value.iter().any(is_image) { return Ok(crate::openai::convert::tool_result_to_content(output)); } - let mut parts = Vec::new(); + let mut content = Vec::new(); for part in value { match part { ToolResultContent::Text(text) => { - parts.push(json!({ "type": "text", "text": text.text })) + content.push(json!({ "type": "text", "text": text.text })); } - ToolResultContent::File(file) if is_image(file) => { - parts.push(convert_image_part(file, provider_options_name, false)?); + ToolResultContent::File(file) if is_image(part) => { + content.push(convert_image_part(file, provider_options_name, true)?); + } + other => { + let kind = match other { + ToolResultContent::File(_) => "file", + ToolResultContent::Custom { .. } => "custom", + ToolResultContent::Text(_) => unreachable!(), + }; + warnings.push(Warning::Unsupported { + feature: format!("tool result content part type: {kind}"), + details: None, + }); } - _ => warnings.push(Warning::Unsupported { - feature: format!( - "tool result content part type: {}", - match part { - ToolResultContent::Text(_) => "text", - ToolResultContent::File(_) => "file", - ToolResultContent::Custom { .. } => "custom", - } - ), - details: None, - }), } } - Ok(Value::Array(parts)) + Ok(Value::Array(content)) } diff --git a/aimux-providers/src/deepseek/model.rs b/aimux-providers/src/deepseek/model.rs index fecea4fc6..29e9fbdbb 100644 --- a/aimux-providers/src/deepseek/model.rs +++ b/aimux-providers/src/deepseek/model.rs @@ -464,11 +464,10 @@ impl LanguageModel for DeepSeekChatLanguageModel { }, usage: convert_deepseek_usage(raw.get("usage")), warnings: built.warnings, - provider_metadata: Some( - provider_namespace(self.provider_options_name(), Value::Object(metadata)) - .expect("provider metadata must be an object"), - ), - request: Some(RequestInfo { body: Some(body) }), + provider_metadata: Some(provider_namespace( + self.provider_options_name(), + Value::Object(metadata), + )?), response: Some(ResponseInfo { id: data.id, timestamp: timestamp(data.created), @@ -476,6 +475,7 @@ impl LanguageModel for DeepSeekChatLanguageModel { headers: Some(response_headers), body: Some(raw), }), + request: Some(RequestInfo { body: Some(body) }), }) } @@ -497,7 +497,7 @@ impl LanguageModel for DeepSeekChatLanguageModel { ) .await?; let response_headers = resp.response_headers; - let mut sse_stream = resp.value; + let sse_stream = resp.value; let provider_options_name = self.provider_options_name().to_string(); let emit_raw_chunks = options.include_raw_chunks == Some(true); @@ -522,7 +522,8 @@ impl LanguageModel for DeepSeekChatLanguageModel { let mut content_logprobs: Vec = Vec::new(); let mut reasoning_logprobs: Vec = Vec::new(); - while let Some(event) = sse_stream.next().await { + let mut events = sse_stream; + while let Some(event) = events.next().await { let parsed = match event { Ok(parsed) => parsed, Err(error) => { @@ -759,7 +760,7 @@ impl LanguageModel for DeepSeekChatLanguageModel { yield Ok(StreamPart::Finish { finish_reason, usage: convert_deepseek_usage(usage.as_ref()), - provider_metadata: Some(provider_namespace(&provider_options_name, Value::Object(metadata)).expect("provider metadata must be an object")), + provider_metadata: Some(provider_namespace(&provider_options_name, Value::Object(metadata)).expect("metadata payload is an object")), }); }; diff --git a/aimux-providers/src/deepseek/prepare_tools.rs b/aimux-providers/src/deepseek/prepare_tools.rs index d5f35206f..5403a334c 100644 --- a/aimux-providers/src/deepseek/prepare_tools.rs +++ b/aimux-providers/src/deepseek/prepare_tools.rs @@ -80,7 +80,7 @@ pub(crate) fn prepare_tools( } } - let tool_choice = tool_choice.map(|tool_choice| match tool_choice { + let tool_choice = tool_choice.map(|choice| match choice { ToolChoice::Auto => json!("auto"), ToolChoice::None => json!("none"), ToolChoice::Required => json!("required"), diff --git a/aimux-providers/src/default_providers.rs b/aimux-providers/src/default_providers.rs index 1edebc13f..c0aac3205 100644 --- a/aimux-providers/src/default_providers.rs +++ b/aimux-providers/src/default_providers.rs @@ -124,6 +124,7 @@ pub fn create_provider( macro_rules! package { ($module:ident, $create:ident, $settings:ident) => { Ok(Arc::new(crate::$module::$create( + #[allow(clippy::needless_update)] crate::$module::$settings { api_key, base_url, @@ -186,7 +187,9 @@ pub fn create_provider( Default::default() ) } - "azure" => package!(azure, create_azure, AzureOpenAIProviderSettings), + "azure" => { + package_without_transform!(azure, create_azure, AzureOpenAIProviderSettings, api_key) + } "cohere" => package_with_string_key!( cohere, create_cohere, @@ -203,18 +206,22 @@ pub fn create_provider( ), "groq" => package_with_string_key!(groq, create_groq, GroqProviderSettings), "huggingface" => package!(huggingface, create_huggingface, HuggingFaceProviderSettings), - "mistral" => package_with_string_key!( - mistral, - create_mistral, - MistralProviderSettings, - Default::default() - ), - "openai" => package_with_string_key!( - openai, - create_openai, - OpenAIProviderSettings, - Default::default() - ), + "mistral" => { + package_with_string_key!( + mistral, + create_mistral, + MistralProviderSettings, + Default::default() + ) + } + "openai" => { + package_with_string_key!( + openai, + create_openai, + OpenAIProviderSettings, + Default::default() + ) + } "voyage" => package!(voyage, create_voyage, VoyageProviderSettings), "xai" => package_with_string_key!(xai, create_xai, XAIProviderSettings), "assemblyai" => package_without_transform!( diff --git a/aimux-providers/src/google/convert.rs b/aimux-providers/src/google/convert.rs index 43762c3b2..302f2c4c0 100644 --- a/aimux-providers/src/google/convert.rs +++ b/aimux-providers/src/google/convert.rs @@ -2,7 +2,7 @@ //! //! Mirrors the TS SDK's `convert-to-google-messages.ts` and the request-body //! construction inside `google-language-model.ts`'s `getArgs`. The Gemini -//! request shape is fundamentally different from OpenAI/Anthropic: +//! request shape lifts system content out of the message list: //! //! - System messages are lifted out of `contents` into a top-level //! `systemInstruction` field (a `{ parts: [{ text }] }` object). @@ -73,9 +73,8 @@ pub struct GooglePrompt { /// - Provider-executed calls and results use their provider metadata to replay /// native `toolCall` / `toolResponse` or `executableCode` / /// `codeExecutionResult` parts. -/// - Tool-result `output` is serialized into `functionResponse.response.content` -/// as a string (JSON-stringified for non-string outputs, matching the TS -/// `output.type === 'json'` path). +/// - Tool-result JSON stays structured, except values containing `$ref`, which +/// are stringified to avoid Google's multimodal reference handling. /// /// Thought signatures from provider options are echoed as a /// `thoughtSignature` sibling of the `functionCall` part. @@ -111,15 +110,13 @@ fn convert_to_google_messages_for_namespace( } LanguageModelMessage::User { content, .. } => { system_messages_allowed = false; - let parts = convert_user_parts(content, namespace); + let parts = convert_user_parts(content); contents.push(json!({ "role": "user", "parts": parts })); } LanguageModelMessage::Assistant { content, .. } => { system_messages_allowed = false; let parts = convert_assistant_parts(content, namespace); - if !parts.is_empty() { - contents.push(json!({ "role": "model", "parts": parts })); - } + contents.push(json!({ "role": "model", "parts": parts })); } LanguageModelMessage::Tool { content, .. } => { system_messages_allowed = false; @@ -129,34 +126,32 @@ fn convert_to_google_messages_for_namespace( // entry per tool-role message). let mut ordinary = Vec::new(); for part in content { - if let ToolPart::ToolResult(result) = part - && let Some(options) = namespace.read(result.provider_options.as_ref()) - && let (Some(server_id), Some(server_type)) = ( - options.get("serverToolCallId"), - options.get("serverToolType"), - ) + let ToolPart::ToolResult(ToolResultPart { + output, + provider_options, + .. + }) = part + else { + continue; + }; + if let Some(opts) = namespace.read_part(provider_options.as_ref()) + && let (Some(id), Some(kind)) = + (opts.get("serverToolCallId"), opts.get("serverToolType")) && let Some(last) = contents.last_mut() - && last.get("role").and_then(Value::as_str) == Some("model") - && let Some(parts) = last.get_mut("parts").and_then(Value::as_array_mut) + && last["role"] == "model" { - let response = match &result.output { - ToolResultOutput::Json { value, .. } => value.clone(), - _ => json!({}), - }; - let mut value = json!({ "toolResponse": { "toolType": server_type, "response": response, "id": server_id } }); - if let Some(signature) = options.get("thoughtSignature") { - value["thoughtSignature"] = signature.clone(); + let mut response = json!({ "toolResponse": { "toolType": kind, "response": match output { ToolResultOutput::Json { value, .. } => value.clone(), _ => json!({}) }, "id": id } }); + if let Some(signature) = opts.get("thoughtSignature") { + response["thoughtSignature"] = signature.clone(); } - parts.push(value); + last["parts"].as_array_mut().unwrap().push(response); } else { ordinary.push(part.clone()); } } let parts = convert_tool_parts(&ordinary, namespace, supports_function_response_parts); - if !parts.is_empty() { - contents.push(json!({ "role": "user", "parts": parts })); - } + contents.push(json!({ "role": "user", "parts": parts })); } } } @@ -174,65 +169,50 @@ fn convert_to_google_messages_for_namespace( } /// Convert user-role content parts into Google parts. -fn convert_user_parts(content: &[UserPart], namespace: Namespace) -> Vec { - let mut parts = Vec::new(); - for part in content { - match part { - UserPart::Text(TextPart { text, .. }) => { - parts.push(json!({ "text": text })); - } - UserPart::File(FilePart { - data, media_type, .. - }) => { - let inline_media_type = if matches!(data, FileData::Text { .. }) - && !aimux_provider_utils::is_full_media_type(media_type) - { - "text/plain" - } else { - media_type.as_str() - }; - let data = match data { - FileData::Data { - data: FileBytes::Binary(bytes), - } => base64::engine::general_purpose::STANDARD.encode(bytes), - FileData::Text { text } => { - base64::engine::general_purpose::STANDARD.encode(text.as_bytes()) - } - FileData::Data { - data: FileBytes::Base64(data), - } => data.clone(), - FileData::Url { url, original_url } => { - let uri = if url.starts_with("gs:") { - original_url.as_deref().unwrap_or(url) - } else { - url - }; - parts.push(json!({ - "fileData": { "mimeType": media_type, "fileUri": uri } - })); - continue; - } - FileData::Reference { reference } => { - if namespace == Namespace::Google - && let Some(reference) = reference.get(GOOGLE) - { - parts.push(json!({ - "fileData": { "mimeType": media_type, "fileUri": reference } - })); - } - continue; - } - }; - parts.push(json!({ - "inlineData": { "mimeType": inline_media_type, "data": data } - })); - } +fn convert_user_parts(content: &[UserPart]) -> Vec { + content + .iter() + .map(|part| match part { + UserPart::Text(TextPart { text, .. }) => json!({ "text": text }), + UserPart::File(file) => convert_file_part(file), + }) + .collect() +} + +fn convert_file_part(file: &FilePart) -> Value { + let media_type = if matches!(file.data, FileData::Text { .. }) { + if aimux_provider_utils::is_full_media_type(&file.media_type) { + file.media_type.clone() + } else { + "text/plain".into() } - } - if parts.is_empty() { - parts.push(json!({ "text": "" })); - } - parts + } else { + aimux_provider_utils::resolve_full_media_type(file) + .unwrap_or_else(|_| file.media_type.clone()) + }; + let data = match &file.data { + FileData::Data { + data: FileBytes::Binary(bytes), + } => base64::engine::general_purpose::STANDARD.encode(bytes), + FileData::Data { + data: FileBytes::Base64(data), + } => data.clone(), + FileData::Text { text } => { + base64::engine::general_purpose::STANDARD.encode(text.as_bytes()) + } + FileData::Url { url, original_url } => { + let uri = if url.starts_with("gs:") { + original_url.as_deref().unwrap_or(url) + } else { + url + }; + return json!({ "fileData": { "mimeType": media_type, "fileUri": uri } }); + } + FileData::Reference { reference } => { + return json!({ "fileData": { "mimeType": media_type, "fileUri": reference.get(GOOGLE) } }); + } + }; + json!({ "inlineData": { "mimeType": media_type, "data": data } }) } /// Convert assistant-role content parts into Google `model`-role parts. @@ -252,7 +232,7 @@ fn convert_assistant_parts(content: &[AssistantPart], namespace: Namespace) -> V // Echo thoughtSignature from provider_options if present // (upstream convert-to-google-messages.ts:355-377). if let Some(sig) = namespace - .read(provider_options.as_ref()) + .read_part(provider_options.as_ref()) .and_then(|g| g.get("thoughtSignature")) .and_then(|v| v.as_str()) { @@ -268,9 +248,9 @@ fn convert_assistant_parts(content: &[AssistantPart], namespace: Namespace) -> V if !text.is_empty() { let mut p = json!({ "text": text, "thought": true }); if let Some(sig) = namespace - .read(provider_options.as_ref()) + .read_part(provider_options.as_ref()) .and_then(|g| g.get("thoughtSignature")) - .and_then(Value::as_str) + .and_then(|v| v.as_str()) { p["thoughtSignature"] = json!(sig); } @@ -282,9 +262,18 @@ fn convert_assistant_parts(content: &[AssistantPart], namespace: Namespace) -> V tool_name, input, provider_options, - .. + provider_executed, }) => { - let google_options = namespace.read(provider_options.as_ref()); + if *provider_executed == Some(true) && tool_name == "code_execution" { + let input = if let Value::String(input) = input { + serde_json::from_str(input).unwrap_or(Value::Null) + } else { + input.clone() + }; + parts.push(json!({ "executableCode": input })); + continue; + } + let google_options = namespace.read_part(provider_options.as_ref()); let server_tool_call_id = google_options .and_then(|options| options.get("serverToolCallId")) .and_then(|value| value.as_str()); @@ -317,7 +306,7 @@ fn convert_assistant_parts(content: &[AssistantPart], namespace: Namespace) -> V } } else { let mut function_call = Map::new(); - if !tool_call_id.is_empty() && namespace != Namespace::Vertex { + if namespace == Namespace::Google && !tool_call_id.is_empty() { function_call.insert("id".to_string(), json!(tool_call_id)); } function_call.insert("name".to_string(), json!(tool_name)); @@ -330,8 +319,8 @@ fn convert_assistant_parts(content: &[AssistantPart], namespace: Namespace) -> V parts.push(part_value); } AssistantPart::ToolResult(ToolResultPart { - output, tool_name, + output, provider_options, .. }) => { @@ -348,7 +337,7 @@ fn convert_assistant_parts(content: &[AssistantPart], namespace: Namespace) -> V // upstream convert-to-google-messages.ts:518-540. // If it carries serverToolCallId + serverToolType, emit as // a toolResponse; otherwise skip (upstream returns undefined). - if let Some(opts) = namespace.read(provider_options.as_ref()) { + if let Some(opts) = namespace.read_part(provider_options.as_ref()) { let server_id = opts.get("serverToolCallId").and_then(|v| v.as_str()); let server_type = opts.get("serverToolType").and_then(|v| v.as_str()); if let (Some(sid), Some(st)) = (server_id, server_type) { @@ -383,7 +372,7 @@ fn convert_assistant_parts(content: &[AssistantPart], namespace: Namespace) -> V }; let mut value = json!({ "inlineData": { "mimeType": part.media_type, "data": data }, "thought": true }); if let Some(signature) = namespace - .read(part.provider_options.as_ref()) + .read_part(part.provider_options.as_ref()) .and_then(|options| options.get("thoughtSignature")) { value["thoughtSignature"] = signature.clone(); @@ -391,40 +380,19 @@ fn convert_assistant_parts(content: &[AssistantPart], namespace: Namespace) -> V parts.push(value); } } + AssistantPart::Custom(_) => {} AssistantPart::File(file) => { - let mut value = match &file.data { - FileData::Data { data } => { - let data = match data { - FileBytes::Binary(bytes) => { - base64::engine::general_purpose::STANDARD.encode(bytes) - } - FileBytes::Base64(data) => data.clone(), - }; - json!({ "inlineData": { "mimeType": file.media_type, "data": data } }) - } - FileData::Text { text } => { - json!({ "inlineData": { "mimeType": if aimux_provider_utils::is_full_media_type(&file.media_type) { file.media_type.as_str() } else { "text/plain" }, "data": base64::engine::general_purpose::STANDARD.encode(text.as_bytes()) } }) - } - FileData::Reference { reference } => { - if let Some(uri) = reference.get(GOOGLE) { - json!({ "fileData": { "mimeType": file.media_type, "fileUri": uri } }) - } else { - continue; - } - } - FileData::Url { .. } => continue, - }; - if let Some(options) = namespace.read(file.provider_options.as_ref()) { - if options.get("thought").and_then(Value::as_bool) == Some(true) { + let mut value = convert_file_part(file); + if let Some(opts) = namespace.read_part(file.provider_options.as_ref()) { + if opts.get("thought") == Some(&Value::Bool(true)) { value["thought"] = json!(true); } - if let Some(signature) = options.get("thoughtSignature") { + if let Some(signature) = opts.get("thoughtSignature") { value["thoughtSignature"] = signature.clone(); } } parts.push(value); } - AssistantPart::Custom(_) => {} } } parts @@ -434,47 +402,11 @@ fn convert_assistant_parts(content: &[AssistantPart], namespace: Namespace) -> V /// /// The TS SDK uses the `functionResponse` shape: /// `{ functionResponse: { id?, name, response: { name, content } } }`. -/// `content` is the tool's output serialized to a string (string outputs -/// pass through; JSON outputs are stringified). -pub(crate) fn tool_file_media_type( - file: &FilePart, -) -> Result { +/// `content` preserves the tool output unless it contains a JSON Schema `$ref`. +pub(crate) fn tool_file_media_type(file: &FilePart) -> Result { aimux_provider_utils::resolve_full_media_type(file) } -pub(crate) fn validate_tool_result_files( - prompt: &LanguageModelPrompt, -) -> Result<(), aimux_core::error::AiMuxError> { - for message in prompt { - if let LanguageModelMessage::Tool { content, .. } = message { - for part in content { - if let ToolPart::ToolResult(part) = part - && let ToolResultOutput::Content { value } = &part.output - { - for item in value { - if let ToolResultContent::File(file) = item - && matches!(file.data, FileData::Data { .. }) - { - tool_file_media_type(file)?; - } - } - } - } - } - } - Ok(()) -} - -fn contains_schema_reference(value: &Value) -> bool { - match value { - Value::Object(object) => { - object.contains_key("$ref") || object.values().any(contains_schema_reference) - } - Value::Array(array) => array.iter().any(contains_schema_reference), - _ => false, - } -} - fn convert_tool_parts( content: &[ToolPart], namespace: Namespace, @@ -520,6 +452,11 @@ fn convert_tool_parts( if let Some((media_type, data)) = url .strip_prefix("data:") .and_then(|url| url.split_once(";base64,")) + .filter(|(media_type, data)| { + !media_type.is_empty() + && !media_type.contains([';', ',']) + && !data.is_empty() + }) { files.push(json!({ "inlineData": { "mimeType": media_type, "data": data } })); } else { @@ -619,9 +556,8 @@ pub struct PreparedTools { /// Prepare `FunctionTool`s into the Google `tools` / `toolConfig` JSON shape. /// -/// Mirrors the function-tools path of the TS `prepareTools`. Provider-defined -/// tools (`google_search`, `code_execution`, …) are out of scope for the Rust -/// port — only `FunctionTool`s are supported. +/// Mirrors the function-tools path of the upstream `prepareTools`. +/// Use [`prepare_all_tools`] for provider-defined tools. #[must_use] pub fn prepare_tools( tools: &Option>, @@ -659,30 +595,10 @@ pub fn prepare_tools( } } Some(ToolChoice::None) => Some(json!({ "functionCallingConfig": { "mode": "NONE" } })), - Some(ToolChoice::Required) => { - if has_strict { - Some(json!({ "functionCallingConfig": { "mode": "VALIDATED" } })) - } else { - Some(json!({ "functionCallingConfig": { "mode": "ANY" } })) - } - } - Some(ToolChoice::Tool { tool_name }) => { - if has_strict { - Some(json!({ - "functionCallingConfig": { - "mode": "VALIDATED", - "allowedFunctionNames": [tool_name] - } - })) - } else { - Some(json!({ - "functionCallingConfig": { - "mode": "ANY", - "allowedFunctionNames": [tool_name] - } - })) - } - } + Some(ToolChoice::Required) => Some(json!({ "functionCallingConfig": { "mode": "ANY" } })), + Some(ToolChoice::Tool { tool_name }) => Some(json!({ + "functionCallingConfig": { "mode": "ANY", "allowedFunctionNames": [tool_name] } + })), }; PreparedTools { @@ -974,6 +890,20 @@ fn push_provider_tool( ))); } } + "google.vertex_rag_store" => { + if caps.supports_gemini_2_tools { + let mut rag = + json!({ "rag_resources": { "rag_corpus": tool.args.get("ragCorpus") } }); + if let Some(top_k) = tool.args.get("topK") { + rag["similarity_top_k"] = top_k.clone(); + } + google_tools.push(json!({ "retrieval": { "vertex_rag_store": rag } })); + } else { + warnings.push(unsupported(Some( + "The RAG store tool is not supported with other Gemini models than Gemini 2.", + ))); + } + } "google.google_maps" => { if caps.supports_gemini_2_tools { google_tools.push(json!({ "googleMaps": {} })); @@ -993,8 +923,7 @@ fn push_provider_tool( /// /// The tool's input schema goes out as `parametersJsonSchema`, unchanged, the /// way `@ai-sdk/google` sends it (Gemini accepts JSON Schema natively). An -/// empty root object schema (`{ type: "object", properties: {} }`) is -/// omitted. +/// empty root object schema is preserved. fn build_function_declaration(ft: &FunctionTool) -> Value { let mut decl = Map::new(); decl.insert("name".to_string(), json!(ft.name)); @@ -1002,223 +931,51 @@ fn build_function_declaration(ft: &FunctionTool) -> Value { "description".to_string(), json!(ft.description.as_deref().unwrap_or("")), ); - // The OpenAPI conversion is the emptiness test: it returns null for an - // empty root object schema. - if !convert_json_schema_to_openapi_schema(&ft.input_schema, true).is_null() { - decl.insert("parametersJsonSchema".to_string(), ft.input_schema.clone()); - } + decl.insert("parametersJsonSchema".to_string(), ft.input_schema.clone()); Value::Object(decl) } -// ── JSON Schema → OpenAPI Schema ───────────────────────────────────────────── - -/// Convert a JSON Schema value into the OpenAPI 3.0 schema variant that -/// Gemini's `functionDeclarations.parameters` expects. -/// -/// Mirrors `convert-json-schema-to-openapi-schema.ts` for the cases that -/// matter for tool parameter schemas: -/// - Empty object schemas become `undefined` at the root and -/// `{ type: "object" }` when nested. -/// - `$schema` / `additionalProperties` / `definitions` etc. are dropped. -/// - `type`, `description`, `required`, `properties`, `items`, `enum`, -/// `format`, `const`, `anyOf`/`oneOf`/`allOf` are preserved. -#[must_use] -pub fn convert_json_schema_to_openapi_schema(schema: &Value, is_root: bool) -> Value { - if schema.is_null() { - return Value::Null; - } - - // Boolean schema. - if let Some(b) = schema.as_bool() { - let _ = b; - return json!({ "type": "boolean", "properties": {} }); +fn contains_schema_reference(value: &Value) -> bool { + match value { + Value::Object(object) => object + .iter() + .any(|(key, value)| key == "$ref" || contains_schema_reference(value)), + Value::Array(array) => array.iter().any(contains_schema_reference), + _ => false, } +} - let obj = match schema.as_object() { - Some(o) => o, - None => return schema.clone(), +/// Preserve JSON Schema while replacing unsupported `const` constraints. +#[must_use] +pub fn sanitize_response_json_schema(schema: &Value) -> Value { + let Some(object) = schema.as_object() else { + return schema.clone(); }; - - // Empty object schema: `{ type: "object" }` with no/empty properties - // and no `additionalProperties`. - if is_empty_object_schema(obj) { - if is_root { - return Value::Null; - } - if obj.contains_key("description") { - return json!({ "type": "object", "description": obj["description"] }); - } - return json!({ "type": "object" }); - } - - let mut result = Map::new(); - - if let Some(desc) = obj.get("description") { - result.insert("description".to_string(), desc.clone()); - } - if let Some(req) = obj.get("required") { - result.insert("required".to_string(), req.clone()); - } - if let Some(format) = obj.get("format") { - result.insert("format".to_string(), format.clone()); - } - - if let Some(const_val) = obj.get("const") { - result.insert("enum".to_string(), json!([const_val])); + let mut result = object.clone(); + if let Some(value) = result.remove("const") { + result.insert("enum".into(), json!([value])); } - - // type: string | array - if let Some(type_val) = obj.get("type") { - if let Some(type_str) = type_val.as_str() { - result.insert("type".to_string(), json!(type_str)); - } else if let Some(types) = type_val.as_array() { - let has_null = types.iter().any(|t| t.as_str() == Some("null")); - let non_null: Vec<&Value> = types - .iter() - .filter(|t| t.as_str() != Some("null")) - .collect(); - if non_null.is_empty() { - // Only null type. - result.insert("type".to_string(), json!("null")); - } else { - // One or more non-null types: always use anyOf (matching TS). - let any_of: Vec = non_null.iter().map(|t| json!({ "type": t })).collect(); - result.insert("anyOf".to_string(), Value::Array(any_of)); - if has_null { - result.insert("nullable".to_string(), json!(true)); - } + for key in ["properties", "$defs"] { + if let Some(Value::Object(definitions)) = result.get_mut(key) { + for value in definitions.values_mut() { + *value = sanitize_response_json_schema(value); } } } - - if let Some(enum_vals) = obj.get("enum") { - result.insert("enum".to_string(), enum_vals.clone()); - } - - if let Some(props) = obj.get("properties").and_then(|p| p.as_object()) { - let mut out = Map::new(); - for (k, v) in props { - out.insert(k.clone(), convert_json_schema_to_openapi_schema(v, false)); - } - result.insert("properties".to_string(), Value::Object(out)); - } - - if let Some(items) = obj.get("items") { - if let Some(arr) = items.as_array() { - let converted: Vec = arr - .iter() - .map(|i| convert_json_schema_to_openapi_schema(i, false)) - .collect(); - result.insert("items".to_string(), Value::Array(converted)); - } else { - result.insert( - "items".to_string(), - convert_json_schema_to_openapi_schema(items, false), - ); - } - } - - // allOf / oneOf: recursively convert each element. - for combinator in ["allOf", "oneOf"] { - if let Some(arr) = obj.get(combinator).and_then(|v| v.as_array()) { - let converted: Vec = arr - .iter() - .map(|i| convert_json_schema_to_openapi_schema(i, false)) - .collect(); - result.insert(combinator.to_string(), Value::Array(converted)); - } - } - - // anyOf: collapse a `null`-typed branch into `nullable: true` (matching - // the TS SDK, which folds `anyOf: [T, {type:null}]` into `T + nullable`). - if let Some(arr) = obj.get("anyOf").and_then(|v| v.as_array()) { - let has_null = arr.iter().any(|s| { - s.as_object() - .and_then(|o| o.get("type")) - .and_then(|t| t.as_str()) - == Some("null") - }); - if has_null { - let non_null: Vec<&Value> = arr - .iter() - .filter(|s| { - s.as_object() - .and_then(|o| o.get("type")) - .and_then(|t| t.as_str()) - != Some("null") - }) - .collect(); - if non_null.len() == 1 { - // Single non-null schema: convert it and merge into result. - let converted = convert_json_schema_to_openapi_schema(non_null[0], false); - result.insert("nullable".to_string(), json!(true)); - if let Some(obj2) = converted.as_object() { - for (k, v) in obj2 { - result.insert(k.clone(), v.clone()); - } + for key in ["items", "additionalProperties", "anyOf", "oneOf"] { + if let Some(value) = result.get_mut(key) { + if let Value::Array(array) = value { + for item in array { + *item = sanitize_response_json_schema(item); } } else { - let converted: Vec = non_null - .iter() - .map(|i| convert_json_schema_to_openapi_schema(i, false)) - .collect(); - result.insert("anyOf".to_string(), Value::Array(converted)); - result.insert("nullable".to_string(), json!(true)); + *value = sanitize_response_json_schema(value); } - } else { - let converted: Vec = arr - .iter() - .map(|i| convert_json_schema_to_openapi_schema(i, false)) - .collect(); - result.insert("anyOf".to_string(), Value::Array(converted)); } } - - if let Some(min_len) = obj.get("minLength") { - result.insert("minLength".to_string(), min_len.clone()); - } - Value::Object(result) } -fn is_empty_object_schema(obj: &Map) -> bool { - obj.get("type").and_then(|v| v.as_str()) == Some("object") - && (obj.get("properties").is_none() - || obj - .get("properties") - .and_then(|p| p.as_object()) - .map(serde_json::Map::is_empty) - .unwrap_or(true)) - && !obj.contains_key("additionalProperties") -} - -// Image wrappers delegate these fields through the upstream language schema. -#[must_use] -pub(crate) fn image_generation_option(key: &str, value: &Value) -> Value { - let fields: &[&str] = match key { - "thinkingConfig" => &["thinkingBudget", "includeThoughts", "thinkingLevel"], - "imageConfig" => &[ - "aspectRatio", - "imageSize", - "personGeneration", - "prominentPeople", - "imageOutputOptions", - ], - "imageOutputOptions" => &["mimeType", "compressionQuality"], - _ => return value.clone(), - }; - let Some(object) = value.as_object() else { - return value.clone(); - }; - Value::Object( - object - .iter() - .filter(|(key, _)| fields.contains(&key.as_str())) - .map(|(key, value)| (key.clone(), image_generation_option(key, value))) - .collect(), - ) -} - // ── build_request_body ─────────────────────────────────────────────────────── /// Build the Gemini `generateContent` request body from `CallOptions`. @@ -1287,13 +1044,104 @@ fn build_request_body_with_warnings_for_namespace( } }; let GooglePrompt { - system_instruction, - contents, + mut system_instruction, + mut contents, } = convert_to_google_messages_for_namespace( &options.prompt, namespace, - model_id.contains("gemini-3"), + get_google_model_capabilities(model_id).uses_gemini_3_features, ); + let mut warnings = Vec::new(); + let provider_options = namespace.read_in(options.provider_options.as_ref()); + let option = |name: &str| { + provider_options + .and_then(|o| o.get(name)) + .filter(|v| !v.is_null()) + }; + let is_vertex = namespace == Namespace::Vertex; + if !is_vertex + && options.tools.as_ref().is_some_and(|tools| { + tools.iter().any( + |tool| matches!(tool, Tool::Provider(tool) if tool.id == "google.vertex_rag_store"), + ) + }) + { + warnings.push(Warning::Other { message: "The 'vertex_rag_store' tool is only supported with the Google Vertex provider and might not be supported or could behave unexpectedly with the current Google provider (google.generative-ai).".into() }); + } + + if !is_vertex && option("streamFunctionCallArguments") == Some(&Value::Bool(true)) { + warnings.push(Warning::Other { message: "'streamFunctionCallArguments' is only supported on the Vertex AI API and will be ignored with the current Google provider (google.generative-ai). See https://docs.cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling#streaming-fc".into() }); + } + if is_vertex && option("serviceTier").is_some() { + warnings.push(Warning::Other { message: "'serviceTier' is a Gemini API option and is not supported on Vertex AI. Use 'sharedRequestType' (and optionally 'requestType') instead. See https://docs.cloud.google.com/vertex-ai/generative-ai/docs/priority-paygo".into() }); + } + if !is_vertex && (option("sharedRequestType").is_some() || option("requestType").is_some()) { + warnings.push(Warning::Other { message: "'sharedRequestType' and 'requestType' are Vertex AI options and are ignored with the current Google provider (google.generative-ai).".into() }); + } + let option_warning_count = warnings.len(); + let mut prompt_warnings = Vec::new(); + let mut thinking_warnings = Vec::new(); + let omit_penalties = + !is_vertex && matches_prefix_boundary(&model_id.to_lowercase(), "gemini-2.5"); + if model_id.to_lowercase().starts_with("gemma-") + && let Some(system) = system_instruction.take() + && let Some(first) = contents.first_mut() + && first["role"] == "user" + && let Some(parts) = first["parts"].as_array_mut() + { + let texts = system["parts"] + .as_array() + .unwrap() + .iter() + .filter_map(|part| part["text"].as_str()) + .collect::>() + .join("\n\n"); + parts.insert(0, json!({ "text": format!("{texts}\n\n") })); + } + if get_google_model_capabilities(model_id).uses_gemini_3_features { + let mut missing = Vec::new(); + for content in &mut contents { + if content["role"] != "model" { + continue; + } + let mut has_signed_call = false; + for part in content["parts"].as_array_mut().unwrap() { + let standard = part.get("functionCall").is_some(); + let server = part.get("toolCall").is_some(); + if !standard && !server { + continue; + } + if part.get("thoughtSignature").is_some() { + if standard { + has_signed_call = true; + } + } else if server || !has_signed_call { + part["thoughtSignature"] = json!("skip_thought_signature_validator"); + missing.push( + part.get("functionCall") + .or_else(|| part.get("toolCall")) + .unwrap() + .get("name") + .or_else(|| part["toolCall"].get("toolType")) + .and_then(Value::as_str) + .unwrap_or("") + .to_string(), + ); + } + } + } + if !missing.is_empty() { + let mut names = Vec::new(); + for name in &missing { + if !names.contains(name) { + names.push(name.clone()); + } + } + prompt_warnings.push(Warning::Other { message: format!( + "Replayed {} `functionCall` part(s) for a Gemini 3 model without a `thoughtSignature` (tools: {}). Injected the documented `skip_thought_signature_validator` sentinel to keep the request from failing with HTTP 400. The likely cause is application code that drops `providerOptions.google.thoughtSignature` when persisting or serializing assistant tool-call messages. See https://ai.google.dev/gemini-api/docs/thought-signatures.", + missing.len(), names.iter().map(|name| format!("`{name}`")).collect::>().join(", ")) }); + } + } let mut generation_config = Map::new(); @@ -1309,11 +1157,20 @@ fn build_request_body_with_warnings_for_namespace( if let Some(top_k) = options.top_k { generation_config.insert("topK".to_string(), json!(top_k)); } - if let Some(presence) = options.presence_penalty { - generation_config.insert("presencePenalty".to_string(), json!(presence)); - } - if let Some(frequency) = options.frequency_penalty { - generation_config.insert("frequencyPenalty".to_string(), json!(frequency)); + for (name, value) in [ + ("frequencyPenalty", options.frequency_penalty), + ("presencePenalty", options.presence_penalty), + ] { + if let Some(value) = value { + if omit_penalties { + warnings.push(Warning::Unsupported { + feature: name.into(), + details: None, + }); + } else { + generation_config.insert(name.into(), json!(value)); + } + } } if let Some(stop) = &options.stop_sequences { generation_config.insert("stopSequences".to_string(), json!(stop)); @@ -1322,42 +1179,87 @@ fn build_request_body_with_warnings_for_namespace( generation_config.insert("seed".to_string(), json!(seed)); } - // Response format: Gemini uses `responseMimeType: "application/json"` - // (and an optional `responseSchema`) rather than OpenAI's `response_format`. + // JSON output uses responseMimeType and the sanitized responseJsonSchema. if let Some(rf) = &options.response_format && let ResponseFormat::Json { schema, .. } = rf { generation_config.insert("responseMimeType".to_string(), json!("application/json")); - if let Some(s) = schema + if let Some(schema) = schema && structured_outputs { - let openapi = convert_json_schema_to_openapi_schema(s, true); - if !openapi.is_null() { - generation_config.insert("responseSchema".to_string(), openapi); - } + generation_config.insert( + "responseJsonSchema".into(), + sanitize_response_json_schema(schema), + ); } } // Provider options (`providerOptions.google`, or `googleVertex` for // Vertex): the generation-config members, then the top-level ones. - let provider_options = namespace.read_in(options.provider_options.as_ref()); - let option = |name: &str| { - provider_options - .and_then(|o| o.get(name)) - .filter(|v| !v.is_null()) - }; - for name in [ - "responseModalities", - "thinkingConfig", - "audioTimestamp", - "mediaResolution", - "imageConfig", + for name in ["responseModalities", "audioTimestamp", "mediaResolution"] { + if let Some(value) = option(name) + && (name != "audioTimestamp" || value == &Value::Bool(true)) + { + generation_config.insert(name.into(), value.clone()); + } + } + let mut thinking = resolve_thinking(options.reasoning, model_id, &mut thinking_warnings); + if let Some(explicit) = option("thinkingConfig").and_then(Value::as_object) { + thinking + .get_or_insert_with(Map::new) + .extend(explicit.clone()); + } + if let Some(thinking) = thinking { + generation_config.insert("thinkingConfig".into(), Value::Object(thinking)); + } + if let Some(image) = option("imageConfig") { + let mut image = image.clone(); + if !is_vertex && let Some(config) = image.as_object_mut() { + let mut dropped = Vec::new(); + for name in ["personGeneration", "prominentPeople", "imageOutputOptions"] { + if config.remove(name).is_some() { + dropped.push(format!("'imageConfig.{name}'")); + } + } + if !dropped.is_empty() { + warnings.insert(option_warning_count, Warning::Other { + message: format!( + "{} {} ignored with the current Google provider (google.generative-ai).", + dropped.join(", "), + if dropped.len() == 1 { + "is a Vertex AI option and is" + } else { + "are Vertex AI options and are" + } + ), + }); + } + } + generation_config.insert("imageConfig".into(), image); + } + for (name, keys) in [ + ( + "thinkingConfig", + &["thinkingBudget", "includeThoughts", "thinkingLevel"][..], + ), + ( + "imageConfig", + &[ + "aspectRatio", + "imageSize", + "personGeneration", + "prominentPeople", + "imageOutputOptions", + ][..], + ), ] { - if let Some(value) = option(name) { - generation_config.insert(name.to_string(), value.clone()); + if let Some(Value::Object(config)) = generation_config.get_mut(name) { + config.retain(|key, _| keys.contains(&key.as_str())); + if let Some(Value::Object(output)) = config.get_mut("imageOutputOptions") { + output.retain(|key, _| ["mimeType", "compressionQuality"].contains(&key.as_str())); + } } } - let mut body = Map::new(); body.insert("contents".to_string(), Value::Array(contents)); if let Some(sys) = system_instruction { @@ -1369,22 +1271,57 @@ fn build_request_body_with_warnings_for_namespace( Value::Object(generation_config), ); for name in ["safetySettings", "cachedContent", "labels", "serviceTier"] { - if let Some(value) = option(name) { + if let Some(value) = option(name) + && !(is_vertex && name == "serviceTier") + { body.insert(name.to_string(), value.clone()); } } + if option("safetySettings").is_none() + && let Some(threshold) = option("threshold") + { + body.insert( + "safetySettings".into(), + json!( + [ + "HARM_CATEGORY_HATE_SPEECH", + "HARM_CATEGORY_DANGEROUS_CONTENT", + "HARM_CATEGORY_HARASSMENT", + "HARM_CATEGORY_SEXUALLY_EXPLICIT" + ] + .map(|category| json!({ "category": category, "threshold": threshold })) + ), + ); + } + if let Some(Value::Array(settings)) = body.get_mut("safetySettings") { + for setting in settings { + if let Some(setting) = setting.as_object_mut() { + setting.retain(|key, _| ["category", "threshold"].contains(&key.as_str())); + } + } + } let prepared = prepare_all_tools(&options.tools, options.tool_choice.as_ref(), model_id); if let Some(tools) = prepared.tools { body.insert("tools".to_string(), Value::Array(tools)); } let mut tool_config = prepared.tool_config; + if is_vertex && let Some(Value::Object(config)) = tool_config.as_mut() { + config.remove("includeServerSideToolInvocations"); + } if let Some(retrieval) = option("retrievalConfig") { let mut config = match tool_config.take() { Some(Value::Object(config)) => config, _ => Map::new(), }; - config.insert("retrievalConfig".to_string(), retrieval.clone()); + let mut retrieval = retrieval.clone(); + if let Some(retrieval) = retrieval.as_object_mut() { + retrieval.retain(|key, _| key == "latLng"); + if let Some(Value::Object(coordinates)) = retrieval.get_mut("latLng") { + coordinates.retain(|key, _| ["latitude", "longitude"].contains(&key.as_str())); + } + } + config.insert("retrievalConfig".to_string(), retrieval); tool_config = Some(Value::Object(config)); } if let Some(tc) = tool_config { @@ -1396,7 +1333,361 @@ fn build_request_body_with_warnings_for_namespace( // the TS SDK. (Some callers include it; the API ignores extra fields.) let _ = model_id; - Ok((Value::Object(body), prepared.warnings)) + warnings.extend(prompt_warnings); + warnings.extend(thinking_warnings); + warnings.extend(prepared.warnings); + Ok((Value::Object(body), warnings)) +} + +/// Validate prompt constraints and provider options before sending a request. +pub(crate) fn validate_call_options( + options: &CallOptions, +) -> Result<(), aimux_core::error::AiMuxError> { + validate_call_options_for_namespace(options, Namespace::Google) +} + +pub(crate) fn validate_call_options_for_namespace( + options: &CallOptions, + namespace: Namespace, +) -> Result<(), aimux_core::error::AiMuxError> { + use aimux_core::error::AiMuxError; + let mut initial = true; + for message in &options.prompt { + if matches!(message, LanguageModelMessage::System { .. }) { + if !initial { + return Err(AiMuxError::UnsupportedFunctionality( + "system messages are only supported at the beginning of the conversation" + .into(), + )); + } + } else { + initial = false; + } + match message { + LanguageModelMessage::User { content, .. } => { + for part in content { + if let UserPart::File(file) = part { + if namespace == Namespace::Vertex + && matches!(file.data, FileData::Reference { .. }) + { + return Err(AiMuxError::UnsupportedFunctionality( + "file parts with provider references".into(), + )); + } + validate_file_part(file, false)?; + } + } + } + LanguageModelMessage::Assistant { content, .. } => { + for part in content { + match part { + AssistantPart::File(file) => { + if namespace == Namespace::Vertex + && matches!(file.data, FileData::Reference { .. }) + { + return Err(AiMuxError::UnsupportedFunctionality( + "file parts with provider references".into(), + )); + } + validate_file_part(file, true)?; + } + AssistantPart::ReasoningFile(file) + if matches!( + file.data, + aimux_core::shared::GeneratedFileData::Url { .. } + ) => + { + return Err(AiMuxError::UnsupportedFunctionality( + "File data URLs in assistant messages are not supported".into(), + )); + } + _ => {} + } + } + } + LanguageModelMessage::Tool { content, .. } => { + for part in content { + if let ToolPart::ToolResult(part) = part + && let ToolResultOutput::Content { value } = &part.output + { + for part in value { + if let ToolResultContent::File(file) = part + && matches!(file.data, FileData::Data { .. }) + { + aimux_provider_utils::resolve_full_media_type(file)?; + } + } + } + } + } + _ => {} + } + } + let Some(values) = Namespace::Google.read_in(options.provider_options.as_ref()) else { + return Ok(()); + }; + let invalid = + |name: &str| AiMuxError::InvalidArgument(format!("Invalid Google provider option: {name}")); + for name in [ + "structuredOutputs", + "audioTimestamp", + "streamFunctionCallArguments", + ] { + if values.get(name).is_some_and(|value| !value.is_boolean()) { + return Err(invalid(name)); + } + } + if values + .get("cachedContent") + .is_some_and(|value| !value.is_string()) + { + return Err(invalid("cachedContent")); + } + for (name, allowed) in [ + ( + "threshold", + &[ + "HARM_BLOCK_THRESHOLD_UNSPECIFIED", + "BLOCK_LOW_AND_ABOVE", + "BLOCK_MEDIUM_AND_ABOVE", + "BLOCK_ONLY_HIGH", + "BLOCK_NONE", + "OFF", + ][..], + ), + ( + "mediaResolution", + &[ + "MEDIA_RESOLUTION_UNSPECIFIED", + "MEDIA_RESOLUTION_LOW", + "MEDIA_RESOLUTION_MEDIUM", + "MEDIA_RESOLUTION_HIGH", + ][..], + ), + ("serviceTier", &["standard", "flex", "priority"][..]), + ("sharedRequestType", &["priority", "flex", "standard"][..]), + ("requestType", &["shared"][..]), + ] { + if values + .get(name) + .is_some_and(|value| !value.as_str().is_some_and(|value| allowed.contains(&value))) + { + return Err(invalid(name)); + } + } + if values.get("responseModalities").is_some_and(|value| { + !value.as_array().is_some_and(|items| { + items + .iter() + .all(|item| matches!(item.as_str(), Some("TEXT" | "IMAGE"))) + }) + }) { + return Err(invalid("responseModalities")); + } + if let Some(value) = values.get("thinkingConfig") { + let config = value.as_object().ok_or_else(|| invalid("thinkingConfig"))?; + if config + .get("thinkingBudget") + .is_some_and(|value| !value.is_number()) + || config + .get("includeThoughts") + .is_some_and(|value| !value.is_boolean()) + || config.get("thinkingLevel").is_some_and(|value| { + !matches!(value.as_str(), Some("minimal" | "low" | "medium" | "high")) + }) + { + return Err(invalid("thinkingConfig")); + } + } + if let Some(value) = values.get("labels") + && !value + .as_object() + .is_some_and(|items| items.values().all(Value::is_string)) + { + return Err(invalid("labels")); + } + if let Some(value) = values.get("safetySettings") + && !value.as_array().is_some_and(|items| { + items.iter().all(|item| { + matches!( + item["category"].as_str(), + Some( + "HARM_CATEGORY_UNSPECIFIED" + | "HARM_CATEGORY_HATE_SPEECH" + | "HARM_CATEGORY_DANGEROUS_CONTENT" + | "HARM_CATEGORY_HARASSMENT" + | "HARM_CATEGORY_SEXUALLY_EXPLICIT" + | "HARM_CATEGORY_CIVIC_INTEGRITY" + ) + ) && matches!( + item["threshold"].as_str(), + Some( + "HARM_BLOCK_THRESHOLD_UNSPECIFIED" + | "BLOCK_LOW_AND_ABOVE" + | "BLOCK_MEDIUM_AND_ABOVE" + | "BLOCK_ONLY_HIGH" + | "BLOCK_NONE" + | "OFF" + ) + ) + }) + }) + { + return Err(invalid("safetySettings")); + } + if let Some(value) = values.get("imageConfig") { + let config = value.as_object().ok_or_else(|| invalid("imageConfig"))?; + for (name, allowed) in [ + ( + "aspectRatio", + &[ + "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9", "1:8", + "8:1", "1:4", "4:1", + ][..], + ), + ("imageSize", &["1K", "2K", "4K", "512"][..]), + ( + "personGeneration", + &[ + "PERSON_GENERATION_UNSPECIFIED", + "ALLOW_ALL", + "ALLOW_ADULT", + "ALLOW_NONE", + ][..], + ), + ( + "prominentPeople", + &[ + "PROMINENT_PEOPLE_UNSPECIFIED", + "ALLOW_PROMINENT_PEOPLE", + "BLOCK_PROMINENT_PEOPLE", + ][..], + ), + ] { + if config + .get(name) + .is_some_and(|value| !value.as_str().is_some_and(|value| allowed.contains(&value))) + { + return Err(invalid(name)); + } + } + if let Some(value) = config.get("imageOutputOptions") { + let output = value + .as_object() + .ok_or_else(|| invalid("imageOutputOptions"))?; + if output + .get("mimeType") + .is_some_and(|value| !matches!(value.as_str(), Some("image/jpeg" | "image/png"))) + || output + .get("compressionQuality") + .is_some_and(|value| !value.is_number()) + { + return Err(invalid("imageOutputOptions")); + } + } + } + if let Some(value) = values.get("retrievalConfig") { + let config = value + .as_object() + .ok_or_else(|| invalid("retrievalConfig"))?; + if let Some(value) = config.get("latLng") + && !value.as_object().is_some_and(|coordinates| { + coordinates.get("latitude").is_some_and(Value::is_number) + && coordinates.get("longitude").is_some_and(Value::is_number) + }) + { + return Err(invalid("retrievalConfig.latLng")); + } + } + Ok(()) +} + +fn validate_file_part( + file: &FilePart, + assistant: bool, +) -> Result<(), aimux_core::error::AiMuxError> { + use aimux_core::error::AiMuxError; + if !matches!(file.data, FileData::Text { .. }) { + aimux_provider_utils::resolve_full_media_type(file)?; + } + if assistant && matches!(file.data, FileData::Url { .. }) { + return Err(AiMuxError::UnsupportedFunctionality( + "File data URLs in assistant messages are not supported".into(), + )); + } + if let FileData::Reference { reference } = &file.data + && !reference.contains_key(GOOGLE) + { + return Err(AiMuxError::InvalidArgument(format!( + "No provider reference found for provider 'google'. Available providers: {}", + reference.keys().cloned().collect::>().join(", ") + ))); + } + Ok(()) +} + +fn resolve_thinking( + reasoning: Option, + model_id: &str, + warnings: &mut Vec, +) -> Option> { + use aimux_core::types::ReasoningEffort as R; + let reasoning = reasoning.filter(|effort| *effort != R::ProviderDefault)?; + let lower = model_id.to_lowercase(); + if get_google_model_capabilities(model_id).uses_gemini_3_features + && !model_id.contains("gemini-3-pro-image") + { + let name = lower.rsplit('/').next().unwrap_or(&lower); + let minimum = if name == "gemini-flash-latest" + || name + .strip_prefix("gemini-") + .and_then(|s| s.split_once("-flash")) + .is_some_and(|(version, tail)| { + let Some((major, minor)) = version.split_once('.') else { + return false; + }; + let major = major.parse::().unwrap_or(0); + let minor = minor.parse::().unwrap_or(0); + (tail.is_empty() + || (tail.starts_with('-') + && tail != "-lite" + && !tail.starts_with("-lite-"))) + && (major > 3 || major == 3 && minor >= 7) + }) { + "low" + } else { + "minimal" + }; + let level = match reasoning { + R::None | R::Minimal => minimum, + R::Low => "low", + R::Medium => "medium", + R::High | R::Xhigh => "high", + R::ProviderDefault => unreachable!(), + }; + if reasoning != R::None && level != reasoning.to_string() { + warnings.push(Warning::Compatibility { feature: "reasoning".into(), details: Some(format!("reasoning \"{reasoning}\" is not directly supported by this model. mapped to effort \"{level}\".")) }); + } + return Some(Map::from_iter([("thinkingLevel".into(), json!(level))])); + } + let percentage = match reasoning { + R::None => 0.0, + R::Minimal => 0.02, + R::Low => 0.1, + R::Medium => 0.3, + R::High => 0.6, + R::Xhigh => 0.9, + R::ProviderDefault => unreachable!(), + }; + let max = if lower.contains("2.5-pro") || lower.contains("gemini-3-pro-image") { + 32768 + } else { + 24576 + }; + Some(Map::from_iter([( + "thinkingBudget".into(), + json!(((65536.0_f64 * percentage).round() as u32).min(max)), + )])) } // ── finish reason ──────────────────────────────────────────────────────────── @@ -1433,15 +1724,16 @@ pub fn parse_finish_reason(reason: &str, has_tool_calls: bool) -> FinishReason { /// Convert a `GoogleUsageMetadata` into the core `Usage` type. /// /// Mirrors `convertGoogleUsage`: -/// - `input.total = promptTokenCount` -/// - `input.noCache = promptTokenCount - cachedContentTokenCount` +/// - `input.total = promptTokenCount + toolUsePromptTokenCount` +/// - `input.noCache = input.total - cachedContentTokenCount` /// - `input.cacheRead = cachedContentTokenCount` /// - `output.total = candidatesTokenCount + thoughtsTokenCount` #[must_use] pub fn convert_usage(usage: &super::types::GoogleUsageMetadata) -> aimux_core::types::Usage { use aimux_core::types::Usage; - let prompt = usage.prompt_token_count.unwrap_or(0); + let prompt = + usage.prompt_token_count.unwrap_or(0) + usage.tool_use_prompt_token_count.unwrap_or(0); let candidates = usage.candidates_token_count.unwrap_or(0); let cached = usage.cached_content_token_count.unwrap_or(0); let thoughts = usage.thoughts_token_count.unwrap_or(0); diff --git a/aimux-providers/src/google/embedding.rs b/aimux-providers/src/google/embedding.rs index 733842289..9384c14ec 100644 --- a/aimux-providers/src/google/embedding.rs +++ b/aimux-providers/src/google/embedding.rs @@ -14,7 +14,6 @@ use aimux_core::embedding_model::{ EmbeddingCallOptions, EmbeddingModel, EmbeddingResponse, EmbeddingResult, }; use aimux_core::error::AiMuxError; -use aimux_core::shared::SharedProviderOptions; use super::options::google_options; use crate::shared::EndpointConfig; @@ -56,173 +55,181 @@ impl EmbeddingModel for GoogleEmbeddingModel { &self, options: &EmbeddingCallOptions, ) -> Result { - let google_options = parse_google_provider_options(options.provider_options.as_ref()); - - let exchange = self.config.exchange(options.headers.as_ref()).await?; - - // For single embeddings, use the single endpoint. - if options.values.len() == 1 { - let value = &options.values[0]; - let mut parts = Map::new(); - parts.insert("text".to_string(), json!(value)); - - let mut content = Map::new(); - content.insert( - "parts".to_string(), - Value::Array(vec![Value::Object(parts)]), - ); - - let mut body = Map::new(); - body.insert( - "model".to_string(), - json!(format!("models/{}", self.model_id)), - ); - body.insert("content".to_string(), Value::Object(content)); - if let Some(dim) = google_options.output_dimensionality { - body.insert("outputDimensionality".to_string(), json!(dim)); + let provider_options = google_options(options.provider_options.as_ref()); + let mut settings = Map::new(); + if let Some(opts) = provider_options { + if let Some(dimension) = opts.get("outputDimensionality") { + if !dimension.is_number() { + return Err(AiMuxError::InvalidArgument( + "outputDimensionality must be a number".into(), + )); + } + settings.insert("outputDimensionality".into(), dimension.clone()); } - if let Some(task_type) = google_options.task_type { - body.insert("taskType".to_string(), json!(task_type)); + if let Some(task) = opts.get("taskType") { + if !matches!( + task.as_str(), + Some( + "SEMANTIC_SIMILARITY" + | "CLASSIFICATION" + | "CLUSTERING" + | "RETRIEVAL_DOCUMENT" + | "RETRIEVAL_QUERY" + | "QUESTION_ANSWERING" + | "FACT_VERIFICATION" + | "CODE_RETRIEVAL_QUERY" + ) + ) { + return Err(AiMuxError::InvalidArgument( + "Invalid Google embedding taskType".into(), + )); + } + settings.insert("taskType".into(), task.clone()); } - - let url = exchange.url(&format!("/models/{}:embedContent", self.model_id)); - - let resp = aimux_provider_utils::post_json_to_api( - exchange.request(url, options), - exchange.transform_body(Value::Object(body)), - aimux_provider_utils::create_json_response_handler(), - super::google_failed_response_handler(), - ) - .await?; - - let response_headers = resp.response_headers; - - let raw_value: Value = resp.value; - - // Single embedding: response.embedding.values - let embedding = raw_value - .get("embedding") - .and_then(|e| e.get("values")) - .and_then(|v| v.as_array()) - .map(|vals| { - vals.iter() - .filter_map(|v| v.as_f64().map(|f| f as f32)) - .collect() + } + if options.values.len() > 100 { + return Err(AiMuxError::InvalidArgument(format!( + "Too many values for a single embedding call. The {} model \"{}\" can only embed up to 100 values per call, but {} values were provided.", + self.provider(), + self.model_id, + options.values.len() + ))); + } + let multimodal = provider_options.and_then(|o| o.get("content")); + let multimodal = multimodal + .map(|v| { + v.as_array().ok_or_else(|| { + AiMuxError::InvalidArgument("Google embedding content must be an array".into()) }) - .unwrap_or_default(); - - return Ok(EmbeddingResult { - embeddings: vec![embedding], - usage: None, - provider_metadata: None, - response: Some(EmbeddingResponse { - headers: Some(response_headers), - body: Some(raw_value), - }), - warnings: Vec::new(), - }); + }) + .transpose()?; + if let Some(content) = multimodal { + if content.len() != options.values.len() { + return Err(AiMuxError::InvalidArgument(format!( + "The number of multimodal content entries ({}) must match the number of values ({}).", + content.len(), + options.values.len() + ))); + } + for entry in content.iter().filter(|v| !v.is_null()) { + let parts = entry.as_array().filter(|p| !p.is_empty()).ok_or_else(|| { + AiMuxError::InvalidArgument( + "Google embedding content entries must be nonempty arrays or null".into(), + ) + })?; + for part in parts { + let valid = part.get("text").is_some_and(Value::is_string) + || part.get("inlineData").is_some_and(|p| { + p.get("mimeType").is_some_and(Value::is_string) + && p.get("data").is_some_and(Value::is_string) + }) + || part.get("fileData").is_some_and(|p| { + p.get("mimeType").is_some_and(Value::is_string) + && p.get("fileUri").is_some_and(Value::is_string) + }); + if !valid { + return Err(AiMuxError::InvalidArgument( + "Invalid Google embedding content part".into(), + )); + } + } + } } - - // For multiple values, use the batch endpoint. + let single = options.values.len() == 1; let requests: Vec = options .values .iter() - .map(|value| { - let mut parts = Map::new(); - parts.insert("text".to_string(), json!(value)); - - let mut content = Map::new(); - content.insert("role".to_string(), json!("user")); - content.insert( - "parts".to_string(), - Value::Array(vec![Value::Object(parts)]), - ); - - let mut req = Map::new(); - req.insert( - "model".to_string(), - json!(format!("models/{}", self.model_id)), - ); - req.insert("content".to_string(), Value::Object(content)); - if let Some(dim) = google_options.output_dimensionality { - req.insert("outputDimensionality".to_string(), json!(dim)); + .enumerate() + .map(|(index, value)| { + let extras = multimodal.and_then(|entries| entries[index].as_array()); + let mut parts = Vec::new(); + if extras.is_none() || !value.is_empty() { + parts.push(json!({"text": value})); } - if let Some(task_type) = google_options.task_type.clone() { - req.insert("taskType".to_string(), json!(task_type)); + if let Some(extras) = extras { + parts.extend(extras.iter().map(|part| { + if let Some(text) = part.get("text").and_then(Value::as_str) { json!({"text": text}) } + else if let Some(inline) = part.get("inlineData").filter(|v| v.get("mimeType").is_some_and(Value::is_string) && v.get("data").is_some_and(Value::is_string)) { + json!({"inlineData": {"mimeType": inline["mimeType"], "data": inline["data"]}}) + } else { json!({"fileData": {"mimeType": part["fileData"]["mimeType"], "fileUri": part["fileData"]["fileUri"]}}) } + })); } - Value::Object(req) + let mut content = json!({"parts": parts}); + if !single { + content["role"] = json!("user"); + } + let mut request = settings.clone(); + request.insert("model".into(), json!(format!("models/{}", self.model_id))); + request.insert("content".into(), content); + Value::Object(request) }) .collect(); - - let mut body = Map::new(); - body.insert("requests".to_string(), Value::Array(requests)); - - let url = exchange.url(&format!("/models/{}:batchEmbedContents", self.model_id)); - + let body = if single { + requests[0].clone() + } else { + json!({"requests": requests}) + }; + let endpoint = if single { + "embedContent" + } else { + "batchEmbedContents" + }; + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - exchange.request(url, options), - exchange.transform_body(Value::Object(body)), - aimux_provider_utils::create_json_response_handler(), + exchange.request( + exchange.url(&format!("/models/{}:{endpoint}", self.model_id)), + options, + ), + exchange.transform_body(body), + aimux_provider_utils::create_json_response_handler::(), super::google_failed_response_handler(), ) .await?; - - let response_headers = resp.response_headers; - - let raw_value: Value = resp.value; - - // Batch embeddings: response.embeddings[].values - let embeddings: Vec> = raw_value - .get("embeddings") - .and_then(|e| e.as_array()) - .map(|arr| { - arr.iter() - .map(|item| { - item.get("values") - .and_then(|v| v.as_array()) - .map(|vals| { - vals.iter() - .filter_map(|v| v.as_f64().map(|f| f as f32)) - .collect() - }) - .unwrap_or_default() + let raw_value = resp.value; + let items = if single { + vec![ + raw_value + .get("embedding") + .ok_or_else(|| AiMuxError::InvalidResponseData("Missing embedding".into()))?, + ] + } else { + raw_value + .get("embeddings") + .and_then(Value::as_array) + .ok_or_else(|| AiMuxError::InvalidResponseData("Missing embeddings".into()))? + .iter() + .collect() + }; + let embeddings = items + .into_iter() + .map(|item| { + let values = item + .get("values") + .and_then(Value::as_array) + .ok_or_else(|| { + AiMuxError::InvalidResponseData("Missing embedding values".into()) + })?; + values + .iter() + .map(|v| { + v.as_f64().map(|v| v as f32).ok_or_else(|| { + AiMuxError::InvalidResponseData( + "Embedding values must be numbers".into(), + ) + }) }) - .collect() + .collect::, _>>() }) - .unwrap_or_default(); - + .collect::, _>>()?; Ok(EmbeddingResult { embeddings, usage: None, provider_metadata: None, response: Some(EmbeddingResponse { - headers: Some(response_headers), + headers: Some(resp.response_headers), body: Some(raw_value), }), warnings: Vec::new(), }) } } - -// ── Provider options parsing ───────────────────────────────────────────────── - -struct GoogleEmbeddingProviderOptions { - output_dimensionality: Option, - task_type: Option, -} - -fn parse_google_provider_options( - options: Option<&SharedProviderOptions>, -) -> GoogleEmbeddingProviderOptions { - let provider_opts = google_options(options); - GoogleEmbeddingProviderOptions { - output_dimensionality: provider_opts - .and_then(|o| o.get("outputDimensionality")) - .and_then(serde_json::Value::as_u64) - .map(|d| d as u32), - task_type: provider_opts - .and_then(|o| o.get("taskType")) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - } -} diff --git a/aimux-providers/src/google/files.rs b/aimux-providers/src/google/files.rs index 783776217..d30a17004 100644 --- a/aimux-providers/src/google/files.rs +++ b/aimux-providers/src/google/files.rs @@ -32,38 +32,32 @@ use crate::shared::EndpointConfig; /// Google-specific error structure: `{ "error": { "message": "..." } }`. /// Google provider-specific file upload options. -#[derive(Debug, Clone, Default)] +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(rename_all = "camelCase")] struct GoogleFilesUploadOptions { display_name: Option, - poll_interval_ms: Option, - poll_timeout_ms: Option, + poll_interval_ms: Option, + poll_timeout_ms: Option, } -/// Parse provider options for the Google files provider. fn parse_google_files_options( provider_options: Option<&SharedProviderOptions>, -) -> GoogleFilesUploadOptions { - let mut opts = GoogleFilesUploadOptions::default(); - if let Some(po) = provider_options - && let Some(google) = po.get(GOOGLE) +) -> Result { + let Some(google) = provider_options.and_then(|options| options.get(GOOGLE)) else { + return Ok(GoogleFilesUploadOptions::default()); + }; + let options: GoogleFilesUploadOptions = serde_json::from_value(Value::Object(google.clone())) + .map_err(|error| { + AiMuxError::InvalidArgument(format!("Invalid Google file options: {error}")) + })?; + if options.poll_interval_ms.is_some_and(|value| value <= 0.0) + || options.poll_timeout_ms.is_some_and(|value| value <= 0.0) { - if let Some(display_name) = google.get("displayName").and_then(|v| v.as_str()) { - opts.display_name = Some(display_name.to_string()); - } - if let Some(interval) = google - .get("pollIntervalMs") - .and_then(serde_json::Value::as_u64) - { - opts.poll_interval_ms = Some(interval); - } - if let Some(timeout) = google - .get("pollTimeoutMs") - .and_then(serde_json::Value::as_u64) - { - opts.poll_timeout_ms = Some(timeout); - } + return Err(AiMuxError::InvalidArgument( + "Google file polling intervals must be positive".into(), + )); } - opts + Ok(options) } /// Convert `UploadFileData` to raw bytes. @@ -80,6 +74,35 @@ fn data_to_bytes(data: &UploadFileData) -> Result, AiMuxError> { } } +// Preserve file names in a single path segment, including dot segments. +fn file_poll_path(name: &str) -> String { + fn encode(segment: &str) -> String { + if segment == "." { + return "%252E".to_string(); + } + if segment == ".." { + return "%252E%252E".to_string(); + } + const COMPONENT: &percent_encoding::AsciiSet = &percent_encoding::NON_ALPHANUMERIC + .remove(b'-') + .remove(b'_') + .remove(b'.') + .remove(b'!') + .remove(b'~') + .remove(b'*') + .remove(b'\'') + .remove(b'(') + .remove(b')'); + percent_encoding::utf8_percent_encode(segment, COMPONENT).to_string() + } + match name.strip_prefix("files/") { + Some(segment) if !segment.is_empty() && !segment.contains('/') => { + format!("files/{}", encode(segment)) + } + _ => encode(name), + } +} + /// The file resource returned by the Google upload API. #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] @@ -132,13 +155,13 @@ impl Files for GoogleFiles { &self, options: &UploadFileCallOptions, ) -> Result { - let google_options = parse_google_files_options(options.provider_options.as_ref()); + let google_options = parse_google_files_options(options.provider_options.as_ref())?; let exchange = self.config.exchange(None).await?; // The upload endpoint hangs off the origin: `base_url` without // `/v1beta`. let init_endpoint = format!( "{}/upload/v1beta/files", - exchange.base_url().replace("/v1beta", "") + exchange.base_url().trim_end_matches("/v1beta") ); let mut warnings = Vec::new(); @@ -234,7 +257,7 @@ impl Files for GoogleFiles { // header and receives the user's file bytes; validate it. (AI SDK // fetches this URL unvalidated — kept stricter here deliberately.) let upload_resp = aimux_provider_utils::post_to_api( - HttpRequest { + exchange.with_transport(HttpRequest { url: upload_url, headers: upload_headers, abort_signal: options.abort_signal.clone(), @@ -242,7 +265,7 @@ impl Files for GoogleFiles { trusted_origin: Some(exchange.base_url().to_string()), credentialed_origin: Some(exchange.base_url().to_string()), ..Default::default() - }, + }), HttpBody::Bytes(file_bytes, media_type.clone()), aimux_provider_utils::create_json_response_handler::(), super::google_failed_response_handler(), @@ -265,8 +288,8 @@ impl Files for GoogleFiles { .unwrap_or(Value::Null); // Step 3: Poll if file is PROCESSING. - let poll_interval_ms = google_options.poll_interval_ms.unwrap_or(2000); - let poll_timeout_ms = google_options.poll_timeout_ms.unwrap_or(300000); + let poll_interval_ms = google_options.poll_interval_ms.unwrap_or(2000.0); + let poll_timeout_ms = google_options.poll_timeout_ms.unwrap_or(300000.0); let start_time = Instant::now(); // Seed evidence from the upload response so a file that is already @@ -275,7 +298,7 @@ impl Files for GoogleFiles { let mut last_poll_status: Option = Some(200); let mut last_poll_url = upload_request_url; while file.state == "PROCESSING" { - if start_time.elapsed() > Duration::from_millis(poll_timeout_ms) { + if start_time.elapsed().as_secs_f64() * 1000.0 > poll_timeout_ms { return Err(AiMuxError::Timeout(format!( "Google file upload polling for {} timed out after {}ms", file.name, poll_timeout_ms @@ -283,12 +306,13 @@ impl Files for GoogleFiles { } sleep_or_abort( - Duration::from_millis(poll_interval_ms), + Duration::try_from_secs_f64(poll_interval_ms / 1000.0) + .map_err(|error| AiMuxError::InvalidArgument(error.to_string()))?, options.abort_signal.as_ref(), ) .await?; - let poll_url = exchange.url(&format!("/{}", file.name)); + let poll_url = exchange.url(&format!("/{}", file_poll_path(&file.name))); let poll_resp = aimux_provider_utils::get_from_api( exchange.with_transport(HttpRequest { diff --git a/aimux-providers/src/google/image.rs b/aimux-providers/src/google/image.rs index e113c32cb..4ae1a757e 100644 --- a/aimux-providers/src/google/image.rs +++ b/aimux-providers/src/google/image.rs @@ -1,44 +1,26 @@ -//! Google image model — implements the `ImageModel` trait. -//! -//! Aligned with Vercel AI SDK `GoogleImageModel` -//! (`reference/ai/packages/google/src/google-image-model.ts`). -//! -//! Two code paths: -//! - **Imagen** models (non-`gemini-*`): `POST {base_url}/models/{id}:predict` -//! - **Gemini** image models (`gemini-*`): `POST {base_url}/models/{id}:generateContent` - -use std::collections::HashMap; - -use async_trait::async_trait; -use serde_json::{Map, Value, json}; +//! Google image generation delegates to the Gemini language model. +use super::convert::{build_request_body_with_warnings, convert_usage, validate_call_options}; +use super::options::{GOOGLE, google_metadata}; +use super::types::GenerateContentResponse; +use crate::shared::EndpointConfig; use aimux_core::error::AiMuxError; use aimux_core::image_model::{ ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, ImageResponse, ImageResult, ImageUsage, }; -use aimux_core::shared::{SharedProviderMetadata, SharedProviderOptions, Warning}; - -use super::options::{GOOGLE, google_options as read_google_options}; -use crate::shared::EndpointConfig; - -/// Google error structure: `{ "error": { "message": "...", "status": "..." } }`. -/// Returns `true` if the model ID is a Gemini image model. -fn is_gemini_model(model_id: &str) -> bool { - model_id.starts_with("gemini-") -} +use aimux_core::language_model_message::{FilePart, LanguageModelMessage, TextPart, UserPart}; +use aimux_core::options::CallOptions; +use aimux_core::shared::{FileBytes, FileData, Warning}; +use aimux_core::tool::{ProviderTool, Tool}; +use async_trait::async_trait; +use serde_json::json; -/// Settings for the Google image model. #[derive(Debug, Clone, Default)] pub struct GoogleImageSettings { - /// Override the maximum number of images per call. pub max_images_per_call: Option, } -/// A Google image generation model (Imagen or Gemini). -/// -/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the process-wide shared -/// `Client` internally (RFC-0009 §4.1). pub struct GoogleImageModel { model_id: String, settings: GoogleImageSettings, @@ -57,415 +39,194 @@ impl GoogleImageModel { config, } } +} - fn predict_path(&self) -> String { - format!("/models/{}:predict", self.model_id) +#[async_trait] +impl ImageModel for GoogleImageModel { + fn provider(&self) -> &str { + &self.config.provider } - - fn generate_content_path(&self) -> String { - let model_path = if self.model_id.contains('/') { - self.model_id.clone() - } else { - format!("models/{}", self.model_id) - }; - format!("/{model_path}:generateContent") + fn model_id(&self) -> &str { + &self.model_id } - - fn max_images(&self) -> u32 { - if let Some(max) = self.settings.max_images_per_call { - return max; - } - 1 + fn max_images_per_call(&self) -> Option { + Some(self.settings.max_images_per_call.unwrap_or(1)) } - // ── Imagen path ───────────────────────────────────────────────────────── - - async fn do_generate_imagen( - &self, - options: &ImageCallOptions, - ) -> Result { - let mut warnings = Vec::new(); - - // Imagen API endpoints do not support image editing - if let Some(ref files) = options.files - && !files.is_empty() - { - return Err(AiMuxError::UnsupportedFunctionality( - "Google Gemini API does not support image editing with Imagen models. \ - Use Google Vertex AI (@ai-sdk/google-vertex) for image editing capabilities." - .to_string(), - )); + async fn do_generate(&self, options: &ImageCallOptions) -> Result { + if !self.model_id.starts_with("gemini-") { + return Err(AiMuxError::UnsupportedFunctionality("Google image models other than Gemini are no longer supported. Use a model ID that starts with `gemini-`.".into())); } - if options.mask.is_some() { return Err(AiMuxError::UnsupportedFunctionality( - "Google Gemini API does not support image editing with masks. \ - Use Google Vertex AI (@ai-sdk/google-vertex) for image editing capabilities." - .to_string(), + "Gemini image models do not support mask-based image editing.".into(), )); } - - if options.size.is_some() { - warnings.push(Warning::Unsupported { - feature: "size".to_string(), + let warnings = if options.size.is_some() { + vec![Warning::Unsupported { + feature: "size".into(), details: Some( "This model does not support the `size` option. Use `aspectRatio` instead." - .to_string(), + .into(), ), - }); - } - - if options.seed.is_some() { - warnings.push(Warning::Unsupported { - feature: "seed".to_string(), - details: Some( - "This model does not support the `seed` option through this provider." - .to_string(), - ), - }); - } - - let google_options = parse_google_image_options(&options.provider_options); - - let mut parameters = Map::new(); - parameters.insert("sampleCount".to_string(), json!(options.n)); - - if let Some(ar) = options.aspect_ratio { - parameters.insert("aspectRatio".to_string(), json!(ar.to_string())); - } - - if let Some(pg) = google_options.person_generation { - parameters.insert("personGeneration".to_string(), json!(pg)); - } - - if google_options.google_search.is_some() { - warnings.push(Warning::Unsupported { - feature: "googleSearch".to_string(), - details: Some( - "Google Search grounding is only supported on Gemini image models.".to_string(), + }] + } else { + Vec::new() + }; + let mut content = Vec::new(); + if let Some(prompt) = &options.prompt { + content.push(UserPart::Text(TextPart { + text: prompt.clone(), + provider_options: None, + })); + } + for file in options.files.iter().flatten() { + let (data, media_type) = match file { + ImageFile::Url { url } => ( + FileData::Url { + url: url.clone(), + original_url: None, + }, + "image/*".into(), ), - }); - } - - let body = json!({ - "instances": [{ "prompt": options.prompt }], - "parameters": parameters, - }); - - let exchange = self.config.exchange(options.headers.as_ref()).await?; - let resp = aimux_provider_utils::post_json_to_api( - exchange.request(exchange.url(&self.predict_path()), options), - exchange.transform_body(body), - aimux_provider_utils::create_json_response_handler(), - super::google_failed_response_handler(), - ) - .await?; - - let response_headers = resp.response_headers; - - let response_body: Value = resp.value; - - let images = extract_imagen_images(&response_body); - let provider_metadata = extract_imagen_metadata(&response_body); - - Ok(ImageResult { - images, - warnings, - provider_metadata: Some(provider_metadata), - response: ImageResponse { - timestamp: Some(chrono::Utc::now().to_rfc3339()), - model_id: Some(self.model_id.clone()), - headers: Some(response_headers), - }, - usage: None, - }) - } - - // ── Gemini path ───────────────────────────────────────────────────────── - - async fn do_generate_gemini( - &self, - options: &ImageCallOptions, - ) -> Result { - let mut warnings = Vec::new(); - - // Gemini does not support mask-based inpainting - if options.mask.is_some() { - return Err(AiMuxError::UnsupportedFunctionality( - "Gemini image models do not support mask-based image editing.".to_string(), - )); - } - - // Gemini does not support generating multiple images per call via n parameter - if options.n > 1 { - return Err(AiMuxError::UnsupportedFunctionality( - "Gemini image models do not support generating a set number of images per call. \ - Use n=1 or omit the n parameter." - .to_string(), - )); - } - - if options.size.is_some() { - warnings.push(Warning::Unsupported { - feature: "size".to_string(), - details: Some( - "This model does not support the `size` option. Use `aspectRatio` instead." - .to_string(), + ImageFile::File { media_type, data } => ( + FileData::Data { + data: match data { + ImageFileData::Base64(data) => FileBytes::Base64(data.clone()), + ImageFileData::Binary(data) => FileBytes::Binary(data.clone()), + }, + }, + media_type.clone(), ), - }); - } - - let google_options = parse_google_image_options(&options.provider_options); - - // Build contents parts - let mut parts: Vec = Vec::new(); - if let Some(ref prompt) = options.prompt { - parts.push(json!({ "text": prompt })); - } - - if let Some(ref files) = options.files { - for file in files { - match file { - ImageFile::Url { url } => { - return Err(AiMuxError::UnsupportedFunctionality(format!( - "URL-based input images with media type \"image/*\" are not passed as \ - inline bytes. URL: {url}" - ))); - } - ImageFile::File { media_type, data } => { - let data_str = match data { - ImageFileData::Base64(b64) => b64.clone(), - ImageFileData::Binary(bytes) => base64::Engine::encode( - &base64::engine::general_purpose::STANDARD, - bytes, - ), - }; - parts.push(json!({ - "inlineData": { - "mimeType": media_type, - "data": data_str, - } - })); - } - } + }; + content.push(UserPart::File(FilePart { + data, + media_type, + filename: None, + provider_options: None, + })); + } + let mut call = CallOptions::new(vec![LanguageModelMessage::User { + content, + provider_options: None, + }]); + call.seed = options.seed; + call.headers = options.headers.clone(); + call.abort_signal = options.abort_signal.clone(); + call.provider_options = Some(options.provider_options.clone()); + let google = call + .provider_options + .as_mut() + .expect("set above") + .entry(GOOGLE.into()) + .or_default(); + if let Some(search) = google.remove("googleSearch") { + let invalid = + || AiMuxError::InvalidArgument("Invalid Google image googleSearch option".into()); + let object = search.as_object().ok_or_else(invalid)?; + if object.get("searchTypes").is_some_and(|types| { + !types.as_object().is_some_and(|types| { + ["webSearch", "imageSearch"] + .iter() + .all(|key| types.get(*key).is_none_or(serde_json::Value::is_object)) + }) + }) || object.get("timeRangeFilter").is_some_and(|range| { + !range.as_object().is_some_and(|range| { + ["startTime", "endTime"] + .iter() + .all(|key| range.get(*key).is_some_and(serde_json::Value::is_string)) + }) + }) { + return Err(invalid()); } - } - - let contents = json!([{ "role": "user", "parts": parts }]); - - // Build generationConfig - let mut generation_config = Map::new(); - generation_config.insert("responseModalities".to_string(), json!(["IMAGE"])); - - if let Some(ar) = options.aspect_ratio { - generation_config.insert( - "imageConfig".to_string(), - json!({ "aspectRatio": ar.to_string() }), - ); - } - - if let Some(seed) = options.seed { - generation_config.insert("seed".to_string(), json!(seed)); - } - - // Only declared generation-config options survive the upstream language schema. - if let Some(google) = read_google_options(Some(&options.provider_options)) { - for key in [ - "audioTimestamp", - "thinkingConfig", - "mediaResolution", - "imageConfig", - ] { - if let Some(value) = google.get(key) { - generation_config.insert( - key.to_string(), - crate::google::convert::image_generation_option(key, value), - ); - } + call.tools = Some(vec![Tool::Provider(ProviderTool { + id: "google.google_search".into(), + name: "google_search".into(), + args: object.clone(), + })]); + } + google.insert("responseModalities".into(), json!(["IMAGE"])); + let user_config = google.remove("imageConfig").filter(|v| !v.is_null()); + if options.aspect_ratio.is_some() || user_config.is_some() { + let mut config = user_config + .and_then(|v| v.as_object().cloned()) + .unwrap_or_default(); + if let Some(ratio) = options.aspect_ratio { + config.insert("aspectRatio".into(), json!(ratio.to_string())); } + google.insert("imageConfig".into(), serde_json::Value::Object(config)); } - - let mut body = Map::new(); - body.insert("contents".to_string(), contents); - body.insert( - "generationConfig".to_string(), - Value::Object(generation_config), - ); - - // Tools (googleSearch) - if let Some(ref gs) = google_options.google_search { - body.insert("tools".to_string(), json!([{ "googleSearch": gs }])); - } - + validate_call_options(&call)?; + let (body, _) = build_request_body_with_warnings(&self.model_id, &call)?; let exchange = self.config.exchange(options.headers.as_ref()).await?; - let resp = aimux_provider_utils::post_json_to_api( - exchange.request(exchange.url(&self.generate_content_path()), options), - exchange.transform_body(Value::Object(body)), - aimux_provider_utils::create_json_response_handler(), + let response = aimux_provider_utils::post_json_to_api( + exchange.request( + exchange.url(&format!("/models/{}:generateContent", self.model_id)), + options, + ), + exchange.transform_body(body), + aimux_provider_utils::create_json_response_handler::(), super::google_failed_response_handler(), ) .await?; - - let response_headers = resp.response_headers; - - let response_body: Value = resp.value; - - let (images, provider_metadata, usage) = extract_gemini_result(&response_body); - + let data = response.value; + let candidate = data.candidates.into_iter().next().unwrap_or_default(); + let images: Vec = candidate + .content + .as_ref() + .and_then(|c| c.parts.as_ref()) + .into_iter() + .flatten() + .filter_map(|part| { + if part.get("thought").and_then(serde_json::Value::as_bool) == Some(true) { + return None; + } + let inline = part.get("inlineData")?; + inline + .get("mimeType")? + .as_str()? + .starts_with("image/") + .then(|| { + inline + .get("data") + .and_then(serde_json::Value::as_str) + .map(String::from) + }) + .flatten() + }) + .collect(); + let metadata = google_metadata(json!({ + "promptFeedback": data.prompt_feedback, + "groundingMetadata": candidate.grounding_metadata, + "urlContextMetadata": candidate.url_context_metadata, + "safetyRatings": candidate.safety_ratings, + "usageMetadata": data.usage_metadata, + "finishMessage": candidate.finish_message, + "serviceTier": data.usage_metadata.as_ref().and_then(|u| u.service_tier.as_ref()), + "finishReason": candidate.finish_reason.or_else(|| data.prompt_feedback.as_ref().and_then(|f| f.get("blockReason")).and_then(serde_json::Value::as_str).filter(|reason| !reason.is_empty() && *reason != "BLOCK_REASON_UNSPECIFIED" && *reason != "BLOCKED_REASON_UNSPECIFIED").map(String::from)), + "images": images.iter().map(|_| json!({})).collect::>(), + })); + let usage = data + .usage_metadata + .as_ref() + .map(convert_usage) + .unwrap_or_default(); + let input = usage.input_tokens.total; + let output = usage.output_tokens.total; Ok(ImageResult { - images, + images: ImageOutputs::Base64(images), warnings, - provider_metadata: Some(provider_metadata), + provider_metadata: Some(metadata), response: ImageResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), model_id: Some(self.model_id.clone()), - headers: Some(response_headers), + headers: Some(response.response_headers), }, - usage, + usage: Some(ImageUsage { + input_tokens: input, + output_tokens: output, + total_tokens: Some(input.unwrap_or(0) + output.unwrap_or(0)), + }), }) } } - -#[async_trait] -impl ImageModel for GoogleImageModel { - fn provider(&self) -> &str { - &self.config.provider - } - - fn model_id(&self) -> &str { - &self.model_id - } - - fn max_images_per_call(&self) -> Option { - Some(self.max_images()) - } - - async fn do_generate(&self, options: &ImageCallOptions) -> Result { - if is_gemini_model(&self.model_id) { - self.do_generate_gemini(options).await - } else { - self.do_generate_imagen(options).await - } - } -} - -// ── Helpers ───────────────────────────────────────────────────────────────── - -/// Parsed Google image provider options. -struct GoogleImageOptions { - person_generation: Option, - google_search: Option, -} - -/// Parse Google image provider options from the `"google"` key. -fn parse_google_image_options(provider_options: &SharedProviderOptions) -> GoogleImageOptions { - let google = read_google_options(Some(provider_options)); - GoogleImageOptions { - person_generation: google - .and_then(|g| g.get("personGeneration")) - .and_then(|v| v.as_str()) - .map(String::from), - google_search: google.and_then(|g| g.get("googleSearch")).cloned(), - } -} - -/// Extract base64 images from an Imagen response. -fn extract_imagen_images(response: &Value) -> ImageOutputs { - let images: Vec = response - .get("predictions") - .and_then(|p| p.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|p| { - p.get("bytesBase64Encoded") - .and_then(|v| v.as_str()) - .map(String::from) - }) - .collect() - }) - .unwrap_or_default(); - ImageOutputs::Base64(images) -} - -/// Extract provider metadata from an Imagen response. -fn extract_imagen_metadata(response: &Value) -> SharedProviderMetadata { - let mut metadata = HashMap::new(); - let predictions = response - .get("predictions") - .and_then(|p| p.as_array()) - .map(std::vec::Vec::len) - .unwrap_or(0); - - let images: Vec = (0..predictions).map(|_| json!({})).collect(); - let mut google_meta = Map::new(); - google_meta.insert("images".to_string(), json!(images)); - metadata.insert(GOOGLE.to_string(), google_meta); - metadata -} - -/// Extract images, provider metadata, and usage from a Gemini generateContent response. -fn extract_gemini_result( - response: &Value, -) -> (ImageOutputs, SharedProviderMetadata, Option) { - let mut images: Vec = Vec::new(); - let mut grounding_metadata: Option = None; - - if let Some(candidates) = response.get("candidates").and_then(|c| c.as_array()) { - for candidate in candidates { - // Extract grounding metadata - if let Some(gm) = candidate.get("groundingMetadata") { - grounding_metadata = Some(gm.clone()); - } - // Extract images from content parts - if let Some(parts) = candidate - .get("content") - .and_then(|c| c.get("parts")) - .and_then(|p| p.as_array()) - { - for part in parts { - if let Some(inline_data) = part.get("inlineData") - && let Some(mime_type) = - inline_data.get("mimeType").and_then(|m| m.as_str()) - && mime_type.starts_with("image/") - && let Some(data) = inline_data.get("data").and_then(|d| d.as_str()) - { - images.push(data.to_string()); - } - } - } - } - } - - // Usage - let usage = response.get("usageMetadata").map(|u| { - let input = u - .get("promptTokenCount") - .and_then(serde_json::Value::as_u64) - .map(|x| x as u32); - let output = u - .get("candidatesTokenCount") - .and_then(serde_json::Value::as_u64) - .map(|x| x as u32); - let total = u - .get("totalTokenCount") - .and_then(serde_json::Value::as_u64) - .map(|x| x as u32); - ImageUsage { - input_tokens: input, - output_tokens: output, - total_tokens: total.or_else(|| Some(input.unwrap_or(0) + output.unwrap_or(0))), - } - }); - - // Provider metadata - let mut metadata = HashMap::new(); - let mut google_meta = Map::new(); - let image_metas: Vec = images.iter().map(|_| json!({})).collect(); - google_meta.insert("images".to_string(), json!(image_metas)); - if let Some(gm) = grounding_metadata { - google_meta.insert("groundingMetadata".to_string(), gm); - } - metadata.insert(GOOGLE.to_string(), google_meta); - - (ImageOutputs::Base64(images), metadata, usage) -} diff --git a/aimux-providers/src/google/mod.rs b/aimux-providers/src/google/mod.rs index ed4347317..fa74db9e4 100644 --- a/aimux-providers/src/google/mod.rs +++ b/aimux-providers/src/google/mod.rs @@ -184,12 +184,14 @@ pub struct GoogleProviderSettings { /// Extra headers on every request. A `None` value removes the header, /// including `x-goog-api-key`. Per-call headers win over these. pub headers: Option, - /// The provider name, the `provider()` string of the language, embedding, - /// image, video and files models. Default `"google.generative-ai"`. + /// The provider name, the `provider()` string of the language, embedding + /// image, video and files interfaces. Default `"google.generative-ai"`. pub name: Option, /// The transport: a mock, a signing decorator, a proxy-aware client. /// `None` uses the process default, resolved per request. pub fetch: Option, + /// Generates IDs for tool calls and sources that have no provider ID. + pub generate_id: Option String + Send + Sync>>, } impl std::fmt::Debug for GoogleProviderSettings { @@ -204,6 +206,7 @@ impl std::fmt::Debug for GoogleProviderSettings { ) .field("name", &self.name) .field("fetch", &self.fetch.is_some()) + .field("generate_id", &self.generate_id.is_some()) .finish() } } @@ -238,6 +241,7 @@ pub fn create_google(settings: GoogleProviderSettings) -> Result, + generate_id: Option String + Send + Sync>>, } impl GoogleProvider { @@ -290,6 +295,13 @@ impl GoogleProvider { #[must_use] pub fn chat(&self, model_id: &str) -> GoogleModel { GoogleModel::from_config(model_id.to_string(), self.model_config(self.name.clone())) + .with_generate_id(self.generate_id.clone()) + } + + /// Alias of [`chat`](Self::chat), matching `generativeAI`. + #[must_use] + pub fn generative_ai(&self, model_id: &str) -> GoogleModel { + self.chat(model_id) } /// A text embedding model (e.g. `"gemini-embedding-001"`); `provider()` is @@ -302,8 +314,7 @@ impl GoogleProvider { ) } - /// An image generation model (e.g. `"imagen-3.0-generate-002"` or - /// `"gemini-2.5-flash-image"`); `provider()` is the provider name. + /// An image generation model (e.g. `"gemini-2.5-flash-image"`); `provider()` is the provider name. #[must_use] pub fn image(&self, model_id: &str) -> GoogleImageModel { self.image_with_settings(model_id, GoogleImageSettings::default()) diff --git a/aimux-providers/src/google/model.rs b/aimux-providers/src/google/model.rs index 3a77425c1..9ef7e7d09 100644 --- a/aimux-providers/src/google/model.rs +++ b/aimux-providers/src/google/model.rs @@ -22,7 +22,7 @@ use aimux_core::language_model::SupportedUrls; use super::convert::{ build_request_body_with_warnings, code_execution_tool_name, convert_usage, extract_sources, - parse_finish_reason, + parse_finish_reason, validate_call_options, }; use super::options::google_metadata; use super::types::{Candidate, GenerateContentResponse, GoogleStreamEvent}; @@ -35,11 +35,26 @@ use crate::shared::EndpointConfig; pub struct GoogleModel { model_id: String, config: EndpointConfig, + generate_id: std::sync::Arc String + Send + Sync>, } impl GoogleModel { pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { - Self { model_id, config } + Self { + model_id, + config, + generate_id: std::sync::Arc::new(aimux_provider_utils::generate_id), + } + } + + pub(crate) fn with_generate_id( + mut self, + generate_id: Option String + Send + Sync>>, + ) -> Self { + if let Some(generate_id) = generate_id { + self.generate_id = generate_id; + } + self } /// `models/{model}`, or the id itself when it already contains a `/` @@ -68,41 +83,14 @@ impl LanguageModel for GoogleModel { } async fn do_generate(&self, options: &CallOptions) -> Result { - if options - .prompt - .iter() - .filter_map(|message| match message { - aimux_core::language_model_message::LanguageModelMessage::Assistant { - content, - .. - } => Some(content), - _ => None, - }) - .flatten() - .any(|part| { - matches!( - part, - aimux_core::language_model_message::AssistantPart::ReasoningFile( - aimux_core::language_model_message::ReasoningFilePart { - data: aimux_core::shared::GeneratedFileData::Url { .. }, - .. - } - ) | aimux_core::language_model_message::AssistantPart::File( - aimux_core::language_model_message::FilePart { - data: aimux_core::shared::FileData::Url { .. }, - .. - } - ) - ) - }) - { - return Err(AiMuxError::UnsupportedFunctionality( - "File data URLs in assistant messages are not supported".to_string(), - )); - } - crate::google::convert::validate_tool_result_files(&options.prompt)?; + validate_call_options(options)?; let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); - let (body, tool_warnings) = build_request_body_with_warnings(&self.model_id, options)?; + let (body, mut tool_warnings) = build_request_body_with_warnings(&self.model_id, options)?; + for warning in &mut tool_warnings { + if let aimux_core::types::Warning::Other { message } = warning { + *message = message.replace("google.generative-ai", &self.config.provider); + } + } let exchange = self.config.exchange(options.headers.as_ref()).await?; let body = exchange.transform_body(body); let url = exchange.url(&format!("/{}:generateContent", self.model_path())); @@ -119,21 +107,30 @@ impl LanguageModel for GoogleModel { let data: GenerateContentResponse = resp.value; - let candidate = data.candidates.into_iter().next().ok_or_else(|| { - AiMuxError::InvalidResponseData("no candidates in response".to_string()) - })?; + let candidate = data.candidates.into_iter().next().unwrap_or_default(); + let block_reason = confirmed_prompt_block_reason(data.prompt_feedback.as_ref()); - let (content, has_tool_calls) = - extract_content_from_candidate(&candidate, &code_execution_tool_name); + let (content, has_tool_calls) = extract_content_from_candidate( + &candidate, + &code_execution_tool_name, + &*self.generate_id, + ); - let finish_reason = candidate - .finish_reason - .as_deref() - .map(|r| parse_finish_reason(r, has_tool_calls)) - .unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }); + let finish_reason = if candidate.finish_reason.is_none() && block_reason.is_some() { + FinishReason { + unified: FinishReasonUnified::ContentFilter, + raw: block_reason, + } + } else { + candidate + .finish_reason + .as_deref() + .map(|r| parse_finish_reason(r, has_tool_calls)) + .unwrap_or(FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }) + }; let usage = data .usage_metadata @@ -150,6 +147,7 @@ impl LanguageModel for GoogleModel { "safetyRatings": candidate.safety_ratings, "usageMetadata": data.usage_metadata, "finishMessage": candidate.finish_message, + "serviceTier": data.usage_metadata.as_ref().and_then(|usage| usage.service_tier.as_ref()), }))); Ok(GenerateResult { @@ -161,7 +159,7 @@ impl LanguageModel for GoogleModel { response: Some(aimux_core::shared::ResponseInfo { id: data.response_id, timestamp: None, - model_id: data.model_version, + model_id: None, headers: Some(response_headers), body: response_body, }), @@ -170,41 +168,14 @@ impl LanguageModel for GoogleModel { } async fn do_stream(&self, options: &CallOptions) -> Result { - if options - .prompt - .iter() - .filter_map(|message| match message { - aimux_core::language_model_message::LanguageModelMessage::Assistant { - content, - .. - } => Some(content), - _ => None, - }) - .flatten() - .any(|part| { - matches!( - part, - aimux_core::language_model_message::AssistantPart::ReasoningFile( - aimux_core::language_model_message::ReasoningFilePart { - data: aimux_core::shared::GeneratedFileData::Url { .. }, - .. - } - ) | aimux_core::language_model_message::AssistantPart::File( - aimux_core::language_model_message::FilePart { - data: aimux_core::shared::FileData::Url { .. }, - .. - } - ) - ) - }) - { - return Err(AiMuxError::UnsupportedFunctionality( - "File data URLs in assistant messages are not supported".to_string(), - )); - } - crate::google::convert::validate_tool_result_files(&options.prompt)?; + validate_call_options(options)?; let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); - let (body, tool_warnings) = build_request_body_with_warnings(&self.model_id, options)?; + let (body, mut tool_warnings) = build_request_body_with_warnings(&self.model_id, options)?; + for warning in &mut tool_warnings { + if let aimux_core::types::Warning::Other { message } = warning { + *message = message.replace("google.generative-ai", &self.config.provider); + } + } let exchange = self.config.exchange(options.headers.as_ref()).await?; let body = exchange.transform_body(body); let endpoint = exchange.url(&format!( @@ -214,7 +185,7 @@ impl LanguageModel for GoogleModel { let resp = aimux_provider_utils::post_json_to_api( exchange.request(endpoint.clone(), options), body.clone(), - aimux_provider_utils::create_event_source_response_handler::(), + aimux_provider_utils::create_event_source_response_handler::(), super::google_failed_response_handler(), ) .await?; @@ -226,18 +197,12 @@ impl LanguageModel for GoogleModel { Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), first_event => first_event, }; - if let Some(Ok(GoogleStreamEvent::Error(error))) = first_event.as_ref() { - return Err(super::google_stream_error( - &error.error, - &endpoint, - body.clone(), - response_headers, - )); - } let stream_error_url = endpoint.clone(); let stream_request_body = body.clone(); let stream_error_headers = response_headers.clone(); + let generate_id = self.generate_id.clone(); + let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings: tool_warnings }); @@ -250,6 +215,7 @@ impl LanguageModel for GoogleModel { let mut has_tool_calls = false; let mut response_metadata_emitted = false; let mut stream_errored = false; + let mut prompt_blocked = false; // Provider-metadata accumulators (mirrors TS `lastGroundingMetadata` / // `lastUrlContextMetadata` + the finishReason-chunk snapshot). @@ -274,6 +240,15 @@ impl LanguageModel for GoogleModel { if stream_errored { break; } + let event = match event { + Ok(raw_value) => { + if include_raw_chunks { + yield Ok(StreamPart::Raw { raw_value: raw_value.clone() }); + } + serde_json::from_value::(raw_value).map_err(AiMuxError::from) + } + Err(error) => Err(error), + }; match event { Ok(GoogleStreamEvent::Chunk(chunk)) => { @@ -285,7 +260,7 @@ impl LanguageModel for GoogleModel { yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { id: Some(id.clone()), timestamp: None, - model_id: chunk.model_version.clone(), + model_id: None, })); } @@ -294,9 +269,16 @@ impl LanguageModel for GoogleModel { last_usage_metadata_value = serde_json::to_value(usage).ok(); } - if let Some(pf) = &chunk.prompt_feedback { - last_prompt_feedback = Some(pf.clone()); - } + if let Some(pf) = &chunk.prompt_feedback + && !prompt_blocked { + last_prompt_feedback = Some(pf.clone()); + if let Some(reason) = confirmed_prompt_block_reason(Some(pf)) { + prompt_blocked = true; + final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::ContentFilter, raw: Some(reason), + }); + } + } let Some(candidates) = chunk.candidates else { continue; @@ -312,6 +294,14 @@ impl LanguageModel for GoogleModel { last_url_context_metadata = Some(ucm.clone()); } + if let Some(sr) = &candidate.safety_ratings { + last_safety_ratings = serde_json::to_value(sr).ok(); + } + if let Some(fm) = &candidate.finish_message { + last_finish_message = Some(json!(fm)); + } + if prompt_blocked { continue; } + // Extract url sources from this chunk's grounding metadata // (deduplicated across chunks; document sources are not // emitted in the stream, matching TS). @@ -320,13 +310,13 @@ impl LanguageModel for GoogleModel { for src in chunk_sources { if let GenerateContent::Source(Source::Url { url, - id, + id: _, title, provider_metadata: None, }) = src && emitted_source_urls.insert(url.clone()) { yield Ok(StreamPart::Source(Source::Url { - id, + id: generate_id(), url, title, provider_metadata: None, @@ -342,6 +332,7 @@ impl LanguageModel for GoogleModel { let thought_sig_meta: Option = part .get("thoughtSignature") .and_then(|v| v.as_str()) + .filter(|signature| !signature.is_empty()) .map(|s| google_metadata(json!({ "thoughtSignature": s }))); // text part @@ -354,21 +345,15 @@ impl LanguageModel for GoogleModel { // reasoning (mirrors the non-streaming // path, which attaches to the last // content item regardless of type). - if let Some(meta) = &thought_sig_meta { - if let Some(id) = &text_id { + if let Some(meta) = &thought_sig_meta + && let Some(id) = &text_id { yield Ok(StreamPart::TextDelta { id: id.clone(), delta: String::new(), provider_metadata: Some(meta.clone()), }); - } else if let Some(id) = &reasoning_id { - yield Ok(StreamPart::ReasoningDelta { - id: id.clone(), - delta: String::new(), - provider_metadata: Some(meta.clone()), - }); + } - } } else { let is_thought = part.get("thought").and_then(serde_json::Value::as_bool).unwrap_or(false); if is_thought { @@ -413,45 +398,6 @@ impl LanguageModel for GoogleModel { } } } - } else if let Some(fc) = part.get("functionCall") { - // Single-chunk complete function call (the - // common case for the public Gemini API). - let name = fc - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or(""); - let id = fc - .get("id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - .unwrap_or_else(|| format!("call-{block_counter}")); - block_counter += 1; - let args = fc.get("args").cloned().unwrap_or(json!({})); - - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: name.to_string(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: thought_sig_meta.clone(), - }); - let args_str = args.to_string(); - yield Ok(StreamPart::ToolInputDelta { - id: id.clone(), - delta: args_str, - provider_metadata: thought_sig_meta.clone(), - }); - yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: thought_sig_meta.clone()}); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: name.to_string(), - input: args.to_string(), - provider_executed: None, - dynamic: None, - provider_metadata: thought_sig_meta.clone(), - })); - has_tool_calls = true; } else if let Some(ec) = part.get("executableCode") { // Provider-executed code execution. let has_code = ec @@ -460,8 +406,7 @@ impl LanguageModel for GoogleModel { .map(|s| !s.is_empty()) .unwrap_or(false); if has_code { - let id = format!("call-{block_counter}"); - block_counter += 1; + let id = generate_id(); last_code_execution_tool_call_id = Some(id.clone()); yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id.clone(), @@ -469,11 +414,7 @@ impl LanguageModel for GoogleModel { input: ec.to_string(), provider_executed: Some(true), dynamic: None, - provider_metadata: Some(server_tool_metadata( - &id, - "code_execution", - None, - )), + provider_metadata: None, })); // provider-executed → does NOT set has_tool_calls } @@ -499,11 +440,7 @@ impl LanguageModel for GoogleModel { is_error: None, preliminary: None, dynamic: None, - provider_metadata: Some(server_tool_metadata( - call_id, - "code_execution", - None, - )), + provider_metadata: None, })); } } else if let Some(tc) = part.get("toolCall") { @@ -515,9 +452,9 @@ impl LanguageModel for GoogleModel { let id = tc .get("id") .and_then(|v| v.as_str()) + .filter(|id| !id.is_empty()) .map(std::string::ToString::to_string) - .unwrap_or_else(|| format!("call-{block_counter}")); - block_counter += 1; + .unwrap_or_else(|| generate_id()); last_server_tool_call_id = Some(id.clone()); let args = tc.get("args").cloned().unwrap_or(json!({})); let server_meta = server_tool_metadata( @@ -544,8 +481,8 @@ impl LanguageModel for GoogleModel { .and_then(|v| v.as_str()) .map(std::string::ToString::to_string) }) - .unwrap_or_else(|| format!("call-{block_counter}")); - block_counter += 1; + .filter(|id| !id.is_empty()) + .unwrap_or_else(|| generate_id()); let response = tr.get("response").cloned().unwrap_or(json!({})); let server_meta = server_tool_metadata( @@ -584,24 +521,58 @@ impl LanguageModel for GoogleModel { } } } + for part in parts { + let thought_sig_meta = part.get("thoughtSignature").and_then(Value::as_str) + .filter(|signature| !signature.is_empty()) + .map(|signature| google_metadata(json!({"thoughtSignature": signature}))); + if let Some(fc) = part.get("functionCall") { + // Single-chunk complete function call (the + // common case for the public Gemini API). + if fc.get("name").and_then(Value::as_str).is_none() || fc.get("partialArgs").is_some_and(|value| !value.is_null()) || fc.get("willContinue").and_then(Value::as_bool) == Some(true) { continue; } + let name = fc + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let id = fc + .get("id") + .and_then(|v| v.as_str()) + .filter(|id| !id.is_empty()) + .filter(|id| !id.is_empty()) + .map(std::string::ToString::to_string) + .unwrap_or_else(|| generate_id()); + let args = fc.get("args").filter(|value| !value.is_null()).cloned().unwrap_or(json!({})); + + yield Ok(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: name.to_string(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: thought_sig_meta.clone(), + }); + let args_str = args.as_str().map(str::to_owned).unwrap_or_else(|| args.to_string()); + if fc.get("args").is_some_and(|value| !value.is_null()) { + yield Ok(StreamPart::ToolInputDelta { + id: id.clone(), + delta: args_str.clone(), + provider_metadata: thought_sig_meta.clone(), + }); + } + yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: thought_sig_meta.clone()}); + yield Ok(StreamPart::ToolCall(RawToolCall { + tool_call_id: id, + tool_name: name.to_string(), + input: args_str, + provider_executed: None, + dynamic: None, + provider_metadata: thought_sig_meta.clone(), + })); + has_tool_calls = true; + } + } } if let Some(reason) = candidate.finish_reason.as_deref() { - // Close any open text/reasoning segment. - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); - } - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None}); - } - // Snapshot the finishReason-chunk metadata. - if let Some(sr) = &candidate.safety_ratings { - last_safety_ratings = - Some(serde_json::to_value(sr).unwrap_or(Value::Null)); - } - if let Some(fm) = &candidate.finish_message { - last_finish_message = Some(json!(fm)); - } final_finish_reason = Some(parse_finish_reason(reason, has_tool_calls)); } @@ -620,7 +591,7 @@ impl LanguageModel for GoogleModel { } Err(error) => { let recoverable = error.is_recoverable_stream_error(); - yield Err(error); + yield Ok(StreamPart::Error { error }); if !recoverable { return; } @@ -643,6 +614,7 @@ impl LanguageModel for GoogleModel { "safetyRatings": last_safety_ratings, "usageMetadata": last_usage_metadata_value, "finishMessage": last_finish_message, + "serviceTier": last_usage_metadata_value.as_ref().and_then(|usage| usage.get("serviceTier")), }))); yield Ok(StreamPart::Finish { @@ -653,7 +625,7 @@ impl LanguageModel for GoogleModel { } } else { final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Stop, + unified: FinishReasonUnified::Other, raw: None, }) }, @@ -702,6 +674,7 @@ fn server_tool_metadata( fn extract_content_from_candidate( candidate: &Candidate, code_execution_tool_name: &str, + generate_id: &dyn Fn() -> String, ) -> (Vec, bool) { let mut content = Vec::new(); let mut has_tool_calls = false; @@ -720,6 +693,7 @@ fn extract_content_from_candidate( let thought_sig_meta: Option = part .get("thoughtSignature") .and_then(|v| v.as_str()) + .filter(|signature| !signature.is_empty()) .map(|s| google_metadata(json!({ "thoughtSignature": s }))); // Branch order matches upstream (google-language-model.ts:420-534): @@ -732,7 +706,7 @@ fn extract_content_from_candidate( .map(|s| !s.is_empty()) .unwrap_or(false); if has_code { - let id = format!("call-{}", content.len()); + let id = generate_id(); last_code_execution_tool_call_id = Some(id.clone()); content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id.clone(), @@ -740,7 +714,7 @@ fn extract_content_from_candidate( input: ec.to_string(), provider_executed: Some(true), dynamic: None, - provider_metadata: Some(server_tool_metadata(&id, "code_execution", None)), + provider_metadata: None, })); } } else if let Some(cer) = part.get("codeExecutionResult") { @@ -759,11 +733,7 @@ fn extract_content_from_candidate( is_error: None, preliminary: None, dynamic: None, - provider_metadata: Some(server_tool_metadata( - call_id, - "code_execution", - None, - )), + provider_metadata: None, })); } } else if let Some(text) = part.get("text").and_then(|v| v.as_str()) { @@ -790,7 +760,10 @@ fn extract_content_from_candidate( }); } } - } else if let Some(fc) = part.get("functionCall") { + } else if let Some(fc) = part + .get("functionCall") + .filter(|fc| fc.get("name").and_then(Value::as_str).is_some()) + { let name = fc .get("name") .and_then(|v| v.as_str()) @@ -799,9 +772,14 @@ fn extract_content_from_candidate( let id = fc .get("id") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let input = fc.get("args").cloned().unwrap_or(json!({})); + .filter(|id| !id.is_empty()) + .map(str::to_owned) + .unwrap_or_else(generate_id); + let input = fc + .get("args") + .filter(|value| !value.is_null()) + .cloned() + .unwrap_or(json!({})); content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id, tool_name: name, @@ -842,8 +820,9 @@ fn extract_content_from_candidate( let id = tc .get("id") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); + .filter(|id| !id.is_empty()) + .map(str::to_owned) + .unwrap_or_else(generate_id); last_server_tool_call_id = Some(id.clone()); let input = tc.get("args").cloned().unwrap_or(json!({})); let server_meta = server_tool_metadata( @@ -870,7 +849,8 @@ fn extract_content_from_candidate( .and_then(|v| v.as_str()) .map(std::string::ToString::to_string) }) - .unwrap_or_default(); + .filter(|id| !id.is_empty()) + .unwrap_or_else(generate_id); let response = tr.get("response").cloned().unwrap_or(json!({})); let server_meta = server_tool_metadata( &id, @@ -893,6 +873,13 @@ fn extract_content_from_candidate( // Sources are appended after the parts (mirrors TS). let mut sources = extract_sources(candidate.grounding_metadata.as_ref(), &mut source_id); + for source in &mut sources { + if let GenerateContent::Source(Source::Url { id, .. } | Source::Document { id, .. }) = + source + { + *id = generate_id(); + } + } content.append(&mut sources); (content, has_tool_calls) @@ -937,3 +924,15 @@ fn set_provider_metadata(item: &mut GenerateContent, meta: ProviderMetadata) { } } } + +fn confirmed_prompt_block_reason(feedback: Option<&Value>) -> Option { + feedback + .and_then(|feedback| feedback.get("blockReason")) + .and_then(Value::as_str) + .filter(|reason| { + !reason.is_empty() + && *reason != "BLOCK_REASON_UNSPECIFIED" + && *reason != "BLOCKED_REASON_UNSPECIFIED" + }) + .map(str::to_owned) +} diff --git a/aimux-providers/src/google/options.rs b/aimux-providers/src/google/options.rs index 03580dfd3..c0d7c8edd 100644 --- a/aimux-providers/src/google/options.rs +++ b/aimux-providers/src/google/options.rs @@ -2,10 +2,9 @@ //! //! `@ai-sdk/google` reads `providerOptions.google` and writes response //! metadata under `google`; the Vertex package reuses the same models and -//! reads `googleVertex` (then `google`, because the shared Gemini model reads -//! it) and writes `googleVertex`. Only these canonical keys exist: the SDK's -//! historical `vertex` alias is not read and not written. This module is the -//! one place that knows the keys, so the models never spell them. +//! reads `googleVertex`, then legacy `vertex`, then `google`, and writes +//! metadata under both Vertex keys. Part metadata additionally supports the +//! cross-provider namespace fallback used by the upstream message converter. use aimux_core::shared::{JsonObject, SharedProviderOptions, provider_namespace}; use aimux_core::types::ProviderMetadata; @@ -25,7 +24,7 @@ pub(crate) const VERTEX_VIDEO_METADATA_KEYS: [&str; 3] = [GOOGLE_VERTEX, "google pub(crate) enum Namespace { /// The public Gemini API: reads and writes `google`. Google, - /// Vertex AI: reads `googleVertex`, then `google`; writes `googleVertex`. + /// Vertex AI: reads `googleVertex`, `vertex`, then `google`; writes both Vertex keys. Vertex, } @@ -34,7 +33,7 @@ impl Namespace { pub(crate) fn read_keys(self) -> &'static [&'static str] { match self { Self::Google => &[GOOGLE], - Self::Vertex => &[GOOGLE_VERTEX, GOOGLE], + Self::Vertex => &[GOOGLE_VERTEX, "vertex", GOOGLE], } } @@ -42,7 +41,7 @@ impl Namespace { pub(crate) fn write_keys(self) -> &'static [&'static str] { match self { Self::Google => &[GOOGLE], - Self::Vertex => &[GOOGLE_VERTEX], + Self::Vertex => &[GOOGLE_VERTEX, "vertex"], } } @@ -57,6 +56,15 @@ impl Namespace { .find_map(|key| provider_options.get(*key)) } + pub(crate) fn read_part(self, options: Option<&SharedProviderOptions>) -> Option<&JsonObject> { + let options = options?; + let keys: &[&str] = match self { + Self::Google => &[GOOGLE, GOOGLE_VERTEX, "vertex"], + Self::Vertex => &[GOOGLE_VERTEX, "vertex", GOOGLE], + }; + keys.iter().find_map(|key| options.get(*key)) + } + /// [`read`](Self::read) for `CallOptions::provider_options`. pub(crate) fn read_in( self, @@ -67,8 +75,15 @@ impl Namespace { /// Wrap `payload` as response metadata under every write key. pub(crate) fn metadata(self, payload: Value) -> ProviderMetadata { - provider_namespace(self.write_keys()[0], payload) - .expect("provider metadata must be an object") + let mut metadata = provider_namespace(self.write_keys()[0], payload.clone()) + .expect("provider metadata must be an object"); + for key in &self.write_keys()[1..] { + metadata.extend( + provider_namespace(key, payload.clone()) + .expect("provider metadata must be an object"), + ); + } + metadata } } @@ -85,16 +100,3 @@ pub(crate) fn google_options( pub(crate) fn google_metadata(payload: Value) -> ProviderMetadata { Namespace::Google.metadata(payload) } - -/// The options under the Vertex key only (`googleVertex`), for the Vertex -/// surfaces that never fall back to `google`. -pub(crate) fn vertex_options( - provider_options: Option<&SharedProviderOptions>, -) -> Option<&JsonObject> { - provider_options?.get(GOOGLE_VERTEX) -} - -/// Response metadata under the Vertex key, as a metadata map. -pub(crate) fn vertex_metadata_map(payload: &Value) -> ProviderMetadata { - Namespace::Vertex.metadata(payload.clone()) -} diff --git a/aimux-providers/src/google/types.rs b/aimux-providers/src/google/types.rs index 056376d8e..0abbcb4be 100644 --- a/aimux-providers/src/google/types.rs +++ b/aimux-providers/src/google/types.rs @@ -32,7 +32,7 @@ pub struct GenerateContentResponse { } /// A single candidate in the response. -#[derive(Debug, Deserialize)] +#[derive(Debug, Default, Deserialize)] #[serde(rename_all = "camelCase")] pub struct Candidate { #[serde(default)] @@ -87,6 +87,14 @@ pub struct GoogleUsageMetadata { pub prompt_tokens_details: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub candidates_tokens_details: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_use_prompt_token_count: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cache_tokens_details: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_use_prompt_tokens_details: Option>, + #[serde(flatten)] + pub extra: std::collections::BTreeMap, } #[derive(Debug, Deserialize, Serialize)] @@ -96,6 +104,8 @@ pub struct TokenDetail { pub modality: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub token_count: Option, + #[serde(flatten)] + pub extra: std::collections::BTreeMap, } // ── Streaming chunk (`streamGenerateContent?alt=sse`) ──────────────────────── diff --git a/aimux-providers/src/google/utils.rs b/aimux-providers/src/google/utils.rs index a6c9425f8..1e57b5861 100644 --- a/aimux-providers/src/google/utils.rs +++ b/aimux-providers/src/google/utils.rs @@ -33,49 +33,15 @@ pub fn is_supported_file_url(url: &str) -> bool { return true; } - // YouTube watch URLs: https://(www.)youtube.com/watch?v=...(&...)? - if let Some(rest) = url - .strip_prefix("https://") - .and_then(|r| r.strip_prefix("www.").or(Some(r))) - && let Some(query) = rest.strip_prefix("youtube.com/watch") - { - // Must be followed by `?v=...` (the query string). A bare - // `/watch` with no query, or `/watch/...`, is not a video URL. - if let Some(q) = query.strip_prefix('?') { - return youtube_query_has_v(q); - } - return false; - } - - // youtu.be short URLs: https://youtu.be/(?...)? - if let Some(rest) = url.strip_prefix("https://youtu.be/") { - // The video id is everything up to `?` or end. - let id = rest.split('?').next().unwrap_or(rest); - return is_valid_youtube_id(id); - } - - false -} - -/// Check that a YouTube watch query string starts with `v=`. -fn youtube_query_has_v(query: &str) -> bool { - // The first parameter must be `v=` where id is non-empty and matches - // `[\w-]+`. Additional params may follow separated by `&`. - let mut params = query.split('&'); - let Some(first) = params.next() else { - return false; - }; - let Some(value) = first.strip_prefix("v=") else { - return false; - }; - is_valid_youtube_id(value) -} - -fn is_valid_youtube_id(id: &str) -> bool { - !id.is_empty() - && id - .bytes() - .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-') + static YOUTUBE: std::sync::OnceLock = std::sync::OnceLock::new(); + YOUTUBE + .get_or_init(|| { + regex::Regex::new( + r"^https://(?:(?:www\.)?youtube\.com/watch\?v=[\w-]+(?:&[\w=&.-]*)?|youtu\.be/[\w-]+(?:\?[\w=&.-]*)?)$", + ) + .expect("static YouTube URL pattern") + }) + .is_match(url) } // ── JSON accumulator ───────────────────────────────────────────────────────── @@ -575,22 +541,3 @@ fn parent_must_be(seg: &Segment) -> AiMuxError { "partial args path type conflict at {seg:?}: accumulated value does not match the path" )) } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn smoke() { - let mut acc = GoogleJsonAccumulator::new(); - let r = acc - .process_partial_args(&[PartialArg { - json_path: "$.location".to_string(), - string_value: Some("Boston".to_string()), - ..Default::default() - }]) - .unwrap(); - assert_eq!(r.text_delta, "{\"location\":\"Boston\""); - assert_eq!(r.current_json["location"], "Boston"); - } -} diff --git a/aimux-providers/src/google/video.rs b/aimux-providers/src/google/video.rs index 70c66cffb..c0598ba70 100644 --- a/aimux-providers/src/google/video.rs +++ b/aimux-providers/src/google/video.rs @@ -43,36 +43,124 @@ impl VideoModel for GoogleVideoModel { &self.model_id } fn max_videos_per_call(&self) -> Option { - Some(1) + Some(4) } async fn do_start( &self, options: &VideoCallOptions, ) -> Result { - let mut instances = vec![json!({"prompt": options.prompt})]; - if let Some(ref image) = options.image - && let aimux_core::video_model::VideoFile::Url { url, .. } = image + use aimux_core::video_model::VideoFrameType; + let mut warnings = Vec::new(); + let google = super::options::google_options(Some(&options.provider_options)); + if let Some(google) = google { + for (key, value) in google.iter().filter(|(_, v)| !v.is_null()) { + let valid = match key.as_str() { + "pollIntervalMs" | "pollTimeoutMs" => value.as_f64().is_some_and(|v| v > 0.0), + "personGeneration" => matches!( + value.as_str(), + Some("dont_allow" | "allow_adult" | "allow_all") + ), + "negativePrompt" => value.is_string(), + "referenceImages" => value.as_array().is_some_and(|images| { + images.iter().all(|image| { + image.as_object().is_some_and(|image| { + ["bytesBase64Encoded", "gcsUri"].iter().all(|key| { + image.get(*key).is_none_or(|v| v.is_null() || v.is_string()) + }) + }) + }) + }), + _ => true, + }; + if !valid { + return Err(AiMuxError::InvalidArgument(format!( + "Invalid Google video option: {key}" + ))); + } + } + } + let mut instance = Map::new(); + if let Some(prompt) = &options.prompt { + instance.insert("prompt".into(), json!(prompt)); + } + let frames = options.frame_images.as_deref().unwrap_or_default(); + let first = frames + .iter() + .find(|f| f.frame_type == VideoFrameType::FirstFrame) + .map(|f| &f.image) + .or(options.image.as_ref()); + if let Some(image) = first.and_then(|f| convert_image(f, &mut warnings)) { + instance.insert("image".into(), image); + } + if let Some(image) = frames + .iter() + .find(|f| f.frame_type == VideoFrameType::LastFrame) + .and_then(|f| convert_image(&f.image, &mut warnings)) { - instances[0]["image"] = json!({"gcsUri": url, "mimeType": "image/png"}); + instance.insert("lastFrame".into(), image); } - + if frames.is_empty() + && options + .input_references + .as_ref() + .is_some_and(|r| !r.is_empty()) + { + let references: Vec = options + .input_references + .iter() + .flatten() + .filter_map(|file| { + convert_image(file, &mut warnings) + .map(|image| json!({"image": image, "referenceType": "asset"})) + }) + .collect(); + instance.insert("referenceImages".into(), json!(references)); + } else if let Some(references) = google + .and_then(|g| g.get("referenceImages")) + .and_then(Value::as_array) + { + instance.insert("referenceImages".into(), json!(references.iter().map(|image| { + if let Some(bytes) = image.get("bytesBase64Encoded").and_then(Value::as_str).filter(|s| !s.is_empty()) { json!({"image": {"bytesBase64Encoded": bytes, "mimeType": "image/png"}, "referenceType": "asset"}) } + else if let Some(uri) = image.get("gcsUri").and_then(Value::as_str).filter(|s| !s.is_empty()) { json!({"image": {"gcsUri": uri, "mimeType": "image/png"}, "referenceType": "asset"}) } + else { image.clone() } + }).collect::>())); + } + let instances = vec![Value::Object(instance)]; let mut parameters = Map::new(); - parameters.insert("sampleCount".to_string(), json!(options.n)); + parameters.insert("sampleCount".into(), json!(options.n)); if let Some(ar) = options.aspect_ratio { - parameters.insert("aspectRatio".to_string(), json!(ar.to_string())); + parameters.insert("aspectRatio".into(), json!(ar.to_string())); } - if let Some(seed) = options.seed { - parameters.insert("seed".to_string(), json!(seed)); + if let Some(resolution) = options.resolution { + let resolution = resolution.to_string(); + parameters.insert( + "resolution".into(), + json!(match resolution.as_str() { + "1280x720" => "720p", + "1920x1080" => "1080p", + "3840x2160" => "4k", + _ => &resolution, + }), + ); } - if let Some(duration) = options.duration { - parameters.insert("durationSeconds".to_string(), json!(duration)); + if let Some(seed) = options.seed.filter(|v| *v != 0) { + parameters.insert("seed".into(), json!(seed)); } - if let Some(fps) = options.fps { - parameters.insert("fps".to_string(), json!(fps)); + if let Some(duration) = options.duration.filter(|v| *v != 0) { + parameters.insert("durationSeconds".into(), json!(duration)); } - if let Some(ga) = options.generate_audio { - parameters.insert("generateAudio".to_string(), json!(ga)); + if let Some(google) = google { + for (key, value) in google { + if !matches!( + key.as_str(), + "pollIntervalMs" | "pollTimeoutMs" | "referenceImages" + ) && !(matches!(key.as_str(), "personGeneration" | "negativePrompt") + && value.is_null()) + { + parameters.insert(key.clone(), value.clone()); + } + } } let body = json!({ @@ -96,14 +184,15 @@ impl VideoModel for GoogleVideoModel { let operation_name = predict_response .get("name") .and_then(|v| v.as_str()) + .filter(|name| !name.is_empty()) .ok_or_else(|| { AiMuxError::InvalidResponseData("No operation name returned from API".to_string()) })? .to_string(); Ok(VideoOperationStart { - operation: json!({ "operation_name": operation_name }), - warnings: Vec::new(), + operation: json!({ "operationName": operation_name }), + warnings, provider_metadata: None, response: VideoResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), @@ -119,11 +208,11 @@ impl VideoModel for GoogleVideoModel { options: &VideoCallOptions, ) -> Result { let operation_name = operation - .get("operation_name") + .get("operationName") .and_then(Value::as_str) .ok_or_else(|| { AiMuxError::InvalidArgument( - "google operation reference is missing operation_name".to_string(), + "google operation reference is missing operationName".to_string(), ) })?; @@ -142,7 +231,10 @@ impl VideoModel for GoogleVideoModel { if raw_body.get("done").and_then(serde_json::Value::as_bool) != Some(true) { return Ok(VideoOperationStatus::Pending); } - if let Some(err) = raw_body.get("error") { + + // Check the in-band error first: a terminal response may carry both + // done:true and an error object (provider-declared failure). + if let Some(err) = raw_body.get("error").filter(|v| !v.is_null()) { let msg = err .get("message") .and_then(|v| v.as_str()) @@ -157,52 +249,39 @@ impl VideoModel for GoogleVideoModel { ..ApiCallError::new(msg, poll_url, serde_json::json!({})) }))); } - let outputs = raw_body .get("response") .and_then(|r| r.get("generateVideoResponse")) .and_then(|r| r.get("generatedSamples")) .and_then(Value::as_array) + .filter(|outputs| !outputs.is_empty()) .ok_or_else(|| { AiMuxError::InvalidResponseData(format!( "No videos in response. Response: {raw_body}" )) })?; - if outputs.is_empty() { - return Err(AiMuxError::InvalidResponseData(format!( - "No videos in response. Response: {raw_body}" - ))); - } - let resolved = self.config.exchange(None).await?; - let api_key = resolved - .headers() - .into_iter() - .find_map(|(name, value)| name.eq_ignore_ascii_case("x-goog-api-key").then_some(value)); - let mut videos = Vec::new(); + let endpoint = (self.config.endpoint)().await?; + let api_key = endpoint + .headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("x-goog-api-key")) + .and_then(|(_, value)| value.as_deref()); let mut video_metadata = Vec::new(); - for sample in outputs { - if let Some(uri) = sample - .get("video") - .and_then(|v| v.get("uri")) - .and_then(Value::as_str) - .filter(|s| !s.is_empty()) - { - let url = if let Some(key) = api_key.as_ref().filter(|k| !k.is_empty()) - && aimux_provider_utils::same_origin(uri, resolved.base_url()) - { - let separator = if uri.contains('?') { '&' } else { '?' }; - format!("{uri}{separator}key={key}") - } else { - uri.to_string() - }; - videos.push(VideoData::Url { - url, - media_type: "video/mp4".to_string(), - }); + let videos: Vec = outputs + .iter() + .filter_map(|sample| { + let uri = sample + .get("video")? + .get("uri")? + .as_str() + .filter(|uri| !uri.is_empty())?; video_metadata.push(json!({"uri": uri})); - } - } - + Some(VideoData::Url { + url: authenticated_video_url(uri, &endpoint.base_url, api_key), + media_type: "video/mp4".to_string(), + }) + }) + .collect(); if videos.is_empty() { return Err(AiMuxError::InvalidResponseData( "No valid videos in response".to_string(), @@ -212,16 +291,9 @@ impl VideoModel for GoogleVideoModel { Ok(VideoOperationStatus::Completed(VideoResult { videos, warnings: Vec::new(), - provider_metadata: Some( - [( - crate::google::options::GOOGLE.to_string(), - [("videos".to_string(), json!(video_metadata))] - .into_iter() - .collect(), - )] - .into_iter() - .collect(), - ), + provider_metadata: Some(super::options::google_metadata( + json!({"videos": video_metadata}), + )), response: VideoResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), model_id: Some(self.model_id.clone()), @@ -230,3 +302,45 @@ impl VideoModel for GoogleVideoModel { })) } } + +fn convert_image( + file: &aimux_core::video_model::VideoFile, + warnings: &mut Vec, +) -> Option { + use aimux_core::video_model::{VideoFile, VideoFileData}; + match file { + VideoFile::Url { url, .. } if url.starts_with("gs://") => { + Some(json!({"gcsUri": url, "mimeType": "image/png"})) + } + VideoFile::Url { .. } => { + warnings.push(aimux_core::shared::Warning::Unsupported { feature: "URL-based image input".into(), details: Some("Google Generative AI video models require base64-encoded images or GCS URIs. URL will be ignored.".into()) }); + None + } + VideoFile::File { data, media_type } => { + let data = match data { + VideoFileData::Base64(data) => data.clone(), + VideoFileData::Binary(data) => { + base64::Engine::encode(&base64::engine::general_purpose::STANDARD, data) + } + }; + Some( + json!({"bytesBase64Encoded": data, "mimeType": if media_type.is_empty() { "image/png" } else { media_type }}), + ) + } + } +} + +fn authenticated_video_url(uri: &str, base_url: &str, key: Option<&str>) -> String { + let same_origin = url::Url::parse(uri) + .ok() + .zip(url::Url::parse(base_url).ok()) + .is_some_and(|(uri, base)| uri.origin() == base.origin()); + if let Some(key) = key.filter(|_| same_origin) { + format!( + "{uri}{}key={key}", + if uri.contains('?') { "&" } else { "?" } + ) + } else { + uri.into() + } +} diff --git a/aimux-providers/src/groq/convert.rs b/aimux-providers/src/groq/convert.rs index e8594f711..a2914bc99 100644 --- a/aimux-providers/src/groq/convert.rs +++ b/aimux-providers/src/groq/convert.rs @@ -42,11 +42,16 @@ fn convert_user_part(part: &UserPart) -> Result { "text file parts".to_string(), )); } - _ if !is_image(&file.media_type) => return Err(non_image()), - FileData::Url { url, .. } => json!({ - "type": "image_url", "image_url": { "url": url }, - }), + FileData::Url { url, .. } => { + if !is_image(&file.media_type) { + return Err(non_image()); + } + json!({ "type": "image_url", "image_url": { "url": url } }) + } FileData::Data { data } => { + if !is_image(&file.media_type) { + return Err(non_image()); + } let data = match data { FileBytes::Binary(bytes) => { base64::engine::general_purpose::STANDARD.encode(bytes) @@ -65,10 +70,10 @@ fn convert_assistant_message(content: &[AssistantPart]) -> Value { let mut tool_calls = Vec::new(); for part in content { match part { + // Groq supports reasoning for tool-calls in multi-turn conversations. AssistantPart::Text(TextPart { text: part_text, .. }) => text.push_str(part_text), - // Groq supports reasoning for tool-calls in multi-turn conversations. AssistantPart::Reasoning(ReasoningPart { text, .. }) => reasoning.push_str(text), AssistantPart::ToolCall(ToolCallPart { tool_call_id, diff --git a/aimux-providers/src/groq/model.rs b/aimux-providers/src/groq/model.rs index a466543fa..04fdd0714 100644 --- a/aimux-providers/src/groq/model.rs +++ b/aimux-providers/src/groq/model.rs @@ -339,12 +339,12 @@ impl LanguageModel for GroqChatLanguageModel { usage: convert_groq_usage(response.usage.as_ref()), warnings, provider_metadata: None, - request: Some(RequestInfo { body: Some(body) }), response: Some(ResponseInfo { headers: Some(response_headers), body: resp.raw_value, ..response_metadata(response.id, response.created, response.model).into() }), + request: Some(RequestInfo { body: Some(body) }), }) } @@ -365,24 +365,6 @@ impl LanguageModel for GroqChatLanguageModel { let response_headers = resp.response_headers; let mut sse_stream = resp.value; - // Only the first event is checked before returning the stream: an - // error reported immediately stays inside Core's operation-retry - // boundary. A normal event is chained back and never consumed. - let first_event = match sse_stream.next().await { - Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), - first => first, - }; - if let Some(Ok(event)) = &first_event - && event.get("error").is_some() - { - return Err(create_groq_stream_error( - event, - &endpoint, - body.clone(), - response_headers.clone(), - )); - } - let emit_raw_chunks = options.include_raw_chunks == Some(true); let stream_error_url = endpoint; let stream_error_body = body.clone(); @@ -405,10 +387,7 @@ impl LanguageModel for GroqChatLanguageModel { let mut is_active_text = false; let mut is_active_reasoning = false; - let mut event_iter = - futures::stream::iter(first_event.into_iter()).chain(sse_stream); - - while let Some(event) = event_iter.next().await { + while let Some(event) = sse_stream.next().await { let parsed = match event { Ok(parsed) => parsed, // a chunk that fails to parse is reported as an error part diff --git a/aimux-providers/src/groq/prepare_tools.rs b/aimux-providers/src/groq/prepare_tools.rs index 0087f028c..aa1aa66ff 100644 --- a/aimux-providers/src/groq/prepare_tools.rs +++ b/aimux-providers/src/groq/prepare_tools.rs @@ -63,16 +63,13 @@ pub(crate) fn prepare_tools(options: &CallOptions, model_id: &str) -> PreparedTo } prepared.tools = Some(groq_tools); - prepared.tool_choice = options - .tool_choice - .as_ref() - .map(|tool_choice| match tool_choice { - ToolChoice::Auto => json!("auto"), - ToolChoice::None => json!("none"), - ToolChoice::Required => json!("required"), - ToolChoice::Tool { tool_name } => { - json!({ "type": "function", "function": { "name": tool_name } }) - } - }); + prepared.tool_choice = options.tool_choice.as_ref().map(|choice| match choice { + ToolChoice::Auto => json!("auto"), + ToolChoice::None => json!("none"), + ToolChoice::Required => json!("required"), + ToolChoice::Tool { tool_name } => { + json!({ "type": "function", "function": { "name": tool_name } }) + } + }); prepared } diff --git a/aimux-providers/src/huggingface/options.rs b/aimux-providers/src/huggingface/options.rs index da78f824c..d9acc744b 100644 --- a/aimux-providers/src/huggingface/options.rs +++ b/aimux-providers/src/huggingface/options.rs @@ -20,5 +20,5 @@ pub(crate) fn huggingface_options( /// `{ "huggingface": payload }`. pub(crate) fn huggingface_metadata(payload: Value) -> ProviderMetadata { - provider_namespace(NAMESPACE, payload).expect("provider metadata must be an object") + provider_namespace(NAMESPACE, payload).expect("metadata payload is an object") } diff --git a/aimux-providers/src/huggingface/responses.rs b/aimux-providers/src/huggingface/responses.rs index 818ba840f..e4b4614d5 100644 --- a/aimux-providers/src/huggingface/responses.rs +++ b/aimux-providers/src/huggingface/responses.rs @@ -755,8 +755,8 @@ pub fn convert_to_huggingface_responses_messages( AssistantPart::ToolCall(_) | AssistantPart::ToolResult(_) | AssistantPart::File(_) - | AssistantPart::Custom(_) - | AssistantPart::ReasoningFile(_) => {} + | AssistantPart::ReasoningFile(_) + | AssistantPart::Custom(_) => {} } } } diff --git a/aimux-providers/src/lib.rs b/aimux-providers/src/lib.rs index 4e76f48a3..86707a17a 100644 --- a/aimux-providers/src/lib.rs +++ b/aimux-providers/src/lib.rs @@ -239,9 +239,10 @@ pub use openai::{ create_openai, }; pub use vertex::{ - VertexAnthropicModel, VertexEmbeddingModel, VertexImageModel, VertexModel, VertexProvider, - VertexProviderSettings, VertexTranscriptionModel, VertexVideoModel, create_google_vertex, - google_vertex, + GoogleAuthOptions, GoogleAuthScopes, VertexAnthropicModel, VertexAnthropicProvider, + VertexAnthropicProviderSettings, VertexEmbeddingModel, VertexImageModel, VertexModel, + VertexProvider, VertexProviderSettings, VertexTranscriptionModel, VertexVideoModel, + create_google_vertex, create_google_vertex_anthropic, google_vertex, google_vertex_anthropic, }; pub use voyage::{ VoyageEmbeddingModel, VoyageProvider, VoyageProviderSettings, VoyageRerankingModel, diff --git a/aimux-providers/src/mistral/convert.rs b/aimux-providers/src/mistral/convert.rs index 221909126..ca1bb2e40 100644 --- a/aimux-providers/src/mistral/convert.rs +++ b/aimux-providers/src/mistral/convert.rs @@ -3,15 +3,15 @@ //! Mirrors the TS `convert-to-mistral-chat-messages.ts`, //! `mistral-prepare-tools.ts`, and `map-mistral-finish-reason.ts`. -use aimux_core::error::AiMuxError; +use aimux_core::AiMuxError; use aimux_core::language_model_message::{ - AssistantPart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolCallPart, ToolPart, - ToolResultOutput, ToolResultPart, UserPart, + AssistantPart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultPart, UserPart, }; use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; use aimux_core::shared::{FileBytes, FileData}; use aimux_core::tool::{FunctionTool, Tool}; -use aimux_core::types::{FinishReason, FinishReasonUnified}; +use aimux_core::types::{FinishReason, FinishReasonUnified, ReasoningEffort, Warning}; use aimux_provider_utils::{get_top_level_media_type, resolve_full_media_type}; use serde_json::{Value, json}; @@ -96,183 +96,122 @@ pub fn prepare_tools( // ── Message conversion ────────────────────────────────────────────────────── -/// Convert a `LanguageModelPrompt` to Mistral `messages` array. -/// -/// Differences from OpenAI: -/// - System content is a plain string. -/// - User content is always an array of typed parts. -/// - Assistant content is a plain string; `prefix: true` is set on the last -/// message if it is an assistant message (continuation mode). -/// - Tool messages include `tool_call_id` (no `name` — the Rust data model -/// does not carry the tool name on `ToolResult` parts). +/// Convert provider prompt parts using the upstream Mistral message format. /// /// # Errors -/// Returns an error for text file parts or an unresolved inline file media type. +/// Rejects file formats and assistant content unsupported by Mistral. pub fn convert_prompt_to_mistral_messages( prompt: &LanguageModelPrompt, ) -> Result, AiMuxError> { - let mut result = Vec::new(); - let last_idx = prompt.len().saturating_sub(1); - for (i, msg) in prompt.iter().enumerate() { - let is_last = i == last_idx; - for value in convert_message_to_mistral(msg, is_last)? { - result.push(value); - } - } - Ok(result) -} - -fn convert_message_to_mistral( - msg: &LanguageModelMessage, - is_last: bool, -) -> Result, AiMuxError> { - Ok(match msg { - LanguageModelMessage::System { content, .. } => { - vec![json!({ "role": "system", "content": content })] - } - LanguageModelMessage::User { content, .. } => { - let parts: Vec = content - .iter() - .map(convert_part_to_mistral) - .collect::>()?; - vec![json!({ "role": "user", "content": parts })] - } - LanguageModelMessage::Assistant { content, .. } => { - let text = join_text_parts(content); - let has_reasoning = content.iter().any(|part| matches!(part, AssistantPart::Reasoning(_))); - let mut content_parts = Vec::new(); - for part in content { - match part { - AssistantPart::Text(part) => content_parts.push(json!({ "type": "text", "text": part.text })), - AssistantPart::Reasoning(part) => content_parts.push(json!({ "type": "thinking", "thinking": [{ "type": "text", "text": part.text }], "closed": true })), - AssistantPart::ToolCall(_) => {}, - _ => return Err(AiMuxError::UnsupportedFunctionality("assistant content part".to_string())), - } + let mut messages = Vec::new(); + for (index, message) in prompt.iter().enumerate() { + match message { + LanguageModelMessage::System { content, .. } => { + messages.push(json!({"role":"system", "content": content})) } - let has_tool_calls = content - .iter() - .any(|p| matches!(p, AssistantPart::ToolCall(_))); - - let mut msg_json = json!({ "role": "assistant", "content": text }); - if has_reasoning { msg_json["content"] = json!(content_parts); } - - if has_tool_calls { - let tool_calls: Vec = content + LanguageModelMessage::User { content, .. } => { + let parts = content .iter() - .filter_map(|p| match p { - AssistantPart::ToolCall(ToolCallPart { - tool_call_id, - tool_name, - input, - .. - }) => { - let arguments = input.to_string(); - Some(json!({ - "id": tool_call_id, - "type": "function", - "function": { - "name": tool_name, - "arguments": arguments, - } - })) + .map(convert_user_part) + .collect::, _>>()?; + messages.push(json!({"role":"user", "content":parts})); + } + LanguageModelMessage::Assistant { content, .. } => { + let mut text = String::new(); + let mut parts = Vec::new(); + let mut calls = Vec::new(); + let mut has_reasoning = false; + for part in content { + match part { + AssistantPart::Text(part) => { + text.push_str(&part.text); + parts.push(json!({"type":"text", "text":part.text})); } - _ => None, - }) - .collect(); - msg_json["tool_calls"] = json!(tool_calls); + AssistantPart::Reasoning(ReasoningPart { text, .. }) => { + has_reasoning = true; + parts.push(json!({"type":"thinking", "thinking":[{"type":"text", "text":text}], "closed":true})); + } + AssistantPart::ToolCall(ToolCallPart { tool_call_id, tool_name, input, .. }) => calls.push(json!({ + "id":tool_call_id, "type":"function", "function":{"name":tool_name,"arguments":input.to_string()} + })), + _ => return Err(AiMuxError::UnsupportedFunctionality("Unsupported content type in assistant message".into())), + } + } + let content = if has_reasoning { + json!(parts) + } else { + json!(text) + }; + let mut value = json!({"role":"assistant", "content":content}); + if index + 1 == prompt.len() { + value["prefix"] = json!(true); + } + if !calls.is_empty() { + value["tool_calls"] = json!(calls); + } + messages.push(value); } - - if is_last { - msg_json["prefix"] = json!(true); + LanguageModelMessage::Tool { content, .. } => { + for part in content { + let ToolPart::ToolResult(ToolResultPart { + tool_call_id, + tool_name, + output, + .. + }) = part + else { + continue; + }; + messages.push(json!({"role":"tool", "name":tool_name, "tool_call_id":tool_call_id, "content":crate::openai::convert::tool_result_to_content(output)})); + } } - vec![msg_json] } - LanguageModelMessage::Tool { content, .. } => content - .iter() - .filter_map(|part| { - let ToolPart::ToolResult(ToolResultPart { tool_call_id, tool_name, output, .. }) = part else { return None; }; - Some(json!({ "role": "tool", "name": tool_name, "tool_call_id": tool_call_id, "content": tool_result_to_content(output) })) - }) - .collect(), - }) -} - -fn join_text_parts(content: &[AssistantPart]) -> String { - content - .iter() - .filter_map(|p| match p { - AssistantPart::Text(TextPart { text, .. }) => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join("") + } + Ok(messages) } -fn tool_result_to_content(output: &ToolResultOutput) -> Value { - Value::String(match output { - ToolResultOutput::Text { value, .. } | ToolResultOutput::ErrorText { value, .. } => { - value.clone() - } - ToolResultOutput::ExecutionDenied { reason, .. } => reason - .clone() - .unwrap_or_else(|| "Tool call execution denied.".to_string()), - ToolResultOutput::Json { value, .. } | ToolResultOutput::ErrorJson { value, .. } => { - value.to_string() +fn convert_user_part(part: &UserPart) -> Result { + use base64::Engine; + let file = match part { + UserPart::Text(TextPart { text, .. }) => return Ok(json!({"type":"text", "text":text})), + UserPart::File(file) => file, + }; + let url = match &file.data { + FileData::Reference { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "file parts with provider references".into(), + )); } - ToolResultOutput::Content { value } => { - crate::openai::convert::tool_result_content_value(value).to_string() + FileData::Text { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "text file parts".into(), + )); } - }) -} - -fn convert_part_to_mistral(part: &UserPart) -> Result { - Ok(match part { - UserPart::Text(TextPart { text, .. }) => json!({ "type": "text", "text": text }), - UserPart::File(file) => { - use base64::Engine; - let is_image = get_top_level_media_type(&file.media_type) == "image"; - let url = match &file.data { - FileData::Data { data } => { - let media_type = resolve_full_media_type(file)?; - if !is_image && media_type != "application/pdf" { - return Err(AiMuxError::UnsupportedFunctionality( - "Only images and PDF file parts are supported".to_string(), - )); - } - let b64 = match data { - FileBytes::Binary(bytes) => { - base64::engine::general_purpose::STANDARD.encode(bytes) - } - FileBytes::Base64(data) => data.clone(), - }; - format!("data:{media_type};base64,{b64}") - } - FileData::Url { url, .. } => { - if !is_image && file.media_type != "application/pdf" { - return Err(AiMuxError::UnsupportedFunctionality( - "Only images and PDF file parts are supported".to_string(), - )); - } - url.clone() - } - FileData::Reference { .. } => { - return Err(AiMuxError::UnsupportedFunctionality( - "file parts with provider references".into(), - )); - } - FileData::Text { .. } => { - return Err(AiMuxError::UnsupportedFunctionality( - "text file parts".into(), - )); - } + FileData::Data { data } => { + let full = resolve_full_media_type(file)?; + let data = match data { + FileBytes::Binary(bytes) => base64::engine::general_purpose::STANDARD.encode(bytes), + FileBytes::Base64(data) => data.clone(), }; - if is_image { - json!({ "type": "image_url", "image_url": url }) - } else { - json!({ "type": "document_url", "document_url": url }) - } + format!("data:{full};base64,{data}") + } + FileData::Url { url, .. } => url.clone(), + }; + if get_top_level_media_type(&file.media_type) == "image" { + Ok(json!({"type":"image_url", "image_url":url})) + } else { + let full = if matches!(file.data, FileData::Url { .. }) { + file.media_type.clone() + } else { + resolve_full_media_type(file)? + }; + if full != "application/pdf" { + return Err(AiMuxError::UnsupportedFunctionality( + "Only images and PDF file parts are supported".into(), + )); } - }) + Ok(json!({"type":"document_url", "document_url":url})) + } } // ── Request body ──────────────────────────────────────────────────────────── @@ -280,38 +219,31 @@ fn convert_part_to_mistral(part: &UserPart) -> Result { /// Convert `CallOptions` to a Mistral request body. /// /// # Errors -/// Returns an error for text file parts or an unresolved inline file media type. +/// Rejects invalid provider options and unsupported prompt parts. pub fn build_request_body( model_id: &str, options: &CallOptions, stream: bool, ) -> Result { - let provider_options = options - .provider_options - .as_ref() - .and_then(|namespaces| namespaces.get(crate::mistral::options::NAMESPACE)) - .map(|namespace| { - crate::openai::convert::parse_option_fields( - namespace, - crate::mistral::options::NAMESPACE, - |key, value| { - let valid = match key { - "safePrompt" | "structuredOutputs" | "strictJsonSchema" - | "parallelToolCalls" => value.is_boolean(), - "documentImageLimit" | "documentPageLimit" => value.is_number(), - "promptCacheKey" => value.is_string(), - "reasoningEffort" => value - .as_str() - .is_some_and(|value| matches!(value, "high" | "none")), - _ => return None, - }; - Some(if valid { Ok(value.clone()) } else { Err(()) }) - }, - ) - }) - .transpose()? - .unwrap_or_default(); - let messages = convert_prompt_to_mistral_messages(&options.prompt)?; + let provider_options = + super::options::validated_mistral_options(options.provider_options.as_ref())?; + let mut messages = convert_prompt_to_mistral_messages(&options.prompt)?; + if matches!( + options.response_format, + Some(ResponseFormat::Json { schema: None, .. }) + ) { + let instruction = "You MUST answer with JSON."; + if messages.first().is_some_and(|m| m["role"] == "system") { + let text = messages[0]["content"].as_str().unwrap_or_default(); + messages[0]["content"] = json!(if text.is_empty() { + instruction.to_owned() + } else { + format!("{text}\n\n{instruction}") + }); + } else { + messages.insert(0, json!({"role":"system", "content":instruction})); + } + } let mut body = json!({ "model": model_id, @@ -322,10 +254,6 @@ pub fn build_request_body( body["stream"] = json!(true); } - if let Some(safe_prompt) = provider_options.get("safePrompt") { - body["safe_prompt"] = safe_prompt.clone(); - } - if let Some(max_tokens) = options.max_output_tokens { body["max_tokens"] = json!(max_tokens); } @@ -357,7 +285,12 @@ pub fn build_request_body( name, description, } => { - if schema.is_some() { + if schema.is_some() + && provider_options + .and_then(|o| o.get("structuredOutputs")) + .and_then(Value::as_bool) + .unwrap_or(true) + { let mut schema_obj = json!({}); if let Some(s) = schema { schema_obj["schema"] = s.clone(); @@ -367,7 +300,12 @@ pub fn build_request_body( if let Some(d) = description { schema_obj["description"] = json!(d); } - schema_obj["strict"] = json!(false); + schema_obj["strict"] = json!( + provider_options + .and_then(|o| o.get("strictJsonSchema")) + .and_then(Value::as_bool) + .unwrap_or(false) + ); body["response_format"] = json!({ "type": "json_schema", "json_schema": schema_obj, @@ -379,27 +317,127 @@ pub fn build_request_body( } } - // Tools (delegated to `prepare_tools`). - let function_tools: Option> = options.tools.as_ref().map(|tools| { - tools - .iter() - .filter_map(|t| match t { - Tool::Function(ft) => Some(ft.clone()), - Tool::Provider(_) => None, - }) - .collect() - }); - let prepared = prepare_tools(&function_tools, options.tool_choice.as_ref()); - if let Some(tools) = prepared.tools { - body["tools"] = json!(tools); - if let Some(tc) = prepared.tool_choice { - body["tool_choice"] = tc; + if let Some(provider_options) = provider_options { + for (key, wire_key) in [ + ("safePrompt", "safe_prompt"), + ("documentImageLimit", "document_image_limit"), + ("documentPageLimit", "document_page_limit"), + ("promptCacheKey", "prompt_cache_key"), + ] { + if let Some(value) = provider_options.get(key) { + body[wire_key] = value.clone(); + } + } + } + if supports_reasoning_effort(model_id) { + let effort = provider_options + .and_then(|o| o.get("reasoningEffort")) + .cloned() + .or_else(|| { + options.reasoning.filter(|r| r.is_custom()).map(|r| { + json!(if r == ReasoningEffort::None { + "none" + } else { + "high" + }) + }) + }); + if let Some(effort) = effort { + body["reasoning_effort"] = effort; } } + // Preserve the empty tools array when the supplied tools are all unsupported. + if let Some(tools) = options.tools.as_ref().filter(|tools| !tools.is_empty()) { + let function_tools = Some( + tools + .iter() + .filter_map(|t| match t { + Tool::Function(ft) => Some(ft.clone()), + Tool::Provider(_) => None, + }) + .collect(), + ); + let mut prepared = prepare_tools(&function_tools, options.tool_choice.as_ref()); + if prepared.tools.is_none() { + prepared.tools = Some(Vec::new()); + prepared.tool_choice = options.tool_choice.as_ref().map(|choice| { + json!(match choice { + ToolChoice::Auto => "auto", + ToolChoice::None => "none", + _ => "any", + }) + }); + } + body["tools"] = json!(prepared.tools); + if let Some(choice) = prepared.tool_choice { + body["tool_choice"] = choice; + } + if let Some(value) = provider_options.and_then(|o| o.get("parallelToolCalls")) { + body["parallel_tool_calls"] = value.clone(); + } + } Ok(body) } +fn supports_reasoning_effort(model_id: &str) -> bool { + matches!( + model_id, + "glm-5-2" + | "labs-leanstral-1-5" + | "labs-leanstral-1-5-1" + | "magistral-medium-latest" + | "magistral-small-latest" + | "mistral-medium" + | "mistral-medium-2604" + | "mistral-medium-3" + | "mistral-medium-3-5" + | "mistral-medium-3.5" + | "mistral-medium-latest" + | "mistral-small-2603" + | "mistral-small-latest" + | "mistral-vibe-cli-fast" + | "mistral-vibe-cli-latest" + | "mistral-vibe-cli-with-tools" + | "zai-glm-5-2" + ) +} + +/// The warnings emitted while preparing an upstream Mistral request. +#[must_use] +pub fn request_warnings(options: &CallOptions, model_id: &str) -> Vec { + let mut warnings = Vec::new(); + if options.top_k.is_some() { + warnings.push(Warning::Unsupported { + feature: "topK".into(), + details: None, + }); + } + if let Some(reasoning) = options.reasoning.filter(|r| r.is_custom()) { + if !supports_reasoning_effort(model_id) { + warnings.push(Warning::Unsupported { + feature: "reasoning".into(), + details: Some("This model does not support reasoning configuration.".into()), + }); + } else if !matches!(reasoning, ReasoningEffort::None | ReasoningEffort::High) + && super::options::mistral_options(options.provider_options.as_ref()) + .and_then(|o| o.get("reasoningEffort")) + .is_none() + { + warnings.push(Warning::Compatibility { feature:"reasoning".into(), details:Some(format!("reasoning \"{reasoning}\" is not directly supported by this model. mapped to effort \"high\".")) }); + } + } + for tool in options.tools.iter().flatten() { + if let Tool::Provider(tool) = tool { + warnings.push(Warning::Unsupported { + feature: format!("provider-defined tool {}", tool.id), + details: None, + }); + } + } + warnings +} + /// Parse Mistral finish reason string into `FinishReason`. /// /// Differences from OpenAI: `model_length` is also mapped to `Length`. @@ -409,7 +447,6 @@ pub fn parse_finish_reason(s: &str) -> FinishReason { "stop" => FinishReasonUnified::Stop, "length" | "model_length" => FinishReasonUnified::Length, "tool_calls" => FinishReasonUnified::ToolCalls, - "content_filter" => FinishReasonUnified::ContentFilter, _ => FinishReasonUnified::Other, }; FinishReason { diff --git a/aimux-providers/src/mistral/embedding.rs b/aimux-providers/src/mistral/embedding.rs index e1dd790f1..2243690c7 100644 --- a/aimux-providers/src/mistral/embedding.rs +++ b/aimux-providers/src/mistral/embedding.rs @@ -50,7 +50,15 @@ impl EmbeddingModel for MistralEmbeddingModel { &self, options: &EmbeddingCallOptions, ) -> Result { - let mistral_options = parse_mistral_provider_options(options.provider_options.as_ref()); + if options.values.len() > 32 { + return Err(AiMuxError::InvalidArgument(format!( + "{} model {} supports at most 32 embeddings per call; received {}", + self.provider(), + self.model_id, + options.values.len() + ))); + } + let mistral_options = parse_mistral_provider_options(options.provider_options.as_ref())?; let mut body = Map::new(); body.insert("model".to_string(), json!(self.model_id)); @@ -123,23 +131,46 @@ impl EmbeddingModel for MistralEmbeddingModel { struct MistralEmbeddingProviderOptions { metadata: Option, - output_dimension: Option, + output_dimension: Option, output_dtype: Option, } fn parse_mistral_provider_options( options: Option<&SharedProviderOptions>, -) -> MistralEmbeddingProviderOptions { +) -> Result { let provider_opts = super::options::mistral_options(options); - MistralEmbeddingProviderOptions { - metadata: provider_opts.and_then(|o| o.get("metadata")).cloned(), - output_dimension: provider_opts - .and_then(|o| o.get("outputDimension")) - .and_then(serde_json::Value::as_u64) - .map(|d| d as u32), - output_dtype: provider_opts - .and_then(|o| o.get("outputDtype")) - .and_then(|d| d.as_str()) - .map(std::string::ToString::to_string), + let metadata = provider_opts.and_then(|o| o.get("metadata")).cloned(); + let output_dimension = provider_opts + .and_then(|o| o.get("outputDimension")) + .cloned(); + let output_dtype = provider_opts.and_then(|o| o.get("outputDtype")); + if metadata.as_ref().is_some_and(|value| !value.is_object()) { + return Err(AiMuxError::InvalidArgument( + "mistral.metadata must be a record".into(), + )); + } + if output_dimension.as_ref().is_some_and(|value| { + !value.as_f64().is_some_and(|number| { + number > 0.0 && number.fract() == 0.0 && number <= 9_007_199_254_740_991.0 + }) + }) { + return Err(AiMuxError::InvalidArgument( + "mistral.outputDimension must be a positive integer".into(), + )); + } + if output_dtype.is_some_and(|value| { + !matches!( + value.as_str(), + Some("float" | "int8" | "uint8" | "binary" | "ubinary") + ) + }) { + return Err(AiMuxError::InvalidArgument( + "mistral.outputDtype must be float, int8, uint8, binary, or ubinary".into(), + )); } + Ok(MistralEmbeddingProviderOptions { + metadata, + output_dimension, + output_dtype: output_dtype.and_then(Value::as_str).map(str::to_owned), + }) } diff --git a/aimux-providers/src/mistral/mod.rs b/aimux-providers/src/mistral/mod.rs index be1690333..34eaa3b0d 100644 --- a/aimux-providers/src/mistral/mod.rs +++ b/aimux-providers/src/mistral/mod.rs @@ -94,7 +94,6 @@ pub(crate) fn mistral_stream_error( const DEFAULT_BASE_URL: &str = "https://api.mistral.ai/v1"; const API_KEY_ENV_VAR: &str = "MISTRAL_API_KEY"; -const DEFAULT_NAME: &str = "mistral"; /// The URL patterns the chat model fetches itself (`supportedUrls` of the AI /// SDK's `MistralChatLanguageModel`): `https` PDFs. @@ -124,7 +123,7 @@ pub struct MistralProviderSettings { /// The transport: a mock, a signing decorator, a proxy-aware client. /// `None` uses the process default, resolved per request. pub fetch: Option, - /// Generates unique IDs for chat output. + /// Generates identifiers for reasoning and streamed tool calls. pub generate_id: Option String + Send + Sync>>, } @@ -156,22 +155,22 @@ pub fn create_mistral(settings: MistralProviderSettings) -> Result validate_base_url(url)?, None => DEFAULT_BASE_URL.to_string(), }; - Ok(MistralProvider { - name: DEFAULT_NAME.to_string(), - base_url, - headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( - provider_headers( - Credential::explicit_or_env( - settings.api_key.map(Resolvable::Value), - API_KEY_ENV_VAR, - "Mistral", - ), - Vec::new(), - settings.headers, + let headers = aimux_provider_utils::headers::with_user_agent_suffix_fn( + provider_headers( + Credential::explicit_or_env( + settings.api_key.map(Resolvable::Value), + API_KEY_ENV_VAR, + "Mistral", ), - options::NAMESPACE, - "4.0.54", + Vec::new(), + settings.headers, ), + options::NAMESPACE, + "4.0.54", + ); + Ok(MistralProvider { + base_url, + headers, fetch: settings.fetch, generate_id: settings.generate_id, }) @@ -191,7 +190,6 @@ pub fn mistral() -> &'static MistralProvider { /// A Mistral provider (the AI SDK's `MistralProvider`). Cheap to clone the /// models out of; it holds no HTTP client. pub struct MistralProvider { - name: String, base_url: String, headers: HeadersFn, fetch: Option, @@ -201,7 +199,7 @@ pub struct MistralProvider { impl MistralProvider { fn model_config(&self, method: &str) -> EndpointConfig { EndpointConfig::fixed( - format!("{}.{method}", self.name), + format!("mistral.{method}"), self.base_url.clone(), self.headers.clone(), self.fetch.clone(), @@ -209,19 +207,19 @@ impl MistralProvider { ) } - /// A chat model; `provider()` is `"{name}.chat"`. + /// A chat model; `provider()` is `"mistral.chat"`. #[must_use] pub fn chat(&self, model_id: &str) -> MistralModel { MistralModel::from_config( model_id.to_string(), self.model_config("chat") .with_supported_urls(Arc::new(|_| chat_supported_urls())), + self.generate_id.clone(), ) - .with_generate_id(self.generate_id.clone()) } /// An embedding model (e.g. `"mistral-embed"`); `provider()` is - /// `"{name}.embedding"`. + /// `"mistral.embedding"`. #[must_use] pub fn embedding(&self, model_id: &str) -> MistralEmbeddingModel { MistralEmbeddingModel::from_config(model_id.to_string(), self.model_config("embedding")) diff --git a/aimux-providers/src/mistral/model.rs b/aimux-providers/src/mistral/model.rs index 125009937..703ce3c64 100644 --- a/aimux-providers/src/mistral/model.rs +++ b/aimux-providers/src/mistral/model.rs @@ -13,6 +13,7 @@ use aimux_core::tool::RawToolCall; use async_trait::async_trait; use futures::StreamExt; use serde_json::Value; +use std::sync::Arc; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; @@ -25,32 +26,27 @@ use aimux_provider_utils::{StreamingToolCallDelta, StreamingToolCallTracker}; use crate::shared::EndpointConfig; -use super::convert::{build_request_body, parse_finish_reason}; +use super::convert::{build_request_body, parse_finish_reason, request_warnings}; use super::types::{ChatCompletionResponse, StreamChunk, UsageResponse}; /// A Mistral language model. pub struct MistralModel { model_id: String, config: EndpointConfig, - generate_id: std::sync::Arc String + Send + Sync>, + generate_id: Arc String + Send + Sync>, } impl MistralModel { - pub(crate) fn with_generate_id( - mut self, - generate_id: Option String + Send + Sync>>, + pub(crate) fn from_config( + model_id: String, + config: EndpointConfig, + generate_id_fn: Option String + Send + Sync>>, ) -> Self { - if let Some(generate_id) = generate_id { - self.generate_id = generate_id; - } - self - } - - pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config, - generate_id: std::sync::Arc::new(aimux_provider_utils::generate_id), + generate_id: generate_id_fn + .unwrap_or_else(|| Arc::new(aimux_provider_utils::generate_id)), } } } @@ -65,8 +61,10 @@ impl MistralModel { /// - `input.cacheRead = cache_read_tokens` (or undefined when 0) /// - `output.total = completion_tokens` fn convert_usage(usage: &UsageResponse) -> Usage { - let prompt_tokens = usage.prompt_tokens.unwrap_or(0); - let completion_tokens = usage.completion_tokens.unwrap_or(0); + let raw = usage.raw.clone(); + let usage = &usage.fields; + let prompt_tokens = usage.prompt_tokens; + let completion_tokens = usage.completion_tokens; let cache_read = usage .num_cached_tokens @@ -99,12 +97,11 @@ fn convert_usage(usage: &UsageResponse) -> Usage { }, output_tokens: aimux_core::types::OutputTokenUsage { total: Some(completion_tokens), + text: Some(completion_tokens), ..Default::default() }, // RFC-0015 P0-3: keep the raw provider usage payload. - raw: serde_json::to_value(usage) - .ok() - .and_then(|value| value.as_object().cloned()), + raw: raw.as_object().cloned(), } } @@ -144,46 +141,19 @@ fn extract_text_content(content: &Option) -> Option { } } -/// Extract reasoning content from a Mistral `content` array's `thinking` parts. -/// -/// Each `thinking` part has a `thinking` array of `{type:"text", text}` chunks. -fn extract_reasoning_content(content: &Option) -> Option { - match content { - Some(Value::Array(arr)) => { - let text: String = arr - .iter() - .filter_map(|part| { - if part.get("type").and_then(|t| t.as_str()) == Some("thinking") { - part.get("thinking") - .and_then(|t| t.as_array()) - .map(|thinking| { - thinking - .iter() - .filter_map(|chunk| { - if chunk.get("type").and_then(|t| t.as_str()) - == Some("text") - { - chunk - .get("text") - .and_then(|t| t.as_str()) - .map(std::string::ToString::to_string) - } else { - None - } - }) - .collect::>() - .join("") - }) - } else { - None - } - }) - .collect::>() - .join(""); - if text.is_empty() { None } else { Some(text) } - } - _ => None, +/// Extract the text chunks of one thinking part. +fn extract_reasoning_part(part: &Value) -> Option { + if part.get("type").and_then(Value::as_str) != Some("thinking") { + return None; } + let text: String = part + .get("thinking")? + .as_array()? + .iter() + .filter(|chunk| chunk.get("type").and_then(Value::as_str) == Some("text")) + .filter_map(|chunk| chunk.get("text").and_then(Value::as_str)) + .collect(); + (!text.is_empty()).then_some(text) } // ── LanguageModel impl ─────────────────────────────────────────────────────── @@ -214,7 +184,6 @@ impl LanguageModel for MistralModel { ) .await?; - let response_body = resp.raw_value; let response_headers = resp.response_headers; let data: ChatCompletionResponse = resp.value; @@ -226,21 +195,33 @@ impl LanguageModel for MistralModel { // Build content array. let mut content = Vec::new(); - // Reasoning content (from thinking parts in array content). - // Pushed before text to match upstream ordering (reasoning → text). - if let Some(r) = extract_reasoning_content(&choice.message.content) - && !r.is_empty() - { - content.push(GenerateContent::Reasoning(ReasoningOutput { - text: r, - provider_metadata: None, - })); - } - - // Content can be a string (legacy) or an array of typed parts. - if let Some(text) = extract_text_content(&choice.message.content) - && !text.is_empty() - { + if let Some(Value::Array(parts)) = &choice.message.content { + for part in parts { + match part.get("type").and_then(Value::as_str) { + Some("thinking") => { + if let Some(text) = extract_reasoning_part(part) { + content.push(GenerateContent::Reasoning(ReasoningOutput { + text, + provider_metadata: None, + })); + } + } + Some("text") => { + if let Some(text) = part + .get("text") + .and_then(Value::as_str) + .filter(|text| !text.is_empty()) + { + content.push(GenerateContent::Text { + text: text.to_string(), + provider_metadata: None, + }); + } + } + _ => {} + } + } + } else if let Some(text) = extract_text_content(&choice.message.content) { content.push(GenerateContent::Text { text, provider_metadata: None, @@ -277,7 +258,7 @@ impl LanguageModel for MistralModel { content, finish_reason, usage, - warnings: Vec::new(), + warnings: request_warnings(options, &self.model_id), provider_metadata: None, response: Some(aimux_core::shared::ResponseInfo { id: data.id, @@ -287,7 +268,7 @@ impl LanguageModel for MistralModel { .map(|dt| dt.to_rfc3339()), model_id: data.model, headers: Some(response_headers), - body: response_body, + body: resp.raw_value, }), request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), }) @@ -308,34 +289,21 @@ impl LanguageModel for MistralModel { let response_headers = resp.response_headers; let mut sse_stream = resp.value; - // Peek at the first SSE event to detect early errors. - let first_event = match sse_stream.next().await { - Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), - first_event => first_event, - }; - if let Some(Ok(ref event)) = first_event - && let Some(err_obj) = event.get("error") - { - return Err(super::mistral_stream_error( - err_obj, - &endpoint, - body.clone(), - response_headers.clone(), - )); - } + let warnings = request_warnings(options, &self.model_id); + let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); let stream_error_url = endpoint; let stream_error_body = body.clone(); let stream_response_headers = response_headers.clone(); - let generate_id = self.generate_id.clone(); - let tool_generate_id = generate_id.clone(); + let generate_id_fn = self.generate_id.clone(); let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings: vec![] }); + yield Ok(StreamPart::StreamStart { warnings }); let text_id = 0usize; let mut text_started = false; let mut reasoning_started = false; + let tool_generate_id = generate_id_fn.clone(); let mut tool_calls = StreamingToolCallTracker::new().with_generate_id(move || tool_generate_id()); let mut tool_parts = Vec::new(); let mut reasoning_id: Option = None; @@ -343,18 +311,13 @@ impl LanguageModel for MistralModel { let mut final_finish_reason: Option = None; let mut response_metadata_emitted = false; - let mut event_iter = - futures::stream::iter(first_event.into_iter()).chain(sse_stream); - - let mut stream_errored = false; - - while let Some(event) = event_iter.next().await { - if stream_errored { - break; - } + while let Some(event) = sse_stream.next().await { match event { Ok(parsed) => { + if include_raw_chunks { + yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); + } if let Some(err_obj) = parsed.get("error") { yield Ok(StreamPart::Error { @@ -365,22 +328,19 @@ impl LanguageModel for MistralModel { stream_response_headers.clone(), ), }); - stream_errored = true; - break; + continue; } let chunk: StreamChunk = match serde_json::from_value(parsed) { Ok(c) => c, Err(e) => { - yield Err(e.into()); + yield Ok(StreamPart::Error { error: e.into() }); continue; } }; // Emit ResponseMetadata from the first valid chunk. - if !response_metadata_emitted - && (chunk.id.is_some() || chunk.model.is_some()) - { + if !response_metadata_emitted { response_metadata_emitted = true; yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { id: chunk.id.clone(), @@ -402,9 +362,8 @@ impl LanguageModel for MistralModel { // Upstream processes only the first choice. if let Some(choice) = chunk.choices.into_iter().next() { // Reasoning content (from thinking parts in array content). - if let Some(reasoning_delta) = - extract_reasoning_content(&choice.delta.content) - && !reasoning_delta.is_empty() { + for reasoning_delta in choice.delta.content.as_ref().and_then(Value::as_array).into_iter().flatten() + .filter_map(extract_reasoning_part) { if !reasoning_started { // End any active text before starting reasoning. if text_started { @@ -414,7 +373,7 @@ impl LanguageModel for MistralModel { }); text_started = false; } - let rid = generate_id(); + let rid = generate_id_fn(); reasoning_id = Some(rid.clone()); reasoning_started = true; yield Ok(StreamPart::ReasoningStart { id: rid, @@ -484,33 +443,15 @@ impl LanguageModel for MistralModel { // Finish reason. if let Some(reason) = choice.finish_reason { - if text_started { - yield Ok(StreamPart::TextEnd { - id: format!("{text_id}"), - provider_metadata: None, - }); - text_started = false; - } - if reasoning_started { - if let Some(ref rid) = reasoning_id { - yield Ok(StreamPart::ReasoningEnd { - id: rid.clone(), - provider_metadata: None, - }); - } - reasoning_started = false; - reasoning_id = None; - } final_finish_reason = Some(parse_finish_reason(&reason)); } } } Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; + if include_raw_chunks { + yield Ok(StreamPart::Raw { raw_value: Value::Null }); } + yield Ok(StreamPart::Error { error }); } } } @@ -536,23 +477,12 @@ impl LanguageModel for MistralModel { } yield Ok(StreamPart::Finish { - finish_reason: if stream_errored { - FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - } - } else { - final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }) - }, - usage: if stream_errored { - Usage::default() - } else { - final_usage - }, - provider_metadata: Some(super::options::mistral_metadata(serde_json::json!({}))), + finish_reason: final_finish_reason.unwrap_or(FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }), + usage: final_usage, + provider_metadata: None, }); }; diff --git a/aimux-providers/src/mistral/options.rs b/aimux-providers/src/mistral/options.rs index 8096e700b..9e7b1d3dd 100644 --- a/aimux-providers/src/mistral/options.rs +++ b/aimux-providers/src/mistral/options.rs @@ -4,9 +4,7 @@ //! metadata under `mistral`, whatever the provider is named. This module is //! the one place that spells the key. -use aimux_core::shared::{JsonObject, SharedProviderOptions, provider_namespace}; -use aimux_core::types::ProviderMetadata; -use serde_json::Value; +use aimux_core::shared::{JsonObject, SharedProviderOptions}; /// The providerOptions / providerMetadata key. pub(crate) const NAMESPACE: &str = "mistral"; @@ -18,7 +16,28 @@ pub(crate) fn mistral_options( provider_options?.get(NAMESPACE) } -/// `{ "mistral": payload }`. -pub(crate) fn mistral_metadata(payload: Value) -> ProviderMetadata { - provider_namespace(NAMESPACE, payload).expect("provider metadata must be an object") +/// Validate only the options declared by the upstream Mistral schema. +pub(crate) fn validated_mistral_options( + provider_options: Option<&SharedProviderOptions>, +) -> Result, aimux_core::AiMuxError> { + let Some(options) = mistral_options(provider_options) else { + return Ok(None); + }; + for (key, value) in options { + let valid = match key.as_str() { + "safePrompt" | "structuredOutputs" | "strictJsonSchema" | "parallelToolCalls" => { + value.is_boolean() + } + "documentImageLimit" | "documentPageLimit" => value.is_number(), + "promptCacheKey" => value.is_string(), + "reasoningEffort" => matches!(value.as_str(), Some("high" | "none")), + _ => true, + }; + if !valid { + return Err(aimux_core::AiMuxError::InvalidArgument(format!( + "Invalid mistral provider option {key}" + ))); + } + } + Ok(Some(options)) } diff --git a/aimux-providers/src/mistral/types.rs b/aimux-providers/src/mistral/types.rs index faa05b1e1..2963c8e2f 100644 --- a/aimux-providers/src/mistral/types.rs +++ b/aimux-providers/src/mistral/types.rs @@ -2,7 +2,7 @@ #![allow(dead_code)] -use serde::{Deserialize, Serialize}; +use serde::Deserialize; use serde_json::Value; // ── Non-streaming response ── @@ -50,27 +50,44 @@ pub struct FunctionCallResponse { pub arguments: String, } -#[derive(Debug, Deserialize, Default, Serialize)] +#[derive(Debug)] pub struct UsageResponse { - #[serde(default)] - pub prompt_tokens: Option, - #[serde(default)] - pub completion_tokens: Option, - #[serde(default)] - pub total_tokens: Option, - #[serde(default)] + pub fields: UsageFields, + pub raw: Value, +} + +impl<'de> Deserialize<'de> for UsageResponse { + fn deserialize>(deserializer: D) -> Result { + let raw = Value::deserialize(deserializer)?; + let fields = serde_json::from_value(raw.clone()).map_err(serde::de::Error::custom)?; + Ok(Self { fields, raw }) + } +} + +#[derive(Debug, Deserialize)] +pub struct UsageFields { + pub prompt_tokens: u32, + pub completion_tokens: u32, + pub total_tokens: u32, + pub prompt_audio_seconds: Option, + pub request_count: Option, + pub service_tier: Option, pub num_cached_tokens: Option, - #[serde(default)] pub prompt_tokens_details: Option, - // Mistral has used both spellings at different times. - #[serde(default)] pub prompt_token_details: Option, + pub completion_tokens_details: Option, } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Debug, Deserialize)] pub struct PromptTokensDetails { - #[serde(default)] pub cached_tokens: Option, + pub audio_tokens: Option, + pub messages: Option>, +} + +#[derive(Debug, Deserialize)] +pub struct CompletionTokensDetails { + pub reasoning_tokens: Option, } // ── Streaming response ── diff --git a/aimux-providers/src/open_responses.rs b/aimux-providers/src/open_responses.rs index f84a4027b..20e5c5ec4 100644 --- a/aimux-providers/src/open_responses.rs +++ b/aimux-providers/src/open_responses.rs @@ -275,8 +275,8 @@ impl LanguageModel for OpenResponsesModel { } async fn do_generate(&self, options: &CallOptions) -> Result { - validate_tool_output_media(&options.prompt)?; let exchange = self.config.exchange(options.headers.as_ref()).await?; + validate_tool_output_media(&options.prompt)?; let (body, warnings) = build_request_body(&self.model_id, options, &self.provider_options_name); let body = exchange.transform_body(body); @@ -415,8 +415,8 @@ impl LanguageModel for OpenResponsesModel { } async fn do_stream(&self, options: &CallOptions) -> Result { - validate_tool_output_media(&options.prompt)?; let exchange = self.config.exchange(options.headers.as_ref()).await?; + validate_tool_output_media(&options.prompt)?; let (body, warnings) = build_request_body(&self.model_id, options, &self.provider_options_name); diff --git a/aimux-providers/src/openai/config.rs b/aimux-providers/src/openai/config.rs index f9d29f127..f39d20f0f 100644 --- a/aimux-providers/src/openai/config.rs +++ b/aimux-providers/src/openai/config.rs @@ -166,3 +166,16 @@ impl OpenAIModelConfig { } } } + +/// Native language model URL patterns from the upstream model classes. +pub(crate) fn supported_urls(method: &str) -> SupportedUrls { + let mut urls = HashMap::new(); + if method.ends_with("chat") || method.ends_with("responses") { + let pattern = regex::Regex::new(r"^https?://.*$").expect("valid URL pattern"); + urls.insert("image/*".to_string(), vec![pattern.clone()]); + if method.ends_with("responses") { + urls.insert("application/pdf".to_string(), vec![pattern]); + } + } + SupportedUrls(urls) +} diff --git a/aimux-providers/src/openai/convert.rs b/aimux-providers/src/openai/convert.rs index ccecb76bd..f5165ab73 100644 --- a/aimux-providers/src/openai/convert.rs +++ b/aimux-providers/src/openai/convert.rs @@ -23,6 +23,121 @@ use super::convert_common::{ModelCapabilities, get_model_capabilities}; // `ModelCapabilities` / `SystemMessageMode` / `get_model_capabilities` live in // `super::convert_common` and are shared with the Responses converter (M10). +/// Apply the upstream structured-output compatibility rules only at schema positions. +pub(crate) fn normalize_json_schema( + schema: &Value, + warnings: &mut Vec, +) -> Result { + fn walk(value: &mut Value, removed: &mut [bool; 2]) -> Result<(), AiMuxError> { + let Some(obj) = value.as_object_mut() else { + return Ok(()); + }; + if let Some(names) = obj.get("propertyNames").filter(|v| !v.is_null()) { + if names.get("type").and_then(Value::as_str) != Some("string") { + return Err(AiMuxError::UnsupportedFunctionality( + "JSON Schema propertyNames that does not use a string schema".into(), + )); + } + obj.remove("propertyNames"); + removed[0] = true; + } + if obj + .get("pattern") + .and_then(Value::as_str) + .is_some_and(contains_lookaround) + { + obj.remove("pattern"); + removed[1] = true; + } + for key in [ + "properties", + "patternProperties", + "definitions", + "$defs", + "dependencies", + ] { + if let Some(Value::Object(record)) = obj.get_mut(key) { + for child in record.values_mut().filter(|v| !v.is_array()) { + walk(child, removed)?; + } + } + } + for key in [ + "additionalProperties", + "additionalItems", + "items", + "contains", + "not", + "allOf", + "anyOf", + "oneOf", + "if", + "then", + "else", + ] { + if let Some(child) = obj.get_mut(key) { + if let Value::Array(children) = child { + for child in children { + walk(child, removed)?; + } + } else { + walk(child, removed)?; + } + } + } + Ok(()) + } + let mut result = schema.clone(); + let mut removed = [false; 2]; + walk(&mut result, &mut removed)?; + for (index, feature, details) in [ + ( + 0, + "JSON Schema propertyNames", + "OpenAI does not support JSON Schema propertyNames. It was removed before sending the schema, so OpenAI will not enforce property-name constraints.", + ), + ( + 1, + "JSON Schema pattern with regex lookaround", + "OpenAI does not support regex lookaround in JSON Schema patterns. The pattern was removed before sending the schema, so OpenAI will not enforce that constraint.", + ), + ] { + if removed[index] { + warnings.push(Warning::Compatibility { + feature: feature.into(), + details: Some(details.into()), + }); + } + } + Ok(result) +} + +fn contains_lookaround(pattern: &str) -> bool { + let bytes = pattern.as_bytes(); + let (mut escaped, mut in_class) = (false, false); + for (i, &byte) in bytes.iter().enumerate() { + if escaped { + escaped = false; + continue; + } + match byte { + b'\\' => escaped = true, + b'[' => in_class = true, + b']' => in_class = false, + b'(' if !in_class + && bytes.get(i + 1) == Some(&b'?') + && (matches!(bytes.get(i + 2), Some(b'=' | b'!')) + || (bytes.get(i + 2) == Some(&b'<') + && matches!(bytes.get(i + 3), Some(b'=' | b'!')))) => + { + return true; + } + _ => {} + } + } + false +} + // ── Prepared tools ────────────────────────────────────────────────────────── /// A warning emitted while preparing tools (mirrors the V4 `SharedV4Warning` @@ -222,13 +337,13 @@ fn convert_file_part_to_openai(file: &FilePart, part_index: usize) -> Result (None, Some(url.as_str())), - FileData::Data { data } => { - let b64 = match data { + FileData::Data { data } => ( + Some(match data { FileBytes::Binary(bytes) => base64::engine::general_purpose::STANDARD.encode(bytes), FileBytes::Base64(data) => data.clone(), - }; - (Some(b64), None) - } + }), + None, + ), }; let data_b64 = data_b64.as_deref(); let filename = filename.as_deref(); @@ -348,22 +463,18 @@ fn convert_message_to_openai( json!({ "role": role, "content": content }) } LanguageModelMessage::Tool { content, .. } => { - let mut messages = Vec::new(); - for part in content { - let ToolPart::ToolResult(result) = part else { - continue; - }; - let mut content = tool_result_to_content(&result.output); - if let Some(breakpoint) = tool_result_cache_breakpoint(&result.output) - .or_else(|| get_prompt_cache_breakpoint(&result.provider_options)) - { - content = json!([{"type":"text", "text":content, "prompt_cache_breakpoint":breakpoint}]); - } - messages.push( - json!({"role":"tool", "content":content, "tool_call_id":result.tool_call_id}), - ); - } - return Ok(messages); + return Ok(content + .iter() + .filter_map(|part| { + let ToolPart::ToolResult(result) = part else { return None; }; + let mut content = tool_result_to_content(&result.output); + let breakpoint = tool_result_cache_breakpoint(&result.output).or_else(|| get_prompt_cache_breakpoint(&result.provider_options)); + if let Some(breakpoint) = breakpoint { + content = json!([{ "type": "text", "text": content, "prompt_cache_breakpoint": breakpoint }]); + } + Some(json!({ "role": "tool", "content": content, "tool_call_id": result.tool_call_id })) + }) + .collect()); } LanguageModelMessage::User { content, .. } => { let all_plain_text = content.len() == 1 @@ -409,14 +520,13 @@ fn convert_message_to_openai( text_parts.push(convert_text_part_to_openai(part)); } AssistantPart::ToolCall(part) => { - let arguments = if part.input.is_object() { - part.input.to_string() - } else { + let arguments = if !part.input.is_object() { "{}".to_string() + } else { + part.input.to_string() }; tool_calls.push(json!({ - "type": "function", - "id": part.tool_call_id, + "type": "function", "id": part.tool_call_id, "function": { "name": part.tool_name, "arguments": arguments }, })); } @@ -440,8 +550,6 @@ fn convert_message_to_openai( Ok(vec![message]) } -/// Serialize a tool-result `output` value into the OpenAI tool message -/// `content` string. pub(crate) fn tool_result_to_content(output: &ToolResultOutput) -> Value { Value::String(match output { ToolResultOutput::Text { value, .. } | ToolResultOutput::ErrorText { value, .. } => { @@ -545,7 +653,6 @@ pub(crate) fn tool_result_cache_breakpoint(output: &ToolResultOutput) -> Option< } } -/// Convert a text part to the OpenAI format. fn convert_text_part_to_openai(part: &TextPart) -> Value { let mut value = json!({ "type": "text", "text": part.text }); if let Some(bpt) = get_prompt_cache_breakpoint(&part.provider_options) { @@ -799,9 +906,13 @@ fn insert_sampling_params(body: &mut Value, params: &SamplingParams, options: &C /// `response_format` → body: a schema is sent as `json_schema` (strict), a /// bare JSON request as `json_object`. -fn apply_response_format(body: &mut Value, options: &CallOptions) { +fn apply_response_format( + body: &mut Value, + options: &CallOptions, + warnings: &mut Vec, +) -> Result<(), AiMuxError> { let Some(ref rf) = options.response_format else { - return; + return Ok(()); }; match rf { ResponseFormat::Text => {} @@ -812,16 +923,13 @@ fn apply_response_format(body: &mut Value, options: &CallOptions) { } => { if let Some(schema) = schema { let mut schema_obj = json!({}); - schema_obj["schema"] = schema.clone(); + schema_obj["schema"] = normalize_json_schema(schema, warnings)?; schema_obj["name"] = json!(name.clone().unwrap_or_else(|| "response".to_string())); if let Some(d) = description { schema_obj["description"] = json!(d); } - schema_obj["strict"] = json!( - openai_option(&options.provider_options, "strictJsonSchema") - .and_then(|value| value.as_bool()) - .unwrap_or(true) - ); + schema_obj["strict"] = openai_option(&options.provider_options, "strictJsonSchema") + .unwrap_or(json!(true)); body["response_format"] = json!({ "type": "json_schema", "json_schema": schema_obj, @@ -831,6 +939,7 @@ fn apply_response_format(body: &mut Value, options: &CallOptions) { } } } + Ok(()) } /// Pass through the simple provider-specific options that map 1:1 to a body @@ -856,7 +965,7 @@ fn apply_provider_option_passthrough( set("promptCacheOptions", "prompt_cache_options"); set("safetyIdentifier", "safety_identifier"); if let Some(logprobs) = openai_option(provider_opts, "logprobs") - && (logprobs == json!(true) || logprobs.is_number()) + && (logprobs == true || logprobs.is_number()) { body["logprobs"] = json!(true); body["top_logprobs"] = if logprobs.is_number() { @@ -891,7 +1000,7 @@ fn apply_service_tier( }); } } - "priority" => { + "priority" | "fast" => { if caps.supports_priority_processing { body["service_tier"] = json!(st); } else { @@ -911,7 +1020,11 @@ fn apply_service_tier( } /// Function tools → `tools` / `tool_choice`. -fn apply_tools(body: &mut Value, options: &CallOptions) { +fn apply_tools( + body: &mut Value, + options: &CallOptions, + warnings: &mut Vec, +) -> Result<(), AiMuxError> { let function_tools: Option> = options.tools.as_ref().map(|tools| { tools .iter() @@ -922,6 +1035,26 @@ fn apply_tools(body: &mut Value, options: &CallOptions) { .collect() }); + let mut function_tools = function_tools; + for tool in function_tools.iter_mut().flatten() { + tool.input_schema = normalize_json_schema(&tool.input_schema, warnings)?; + } + for tool in options.tools.iter().flatten() { + if matches!(tool, Tool::Provider(_)) { + warnings.push(Warning::Unsupported { + feature: "tool type: provider".into(), + details: None, + }); + } + } + if options + .tools + .as_ref() + .is_some_and(|tools| !tools.is_empty()) + && function_tools.as_ref().is_some_and(Vec::is_empty) + { + body["tools"] = json!([]); + } let prepared = prepare_tools(&function_tools, options.tool_choice.as_ref()); if let Some(tools) = prepared.tools { body["tools"] = json!(tools); @@ -929,6 +1062,7 @@ fn apply_tools(body: &mut Value, options: &CallOptions) { body["tool_choice"] = tc; } } + Ok(()) } /// Convert `CallOptions` to an OpenAI request body, returning warnings. @@ -972,7 +1106,20 @@ pub(crate) fn build_request_body_with_chat_options( }); let provider_opts = &parsed_provider_opts; - let resolved_reasoning_effort = resolve_reasoning_effort(provider_opts, &options.reasoning); + let mut resolved_reasoning_effort = resolve_reasoning_effort(provider_opts, &options.reasoning); + if let (Some(effort), Some(supported)) = + (&resolved_reasoning_effort, caps.supported_reasoning_efforts) + && !supported.contains(&effort.as_str()) + { + warnings.push(Warning::Unsupported { + feature: "reasoningEffort".into(), + details: Some(format!( + "{model_id} only supports the following reasoning efforts: {}", + supported.join(", ") + )), + }); + resolved_reasoning_effort = None; + } let is_reasoning_model = resolve_is_reasoning_model(provider_opts, &caps); let system_message_mode = resolve_system_message_mode(provider_opts, is_reasoning_model, &caps); @@ -984,6 +1131,17 @@ pub(crate) fn build_request_body_with_chat_options( }); } + if system_message_mode == SystemMessageMode::Remove { + for _ in options + .prompt + .iter() + .filter(|m| matches!(m, LanguageModelMessage::System { .. })) + { + warnings.push(Warning::Other { + message: "system messages are removed for this model".into(), + }); + } + } let messages = convert_prompt_to_openai_messages_with_mode_fallible(&options.prompt, system_message_mode)?; @@ -1009,16 +1167,41 @@ pub(crate) fn build_request_body_with_chat_options( ); insert_sampling_params(&mut body, &sampling, options); - apply_response_format(&mut body, options); + apply_response_format(&mut body, options, &mut warnings)?; apply_provider_option_passthrough(&mut body, provider_opts); - // Reasoning effort (v3 passthrough: no built-in vendor normalization). + if caps.supported_reasoning_efforts.is_some() && body.get("prompt_cache_retention").is_some() { + body.as_object_mut() + .unwrap() + .remove("prompt_cache_retention"); + warnings.push(Warning::Unsupported { feature: "promptCacheRetention".into(), details: Some("promptCacheRetention is not supported by sixth-generation and later models; use promptCacheOptions instead".into()) }); + } + if is_reasoning_model { + let allow_logprobs = resolved_reasoning_effort.as_deref() == Some("none") + && caps.supports_non_reasoning_parameters; + for (key, feature) in [ + ("logprobs", "logprobs"), + ("logit_bias", "logitBias"), + ("top_logprobs", "topLogprobs"), + ] { + if key == "logprobs" && allow_logprobs { + continue; + } + if body.as_object_mut().unwrap().remove(key).is_some() { + warnings.push(Warning::Other { + message: format!("{feature} is not supported for reasoning models"), + }); + } + } + } + + // Provider reasoning takes precedence over the standardized setting. if let Some(ref effort) = resolved_reasoning_effort { body["reasoning_effort"] = json!(effort); } apply_service_tier(&mut body, provider_opts, &caps, &mut warnings); - apply_tools(&mut body, options); + apply_tools(&mut body, options, &mut warnings)?; Ok(RequestBodyResult { body, warnings }) } diff --git a/aimux-providers/src/openai/convert_common.rs b/aimux-providers/src/openai/convert_common.rs index 88bb57869..12d7f648b 100644 --- a/aimux-providers/src/openai/convert_common.rs +++ b/aimux-providers/src/openai/convert_common.rs @@ -31,15 +31,10 @@ pub fn get_gpt_version(model_id: &str) -> Option { .map(|i| (&stripped[..i], &stripped[i..])) .unwrap_or((stripped, "")); if minor_str.is_empty() { - return Some(GptVersion { - major, - minor: None, - variant: if remainder.is_empty() { - None - } else { - Some(remainder.trim_start_matches('-').to_string()) - }, - }); + return None; + } + if !after.is_empty() && (!after.starts_with('-') || after.len() == 1) { + return None; } ( minor_str.parse::().ok(), @@ -50,6 +45,9 @@ pub fn get_gpt_version(model_id: &str) -> Option { }, ) } else { + if !remainder.is_empty() && (!remainder.starts_with('-') || remainder.len() == 1) { + return None; + } ( None, if remainder.is_empty() { @@ -70,10 +68,13 @@ pub fn get_gpt_version(model_id: &str) -> Option { /// Extract o-series version (e.g. `o3-mini` → 3). Mirrors TS `getOSeriesVersion`. pub fn get_o_series_version(model_id: &str) -> Option { let rest = model_id.strip_prefix('o')?; - let (digits, _) = rest + let (digits, remainder) = rest .find(|c: char| !c.is_ascii_digit()) .map(|i| (&rest[..i], &rest[i..])) .unwrap_or((rest, "")); + if !remainder.is_empty() && !remainder.starts_with('-') { + return None; + } if digits.is_empty() { return None; } @@ -87,6 +88,8 @@ pub struct ModelCapabilities { pub supports_flex_processing: bool, pub supports_priority_processing: bool, pub supports_non_reasoning_parameters: bool, + pub supports_configuration_update: bool, + pub supported_reasoning_efforts: Option<&'static [&'static str]>, } /// How system messages are mapped. @@ -125,7 +128,16 @@ pub fn get_model_capabilities(model_id: &str) -> ModelCapabilities { let supports_non_reasoning_parameters = gpt_version .as_ref() - .is_some_and(|v| v.major > 5 || (v.major == 5 && v.minor.unwrap_or(0) >= 1)); + .is_some_and(|v| v.major == 5 && v.minor.unwrap_or(0) >= 1); + + let supported_reasoning_efforts = gpt_version.as_ref().filter(|v| v.major >= 6).map(|v| { + if v.major == 6 && v.minor.is_none() && matches!(v.variant.as_deref(), Some("sol" | "luna")) + { + &["none", "low", "medium", "high", "xhigh", "max"][..] + } else { + &["low", "medium", "high", "xhigh", "max"][..] + } + }); let system_message_mode = if is_reasoning_model { SystemMessageMode::Developer @@ -139,5 +151,7 @@ pub fn get_model_capabilities(model_id: &str) -> ModelCapabilities { supports_flex_processing, supports_priority_processing, supports_non_reasoning_parameters, + supported_reasoning_efforts, + supports_configuration_update: gpt_version.as_ref().is_some_and(|v| v.major >= 6), } } diff --git a/aimux-providers/src/openai/embedding.rs b/aimux-providers/src/openai/embedding.rs index bc3c77053..8b4bef080 100644 --- a/aimux-providers/src/openai/embedding.rs +++ b/aimux-providers/src/openai/embedding.rs @@ -53,7 +53,13 @@ impl EmbeddingModel for OpenAIEmbeddingModel { &self, options: &EmbeddingCallOptions, ) -> Result { - let openai_options = parse_openai_provider_options(options.provider_options.as_ref()); + if options.values.len() > 2048 { + return Err(AiMuxError::InvalidArgument(format!( + "too many embedding values: {} exceeds 2048", + options.values.len() + ))); + } + let openai_options = parse_openai_provider_options(options.provider_options.as_ref())?; let mut body = Map::new(); body.insert("model".to_string(), json!(self.model_id)); @@ -87,42 +93,12 @@ impl EmbeddingModel for OpenAIEmbeddingModel { let raw_value: Value = resp.value; - // Extract embeddings: response.data[].embedding - // The embedding field can be a JSON array of floats (default) or a - // base64-encoded string (when encoding_format="base64"). - let embeddings: Vec> = raw_value - .get("data") - .and_then(|d| d.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|item| { - let emb = item.get("embedding")?; - if let Some(arr) = emb.as_array() { - // Standard format: array of numbers - Some( - arr.iter() - .filter_map(|v| v.as_f64().map(|f| f as f32)) - .collect(), - ) - } else if let Some(s) = emb.as_str() { - // Base64 format: decode to little-endian f32 array - decode_base64_embedding(s) - } else { - None - } - }) - .collect() - }) - .unwrap_or_default(); - - // Extract usage: response.usage.prompt_tokens - let usage = raw_value - .get("usage") - .and_then(|u| u.get("prompt_tokens")) - .and_then(serde_json::Value::as_u64) - .map(|tokens| EmbeddingUsage { - tokens: tokens as u32, - }); + let parsed: OpenAIEmbeddingResponse = serde_json::from_value(raw_value.clone()) + .map_err(|error| AiMuxError::InvalidResponseData(error.to_string()))?; + let embeddings = parsed.data.into_iter().map(|item| item.embedding).collect(); + let usage = parsed.usage.map(|usage| EmbeddingUsage { + tokens: usage.prompt_tokens as u32, + }); Ok(EmbeddingResult { embeddings, @@ -139,49 +115,46 @@ impl EmbeddingModel for OpenAIEmbeddingModel { // ── Provider options parsing ───────────────────────────────────────────────── -/// Parsed `openai` embedding provider options. +#[derive(serde::Deserialize)] +struct OpenAIEmbeddingResponse { + data: Vec, + usage: Option, +} + +#[derive(serde::Deserialize)] +struct OpenAIEmbeddingData { + embedding: Vec, +} + +#[derive(serde::Deserialize)] +struct OpenAIEmbeddingUsage { + prompt_tokens: f64, +} + +#[derive(Default, serde::Deserialize)] struct OpenAIEmbeddingProviderOptions { - dimensions: Option, + dimensions: Option, user: Option, } -/// Extract OpenAI-specific embedding options from the shared provider options. -/// -/// Mirrors the TS `parseProviderOptions({ provider: 'openai', ... })`. fn parse_openai_provider_options( options: Option<&SharedProviderOptions>, -) -> OpenAIEmbeddingProviderOptions { - let provider_opts = options.and_then(|opts| opts.get("openai")); - OpenAIEmbeddingProviderOptions { - dimensions: provider_opts - .and_then(|o| o.get("dimensions")) - .and_then(serde_json::Value::as_u64) - .map(|d| d as u32), - user: provider_opts - .and_then(|o| o.get("user")) - .and_then(|u| u.as_str()) - .map(std::string::ToString::to_string), - } -} - -/// Decode a base64-encoded embedding string into a `Vec`. -/// -/// OpenAI's API returns embeddings as base64-encoded little-endian f32 -/// arrays when `encoding_format: "base64"` is requested. The raw bytes -/// are decoded from base64, then reinterpreted as little-endian f32. -// Rust 1.98 clippy suggests `as_chunks::<4>()`, stabilized in 1.88; the -// workspace MSRV is 1.85. Drop when the MSRV moves past 1.88. -#[allow(unknown_lints, clippy::chunks_exact_to_as_chunks)] -fn decode_base64_embedding(s: &str) -> Option> { - use base64::Engine; - let bytes = base64::engine::general_purpose::STANDARD.decode(s).ok()?; - if bytes.len() % 4 != 0 { - return None; +) -> Result { + match options.and_then(|options| options.get("openai")) { + Some(value) => { + if value + .get("dimensions") + .is_some_and(|value| !value.is_number()) + || value.get("user").is_some_and(|value| !value.is_string()) + { + return Err(AiMuxError::InvalidArgument( + "invalid openai embedding options".into(), + )); + } + serde_json::from_value(Value::Object(value.clone())).map_err(|error| { + AiMuxError::InvalidArgument(format!("invalid openai embedding options: {error}")) + }) + } + None => Ok(OpenAIEmbeddingProviderOptions::default()), } - Some( - bytes - .chunks_exact(4) - .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) - .collect(), - ) } diff --git a/aimux-providers/src/openai/files.rs b/aimux-providers/src/openai/files.rs index 9c88bdb67..6df3c3c21 100644 --- a/aimux-providers/src/openai/files.rs +++ b/aimux-providers/src/openai/files.rs @@ -121,14 +121,17 @@ impl Files for OpenAIFiles { .unwrap_or_else(|| "blob".to_string()); // Build multipart/form-data body manually. + let expiry_seconds = openai_options.expires_after.map(|value| value.to_string()); + let mut text_fields = vec![("purpose", purpose.as_str())]; + if let Some(seconds) = expiry_seconds.as_deref() { + text_fields.push(("expires_after[anchor]", "created_at")); + text_fields.push(("expires_after[seconds]", seconds)); + } let (body, content_type) = build_multipart_form( &format!("form-data; name=\"file\"; filename=\"{filename}\""), &options.media_type, &file_bytes, - &[("purpose", purpose.as_str())], - openai_options - .expires_after - .map(|v| ("expires_after", v.to_string())), + &text_fields, ); let header_list = self.config.request_headers(None).await?; @@ -200,13 +203,11 @@ impl Files for OpenAIFiles { /// - `file_media_type`: the media type for the file part. /// - `file_bytes`: the raw file content. /// - `text_fields`: name/value pairs for simple text fields (e.g. `purpose`). -/// - `extra_text_field`: an optional additional text field. fn build_multipart_form( file_content_disposition: &str, file_media_type: &str, file_bytes: &[u8], text_fields: &[(&str, &str)], - extra_text_field: Option<(&str, String)>, ) -> (Vec, String) { let boundary = generate_boundary(); let mut body = Vec::new(); @@ -223,18 +224,6 @@ fn build_multipart_form( body.extend_from_slice(b"\r\n"); } - // Optional extra text field. - if let Some((name, value)) = extra_text_field { - body.extend_from_slice(b"--"); - body.extend_from_slice(boundary.as_bytes()); - body.extend_from_slice(b"\r\n"); - body.extend_from_slice( - format!("Content-Disposition: form-data; name=\"{name}\"\r\n\r\n").as_bytes(), - ); - body.extend_from_slice(value.as_bytes()); - body.extend_from_slice(b"\r\n"); - } - // File part. body.extend_from_slice(b"--"); body.extend_from_slice(boundary.as_bytes()); diff --git a/aimux-providers/src/openai/image.rs b/aimux-providers/src/openai/image.rs index ea93697ab..7319299b8 100644 --- a/aimux-providers/src/openai/image.rs +++ b/aimux-providers/src/openai/image.rs @@ -115,9 +115,10 @@ impl ImageModel for OpenAIImageModel { let (body_value, response_headers, _response_body) = if options.files.is_some() { // ── Edit path: multipart form data ── + validate_image_options(&options.provider_options, true)?; let openai_options = parse_edit_provider_options(&options.provider_options); let (form_body, content_type) = - build_edit_multipart(&self.model_id, options, &openai_options); + build_edit_multipart(&self.model_id, options, &openai_options).await?; let resp = aimux_provider_utils::post_to_api( self.config @@ -132,6 +133,7 @@ impl ImageModel for OpenAIImageModel { (val, resp.response_headers, None) } else { // ── Generation path: JSON body ── + validate_image_options(&options.provider_options, false)?; let openai_options = parse_generation_provider_options(&options.provider_options); let body = self .config @@ -174,6 +176,57 @@ impl ImageModel for OpenAIImageModel { // ── Helpers ───────────────────────────────────────────────────────────────── +/// Validate the fields accepted by the upstream generation/edit schemas. +fn validate_image_options(options: &SharedProviderOptions, edit: bool) -> Result<(), AiMuxError> { + let Some(options) = options.get("openai") else { + return Ok(()); + }; + let enums: &[(&str, &[&str])] = &[ + ( + "quality", + &[ + "standard", "hd", "low", "medium", "high", "xhigh", "max", "auto", + ], + ), + ("background", &["transparent", "opaque", "auto"]), + ("outputFormat", &["png", "jpeg", "webp"]), + ]; + let validate_enum = |name: &str, values: &[&str]| -> Result<(), AiMuxError> { + if let Some(value) = options.get(name) + && !value.as_str().is_some_and(|value| values.contains(&value)) + { + return Err(AiMuxError::InvalidArgument(format!( + "invalid openai image option: {name}" + ))); + } + Ok(()) + }; + for (name, values) in enums { + validate_enum(name, values)?; + } + if edit { + validate_enum("inputFidelity", &["high", "low"])?; + } else { + validate_enum("style", &["vivid", "natural"])?; + validate_enum("moderation", &["auto", "low"])?; + } + if let Some(value) = options.get("outputCompression") + && !value + .as_f64() + .is_some_and(|value| (0.0..=100.0).contains(&value) && value.fract() == 0.0) + { + return Err(AiMuxError::InvalidArgument( + "invalid openai image option: outputCompression".into(), + )); + } + if options.get("user").is_some_and(|value| !value.is_string()) { + return Err(AiMuxError::InvalidArgument( + "invalid openai image option: user".into(), + )); + } + Ok(()) +} + /// OpenAI generation provider options (camelCase → snake_case mapping). struct GenerationOptions { quality: Option, @@ -269,38 +322,55 @@ fn build_generation_body( body } -/// Decode an `ImageFile` to raw bytes. -fn image_file_to_bytes(file: &ImageFile) -> Result, AiMuxError> { - match file { - ImageFile::File { data, .. } => match data { - ImageFileData::Binary(bytes) => Ok(bytes.clone()), - ImageFileData::Base64(b64) => { - base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64) - .map_err(|e| AiMuxError::InvalidArgument(format!("invalid base64: {e}"))) - } - }, - ImageFile::Url { .. } => Err(AiMuxError::InvalidArgument( - "OpenAI image edit does not support URL files in this implementation".to_string(), - )), - } -} - -/// Get the media type from an `ImageFile`. -fn image_file_media_type(file: &ImageFile) -> &str { +/// Convert inline data or download a URL without forwarding provider credentials. +async fn image_file_to_bytes( + file: &ImageFile, + options: &ImageCallOptions, +) -> Result<(Vec, String), AiMuxError> { match file { - ImageFile::File { media_type, .. } => media_type, - ImageFile::Url { .. } => "application/octet-stream", + ImageFile::File { + data, media_type, .. + } => { + let bytes = match data { + ImageFileData::Binary(bytes) => bytes.clone(), + ImageFileData::Base64(value) => { + base64::Engine::decode(&base64::engine::general_purpose::STANDARD, value) + .map_err(|error| { + AiMuxError::InvalidArgument(format!("invalid base64: {error}")) + })? + } + }; + Ok((bytes, media_type.clone())) + } + ImageFile::Url { url, .. } => { + let response = aimux_provider_utils::get_from_api( + aimux_provider_utils::HttpRequest { + url: url.clone(), + abort_signal: options.abort_signal.clone(), + ..Default::default() + }, + aimux_provider_utils::create_binary_response_handler(), + super::openai_failed_response_handler(), + ) + .await?; + let media_type = response + .response_headers + .get("content-type") + .cloned() + .unwrap_or_default(); + Ok((response.value.to_vec(), media_type)) + } } } /// Build the multipart form-data body for `/images/edits`. /// /// Returns `(body_bytes, content_type)`. -fn build_edit_multipart( +async fn build_edit_multipart( model_id: &str, options: &ImageCallOptions, openai: &EditOptions, -) -> (Vec, String) { +) -> Result<(Vec, String), AiMuxError> { let boundary = generate_boundary(); let mut body = Vec::new(); @@ -346,22 +416,26 @@ fn build_edit_multipart( let files = options.files.as_ref().unwrap(); if files.len() == 1 { let file = &files[0]; - let media_type = image_file_media_type(file); - let bytes = image_file_to_bytes(file).unwrap_or_default(); - write_file_field(&mut body, &boundary, "image", "image", media_type, &bytes); + let (bytes, media_type) = image_file_to_bytes(file, options).await?; + write_file_field(&mut body, &boundary, "image", "image", &media_type, &bytes); } else { for file in files { - let media_type = image_file_media_type(file); - let bytes = image_file_to_bytes(file).unwrap_or_default(); - write_file_field(&mut body, &boundary, "image[]", "image", media_type, &bytes); + let (bytes, media_type) = image_file_to_bytes(file, options).await?; + write_file_field( + &mut body, + &boundary, + "image[]", + "image", + &media_type, + &bytes, + ); } } // Mask (optional). if let Some(ref mask) = options.mask { - let media_type = image_file_media_type(mask); - let bytes = image_file_to_bytes(mask).unwrap_or_default(); - write_file_field(&mut body, &boundary, "mask", "mask", media_type, &bytes); + let (bytes, media_type) = image_file_to_bytes(mask, options).await?; + write_file_field(&mut body, &boundary, "mask", "mask", &media_type, &bytes); } // Closing boundary. @@ -370,7 +444,7 @@ fn build_edit_multipart( body.extend_from_slice(b"--\r\n"); let content_type = format!("multipart/form-data; boundary={boundary}"); - (body, content_type) + Ok((body, content_type)) } fn write_text_field(body: &mut Vec, boundary: &str, name: &str, value: &str) { diff --git a/aimux-providers/src/openai/mod.rs b/aimux-providers/src/openai/mod.rs index cdc8c4ba7..d23104813 100644 --- a/aimux-providers/src/openai/mod.rs +++ b/aimux-providers/src/openai/mod.rs @@ -22,11 +22,11 @@ pub mod model; pub(crate) mod options; pub mod responses; pub mod speech; +pub mod tools; pub mod transcription; mod types; pub use completion::OpenAICompletionModel; -pub use config::TransformRequestBody; pub use decision::OpenAIDecisionModel; pub use embedding::OpenAIEmbeddingModel; pub use files::OpenAIFiles; @@ -34,6 +34,7 @@ pub use image::OpenAIImageModel; pub use model::OpenAIModel; pub use responses::OpenAIResponsesModel; pub use speech::OpenAISpeechModel; +pub use tools::{OpenAIComputerAction, OpenAIComputerSafetyCheck}; pub use transcription::OpenAITranscriptionModel; use std::sync::{Arc, OnceLock}; @@ -86,33 +87,89 @@ pub(crate) fn openai_stream_error( request_body_values: Value, response_headers: std::collections::HashMap, ) -> AiMuxError { - let message = error + let payload = if error.get("type").and_then(Value::as_str) == Some("response.failed") { + error + .get("response") + .and_then(|response| response.get("error")) + .unwrap_or(error) + } else { + error.get("error").unwrap_or(error) + }; + let message = payload .get("message") .and_then(Value::as_str) .unwrap_or("OpenAI stream failed before any output was generated") .to_owned(); - let code = error.get("code").or_else(|| error.get("type")); - let provider_code = code.and_then(|value| match value { + let provider_code = payload.get("code").and_then(|value| match value { Value::String(value) => Some(value.clone()), Value::Number(value) => Some(value.to_string()), _ => None, }); - // Only a numeric HTTP status in the payload is a status; a string code - // ("invalid_api_key") must not be laundered into a retryable 500. - let status_code = code - .and_then(Value::as_u64) - .filter(|status| (400..=599).contains(status)) - .map(|status| status as u16); - - aimux_provider_utils::stream_error_api_call( + let error_type = payload + .get("type") + .and_then(Value::as_str) + .unwrap_or_default(); + let discriminator = format!( + "{} {error_type}", + provider_code.as_deref().unwrap_or_default() + ) + .to_ascii_lowercase(); + let explicit_status = payload + .get("code") + .and_then(|value| { + value + .as_u64() + .and_then(|value| u16::try_from(value).ok()) + .or_else(|| { + value + .as_str() + .filter(|value| { + value.len() == 3 && value.bytes().all(|byte| byte.is_ascii_digit()) + }) + .and_then(|value| value.parse::().ok()) + }) + }) + .filter(|value| (400..=599).contains(value)); + let status = explicit_status.unwrap_or_else(|| { + if ["insufficient_quota", "rate_limit"] + .iter() + .any(|term| discriminator.contains(term)) + { + 429 + } else if discriminator.contains("authentication") { + 401 + } else if discriminator.contains("permission") { + 403 + } else if discriminator.contains("not_found") { + 404 + } else if ["invalid", "bad_request", "context_length"] + .iter() + .any(|term| discriminator.contains(term)) + { + 400 + } else if discriminator.contains("overload") { + 503 + } else if discriminator.contains("timeout") { + 504 + } else { + 500 + } + }); + let quota = provider_code.as_deref() == Some("insufficient_quota") + || error_type == "insufficient_quota"; + let mut result = aimux_provider_utils::stream_error_api_call( message, provider_code, - status_code, + Some(status), error, url, request_body_values, response_headers, - ) + ); + if quota && let AiMuxError::ApiCall(details) = &mut result { + details.is_retryable = false; + } + result } const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1"; @@ -120,13 +177,13 @@ const API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; /// Settings of [`create_openai`] (the AI SDK's `OpenAIProviderSettings`). /// -/// Every field is optional. Nothing here is evaluated when the provider is -/// created except `base_url` and `name`; `api_key` and `headers` are evaluated +/// Every field is optional. The base URL (including `OPENAI_BASE_URL`) and +/// name are resolved when the provider is created; `api_key` and `headers` are evaluated /// on every request. #[derive(Clone, Default)] pub struct OpenAIProviderSettings { - /// Base URL for the API calls. Default `https://api.openai.com/v1`; a - /// trailing slash is removed. + /// Base URL for API calls. Reads `OPENAI_BASE_URL` when absent, then + /// defaults to `https://api.openai.com/v1`. A trailing slash is removed. pub base_url: Option, /// The API key. `None` loads `OPENAI_API_KEY` when a request is made and /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An @@ -177,8 +234,9 @@ impl std::fmt::Debug for OpenAIProviderSettings { /// URL with a host. That is the only way this fails: the key is loaded per /// request, not here. pub fn create_openai(settings: OpenAIProviderSettings) -> Result { - let base_url = match load_optional_setting(settings.base_url.as_deref(), "OPENAI_BASE_URL") { - Some(url) => validate_base_url(&url)?, + let configured_url = load_optional_setting(settings.base_url.as_deref(), "OPENAI_BASE_URL"); + let base_url = match configured_url.as_deref() { + Some(url) => validate_base_url(url)?, None => DEFAULT_BASE_URL.to_string(), }; let name = settings.name.unwrap_or_else(|| "openai".to_string()); @@ -210,12 +268,13 @@ pub fn create_openai(settings: OpenAIProviderSettings) -> Result &'static OpenAIProvider { static DEFAULT: OnceLock = OnceLock::new(); DEFAULT.get_or_init(|| { - create_openai(OpenAIProviderSettings::default()).expect("invalid OPENAI_BASE_URL") + create_openai(OpenAIProviderSettings::default()) + .expect("default OpenAI base URL must be valid") }) } @@ -232,13 +291,18 @@ pub struct OpenAIProvider { impl OpenAIProvider { fn model_config(&self, method: &str) -> OpenAIModelConfig { - OpenAIModelConfig::fixed( + let mut config = OpenAIModelConfig::fixed( format!("{}.{method}", self.name), self.base_url.clone(), self.headers.clone(), self.fetch.clone(), None, - ) + ); + config.supported_urls = config::supported_urls(method); + if method == "responses" { + config.responses.file_id_prefixes = vec!["file-"]; + } + config } /// A chat-completions model; `provider()` is `"{name}.chat"`. diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs index bb392243c..064200ba3 100644 --- a/aimux-providers/src/openai/model.rs +++ b/aimux-providers/src/openai/model.rs @@ -20,7 +20,9 @@ use aimux_core::shared::provider_namespace; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::{HttpRequest, StreamingToolCallDelta, StreamingToolCallTracker}; +use aimux_provider_utils::{ + HttpRequest, StreamingToolCallDelta, StreamingToolCallTracker, TypeValidation, generate_id, +}; use super::config::{OpenAIModelConfig, TransformRequestBody}; use super::convert::{ @@ -211,15 +213,17 @@ pub(crate) async fn execute_generate( let response_value = resp.raw_value.unwrap_or(Value::Null); let data = resp.value; - let choice = data - .choices - .into_iter() - .next() - .ok_or_else(|| AiMuxError::InvalidResponseData("no choices in response".to_string()))?; + let choice = data.choices.into_iter().next().ok_or_else(|| { + AiMuxError::InvalidResponseData("Response did not contain any choices.".to_string()) + })?; // Build content array. let mut content = Vec::new(); - if let Some(text) = choice.message.content + if let Some(text) = choice + .message + .content + .filter(|text| !text.is_empty()) + .or_else(|| choice.message.audio.and_then(|audio| audio.transcript)) && !text.is_empty() { content.push(GenerateContent::Text { @@ -286,7 +290,10 @@ pub(crate) async fn execute_generate( raw: None, }); - let usage = convert_usage(&data.usage, response_value.get("usage")); + let usage = data.usage.as_ref().map_or_else(Usage::default, |usage| { + convert_usage(usage, response_value.get("usage")) + }); + let raw_usage = data.usage.unwrap_or_default(); // Build provider metadata: logprobs + prediction tokens. let mut pm_openai = serde_json::json!({}); @@ -296,7 +303,7 @@ pub(crate) async fn execute_generate( pm_openai["logprobs"] = content_lp.clone(); } // Accepted/rejected prediction tokens from completion_tokens_details. - if let Some(ref details) = data.usage.completion_tokens_details { + if let Some(ref details) = raw_usage.completion_tokens_details { if let Some(apt) = details.accepted_prediction_tokens { pm_openai["acceptedPredictionTokens"] = json!(apt); } @@ -314,12 +321,13 @@ pub(crate) async fn execute_generate( warnings: request_result.warnings, provider_metadata, response: Some(aimux_core::shared::ResponseInfo { - id: Some(data.id), + id: data.id, timestamp: data .created + .filter(|created| *created != 0) .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) .map(|dt| dt.to_rfc3339()), - model_id: Some(data.model), + model_id: data.model, headers: Some(response_headers), body: Some(response_value), }), @@ -410,7 +418,10 @@ pub(crate) async fn execute_stream( // Streamed tool calls, correlated by wire id, index and function name // (the AI SDK's StreamingToolCallTracker) and finalized on flush. - let mut tool_calls = StreamingToolCallTracker::new(); + // Same options as `@ai-sdk/openai`'s chat model. + let mut tool_calls = StreamingToolCallTracker::new() + .with_generate_id(generate_id) + .with_type_validation(TypeValidation::IfPresent); let mut tool_parts = Vec::new(); // Process the first event (already peeked) then the rest. @@ -420,10 +431,6 @@ pub(crate) async fn execute_stream( let mut stream_errored = false; while let Some(event) = event_iter.next().await { - if stream_errored { - break; - } - match event { Ok(parsed) => { @@ -447,7 +454,7 @@ pub(crate) async fn execute_stream( ), }); stream_errored = true; - break; + continue; } // The chunk's top-level `usage`, taken from the raw JSON before @@ -462,20 +469,22 @@ pub(crate) async fn execute_stream( let chunk: StreamChunk = match serde_json::from_value(parsed) { Ok(c) => c, Err(e) => { - yield Err(e.into()); + stream_errored = true; + yield Ok(StreamPart::Error { error: e.into() }); continue; } }; // Emit ResponseMetadata from the first valid chunk. if !response_metadata_emitted - && (chunk.id.is_some() || chunk.model.is_some()) + && (chunk.id.as_ref().is_some_and(|id| !id.is_empty()) || chunk.model.as_ref().is_some_and(|model| !model.is_empty()) || chunk.created.is_some_and(|created| created != 0)) { response_metadata_emitted = true; yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { id: chunk.id.clone(), timestamp: chunk .created + .filter(|created| *created != 0) .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) .map(|dt| dt.to_rfc3339()), model_id: chunk.model.clone(), @@ -556,15 +565,17 @@ pub(crate) async fn execute_stream( for dtc in &tool_call_deltas { let function = dtc.function.as_ref(); let delta = StreamingToolCallDelta { - index: dtc.index, + index: Some(dtc.index), id: dtc.id.as_deref(), + r#type: dtc.r#type.as_deref(), name: function.and_then(|f| f.name.as_deref()), arguments: function.and_then(|f| f.arguments.as_deref()), - ..Default::default() + provider_metadata: None, }; // A malformed delta (new call without a - // function name) is invalid response data, - // as in the AI SDK; the stream ends. + // function name, or a non-`function` type) + // is invalid response data, as in the AI SDK; + // the stream ends. if let Err(error) = tool_calls.process(delta, &mut tool_parts) { yield Err(error.into()); return; @@ -610,22 +621,15 @@ pub(crate) async fn execute_stream( reasoning_started = false; } - // Close any open text segment. - if text_started { - yield Ok(StreamPart::TextEnd { - id: format!("{text_id}"), - provider_metadata: None, - }); - text_started = false; - } - + stream_errored = false; final_finish_reason = Some(parse_finish_reason(&reason)); } } } Err(error) => { let recoverable = error.is_recoverable_stream_error(); - yield Err(error); + stream_errored = true; + yield Ok(StreamPart::Error { error }); if !recoverable { return; } @@ -682,15 +686,11 @@ pub(crate) async fn execute_stream( } } else { final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Stop, + unified: FinishReasonUnified::Other, raw: None, }) }, - usage: if stream_errored { - Usage::default() - } else { - final_usage - }, + usage: final_usage, provider_metadata: Some(provider_metadata), }); }; diff --git a/aimux-providers/src/openai/responses/convert.rs b/aimux-providers/src/openai/responses/convert.rs index 882911c42..a94150da7 100644 --- a/aimux-providers/src/openai/responses/convert.rs +++ b/aimux-providers/src/openai/responses/convert.rs @@ -10,9 +10,9 @@ //! - `convert-openai-responses-usage.ts` -> [`convert_responses_usage`] //! - `map-openai-responses-finish-reason.ts` -> [`map_responses_finish_reason`] +use aimux_core::AiMuxError; use serde_json::{Value, json}; -use aimux_core::error::AiMuxError; use aimux_core::language_model_message::{ AssistantPart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, @@ -21,8 +21,8 @@ use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; use aimux_core::shared::{FileBytes, FileData, JsonObject, SharedProviderOptions}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified, ReasoningEffort, Usage, Warning}; -use aimux_provider_utils::{get_top_level_media_type, resolve_full_media_type}; +use super::tool_args::{validate_tool_args, validate_tool_call, validate_tool_result}; use super::types::ResponsesUsage; // -- Model capabilities ------------------------------------------------------ @@ -95,6 +95,7 @@ pub(crate) struct ResponsesProfile { pub(crate) namespace: ResponsesNamespace, /// Empty: no file part is treated as a file id. pub(crate) file_id_prefixes: Vec<&'static str>, + pub(crate) explicit_message_item_type: bool, } /// Get a value from the host namespace's options (`openai.` for OpenAI). @@ -132,47 +133,86 @@ pub struct ResponsesInputResult { /// `has_previous_response_id` is true, assistant reasoning/function-call items /// that already carry an `itemId` are skipped (they live in the previous /// response chain). -/// /// # Errors -/// Returns an error for text file parts or an unresolved inline image media type. +/// Rejects unsupported file data and unresolved file references. pub fn convert_to_responses_input( ns: ResponsesNamespace, prompt: &LanguageModelPrompt, system_message_mode: SystemMessageMode, store: bool, has_previous_response_id: bool, +) -> Result { + convert_responses_input_with_conversation( + ns, + prompt, + system_message_mode, + store, + has_previous_response_id, + false, + None, + ) +} + +fn convert_responses_input_with_conversation( + ns: ResponsesNamespace, + prompt: &LanguageModelPrompt, + system_message_mode: SystemMessageMode, + store: bool, + has_previous_response_id: bool, has_conversation: bool, + tools: Option<&[Tool]>, ) -> Result { let mut input: Vec = Vec::new(); let mut warnings: Vec = Vec::new(); let mut processed_approval_ids = std::collections::HashSet::new(); + let mut programmatic_tool_call_ids = std::collections::HashSet::new(); for msg in prompt { match msg { - LanguageModelMessage::System { content, .. } => match system_message_mode { - SystemMessageMode::System => { - input.push(json!({ - "role": "system", - "content": content, - })); - } - SystemMessageMode::Developer => { - input.push(json!({ - "role": "developer", - "content": content, - })); + LanguageModelMessage::System { + content, + provider_options, + } => { + if let Some(effort) = + openai_sub_option(ns, provider_options, "reasoningEffortUpdate") + { + input.push(json!({ "type": "configuration_update", "reasoning": { "effort": effort } })); + continue; } - SystemMessageMode::Remove => { - warnings.push(Warning::Other { - message: "system messages are removed for this model".to_string(), - }); + match system_message_mode { + SystemMessageMode::System => { + input.push(json!({ + "role": "system", + "content": system_content(ns, content, provider_options), + })); + } + SystemMessageMode::Developer => { + input.push(json!({ + "role": "developer", + "content": system_content(ns, content, provider_options), + })); + } + SystemMessageMode::Remove => { + warnings.push(Warning::Other { + message: "system messages are removed for this model".to_string(), + }); + } } - }, + } LanguageModelMessage::User { content, .. } => { let content: Vec = content .iter() - .map(|part| convert_user_part(ns, part)) - .collect::>()?; + .enumerate() + .map(|(index, part)| { + let mut value = convert_user_part(ns, part, index)?; + if let Some(breakpoint) = + openai_sub_option(ns, user_part_options(part), "promptCacheBreakpoint") + { + value["prompt_cache_breakpoint"] = breakpoint; + } + Ok(value) + }) + .collect::>()?; input.push(json!({ "role": "user", "content": content })); } LanguageModelMessage::Assistant { content, .. } => { @@ -183,7 +223,7 @@ pub fn convert_to_responses_input( provider_options, }) => { let id = item_id(ns, provider_options); - if (has_previous_response_id || has_conversation) && id.is_some() { + if has_conversation && id.is_some() { continue; } if store && let Some(ref id) = id { @@ -193,11 +233,8 @@ pub fn convert_to_responses_input( let phase = phase_from_provider_options(ns, provider_options); let mut item = json!({ "role": "assistant", - "content": [{ "type": "output_text", "text": text }], + "content": text, }); - if let Some(ref id) = id { - item["id"] = json!(id); - } if let Some(phase) = phase { item["phase"] = json!(phase); } @@ -208,12 +245,58 @@ pub fn convert_to_responses_input( tool_name, input: tool_input, provider_options, + provider_executed, .. }) => { let id = item_id(ns, provider_options); - if (has_previous_response_id || has_conversation) && id.is_some() { + if openai_sub_option(ns, provider_options, "caller") + .is_some_and(|caller| caller["type"] == "program") + { + programmatic_tool_call_ids.insert(tool_call_id.as_str()); + } + if has_conversation && id.is_some() { continue; } + let kind = provider_tool_kind(tools, tool_name); + if provider_executed == &Some(true) { + if store && let Some(id) = &id { + input.push(json!({"type":"item_reference", "id":id})); + } + if store || kind != Some("shell") { + continue; + } + } + if let Some(kind) = kind { + if store + && id.is_some() + && kind != "tool_search" + && kind != "programmatic_tool_calling" + && has_previous_response_id + { + continue; + } + if store && let Some(id) = &id { + input.push(json!({"type":"item_reference", "id":id})); + continue; + } + if let Some(mut item) = provider_tool_call_input( + kind, + tool_name, + tool_call_id, + id.as_deref(), + tool_input, + )? { + if kind == "custom" + && let Some(value) = + openai_sub_option(ns, provider_options, "async") + .filter(|value| !value.is_null()) + { + item["async"] = value; + } + input.push(item); + continue; + } + } let namespace = namespace_from_provider_options(ns, provider_options); let mut item = json!({ "type": "function_call", @@ -224,6 +307,22 @@ pub fn convert_to_responses_input( if let Some(ref namespace) = namespace { item["namespace"] = json!(namespace); } + if let Some(value) = openai_sub_option(ns, provider_options, "async") { + item["async"] = value; + } + if let Some(caller) = openai_sub_option(ns, provider_options, "caller") + { + let mut caller = map_tool_fields( + &caller, + &[("type", "type"), ("callerId", "caller_id")], + ); + if caller["type"] != "program" + && let Some(object) = caller.as_object_mut() + { + object.remove("caller_id"); + } + item["caller"] = caller; + } input.push(item); } AssistantPart::Reasoning(ReasoningPart { @@ -292,14 +391,24 @@ pub fn convert_to_responses_input( { continue; } - if store { - let id = item_id(ns, &part.provider_options) - .unwrap_or_else(|| part.tool_call_id.clone()); - input.push(json!({"type":"item_reference", "id":id})); - } else { - warnings.push(Warning::Other { - message:format!("Results for OpenAI tool {} are not sent to the API when store is false", part.tool_name), - }); + let id = item_id(ns, &part.provider_options) + .unwrap_or_else(|| part.tool_call_id.clone()); + let kind = provider_tool_kind(tools, &part.tool_name); + match kind { + Some("shell") => { + if let ToolResultOutput::Json { value, .. } = &part.output + && let Some(item) = provider_tool_result_input("shell", &part.tool_call_id, value)? + { input.push(item); } + } + _ if store => input.push(json!({"type":"item_reference", "id":id})), + Some("tool_search" | "programmatic_tool_calling") => { + if let ToolResultOutput::Json { value, .. } = &part.output { + let kind = kind.unwrap_or_default(); + validate_tool_result(kind, value)?; + input.push(if kind == "tool_search" { json!({"type":"tool_search_output", "id":id, "execution":"server", "call_id":null, "status":"completed", "tools":value["tools"]}) } else { json!({"type":"program_output", "id":id, "call_id":part.tool_call_id, "result":value["result"], "status":value["status"]}) }); + } + } + _ => warnings.push(Warning::Other { message:format!("Results for OpenAI tool {} are not sent to the API when store is false", part.tool_name) }), } } AssistantPart::Custom(part) if part.kind == "openai.compaction" => { @@ -333,9 +442,9 @@ pub fn convert_to_responses_input( } let ToolPart::ToolResult(ToolResultPart { tool_call_id, + tool_name, output, provider_options, - .. }) = part else { continue; @@ -344,39 +453,56 @@ pub fn convert_to_responses_input( provider_options, .. } = output && openai_sub_option(ns, provider_options, "approvalId") - .and_then(|value| value.as_str().map(|id| !id.is_empty())) - == Some(true) + .and_then(|id| id.as_str().map(|id| !id.is_empty())) + .unwrap_or(false) + { + continue; + } + let kind = provider_tool_kind(tools, tool_name); + if let Some(kind) = kind + && kind != "custom" + && let ToolResultOutput::Json { value, .. } = output + && let Some(item) = provider_tool_result_input(kind, tool_call_id, value)? { + input.push(item); continue; } - let mut content_value = convert_tool_result_output(ns, output, &mut warnings)?; + let custom = kind == Some("custom"); + if !custom + && matches!(output, ToolResultOutput::ExecutionDenied { .. }) + && (programmatic_tool_call_ids.contains(tool_call_id.as_str()) + || openai_sub_option(ns, provider_options, "caller") + .is_some_and(|caller| caller["type"] == "program")) + { + return Err(AiMuxError::UnsupportedFunctionality( + "execution-denied results for programmatic tool calls".into(), + )); + } + let mut content_value = + convert_tool_result_output(ns, output, custom, &mut warnings)?; + let has_output_schema = tools.is_some_and(|tools| tools.iter().any(|tool| { + matches!(tool, Tool::Function(tool) if tool.name == *tool_name && tool.provider_options.as_ref().and_then(|options| options.get("openai")).is_some_and(|options| options.contains_key("outputSchema"))) + })); + if has_output_schema + && matches!( + output, + ToolResultOutput::Text { .. } + | ToolResultOutput::ErrorText { .. } + | ToolResultOutput::ExecutionDenied { .. } + ) + { + content_value = json!(content_value.to_string()); + } if !matches!(output, ToolResultOutput::Content { .. }) - && let Some(breakpoint) = match output { - ToolResultOutput::Text { - provider_options, .. - } - | ToolResultOutput::Json { - provider_options, .. - } - | ToolResultOutput::ErrorText { - provider_options, .. - } - | ToolResultOutput::ErrorJson { - provider_options, .. - } - | ToolResultOutput::ExecutionDenied { - provider_options, .. - } => openai_sub_option(ns, provider_options, "promptCacheBreakpoint"), - ToolResultOutput::Content { .. } => None, - } - .or_else(|| { - openai_sub_option(ns, provider_options, "promptCacheBreakpoint") - }) + && let Some(breakpoint) = scalar_tool_result_cache_breakpoint(ns, output) + .or_else(|| { + openai_sub_option(ns, provider_options, "promptCacheBreakpoint") + }) { content_value = json!([{"type":"input_text", "text":content_value, "prompt_cache_breakpoint":breakpoint}]); } input.push(json!({ - "type": "function_call_output", + "type": if custom { "custom_tool_call_output" } else { "function_call_output" }, "call_id": tool_call_id, "output": content_value, })); @@ -406,78 +532,260 @@ pub fn convert_to_responses_input( Ok(ResponsesInputResult { input, warnings }) } +fn provider_tool_kind<'a>(tools: Option<&'a [Tool]>, name: &str) -> Option<&'a str> { + tools?.iter().find_map(|tool| match tool { + Tool::Provider(tool) if tool.name == name => tool.id.strip_prefix("openai."), + _ => None, + }) +} + +fn provider_tool_call_input( + kind: &str, + name: &str, + call_id: &str, + id: Option<&str>, + args: &Value, +) -> Result, AiMuxError> { + let parsed; + let args = if kind != "custom" && args.is_string() { + parsed = serde_json::from_str(args.as_str().unwrap_or_default()) + .map_err(|error| AiMuxError::InvalidArgument(error.to_string()))?; + &parsed + } else { + args + }; + if kind != "custom" { + validate_tool_call(kind, args)?; + } + let mut value = match kind { + "custom" => { + json!({"type":"custom_tool_call", "call_id":call_id, "name":name, "input":args.as_str().map(str::to_owned).unwrap_or_else(|| args.to_string())}) + } + "local_shell" => { + json!({"type":"local_shell_call", "call_id":call_id, "action": map_tool_fields(&args["action"], &[("command", "command"), ("timeoutMs", "timeout_ms"), ("user", "user"), ("workingDirectory", "working_directory"), ("env", "env")])}) + } + "shell" => { + json!({"type":"shell_call", "call_id":call_id, "status":"completed", "action":map_tool_fields(&args["action"], &[("commands", "commands"), ("timeoutMs", "timeout_ms"), ("maxOutputLength", "max_output_length")])}) + } + "apply_patch" => { + json!({"type":"apply_patch_call", "call_id":args["callId"], "status":"completed", "operation":args["operation"]}) + } + "computer" => { + let mut actions = args + .get("actions") + .and_then(Value::as_array) + .cloned() + .ok_or_else(|| AiMuxError::InvalidArgument("computer actions".into()))?; + for action in &mut actions { + if action["type"] == "scroll" + && let Some(object) = action.as_object_mut() + { + for (key, wire) in [("scrollX", "scroll_x"), ("scrollY", "scroll_y")] { + if let Some(v) = object.remove(key) { + object.insert(wire.into(), v); + } + } + } + } + json!({"type":"computer_call", "call_id":call_id, "status":args["status"], "actions":actions, "pending_safety_checks":args["pendingSafetyChecks"]}) + } + "programmatic_tool_calling" => { + json!({"type":"program", "call_id":call_id, "code":args["code"], "fingerprint":args["fingerprint"]}) + } + "tool_search" => { + let args = if let Some(text) = args.as_str() { + serde_json::from_str(text) + .map_err(|error| AiMuxError::InvalidArgument(error.to_string()))? + } else { + args.clone() + }; + let mut value = map_tool_fields(&args, &[("arguments", "arguments")]); + value["type"] = json!("tool_search_call"); + value["status"] = json!("completed"); + value["execution"] = json!(if args.get("call_id").is_some_and(|id| !id.is_null()) { + "client" + } else { + "server" + }); + value["call_id"] = args.get("call_id").cloned().unwrap_or(Value::Null); + value + } + _ => return Ok(None), + }; + if kind == "local_shell" { + value["action"]["type"] = json!("exec"); + } + if let Some(id) = id { + value["id"] = json!(id); + } else if ["tool_search", "programmatic_tool_calling"].contains(&kind) { + value["id"] = json!(call_id); + } + Ok(Some(value)) +} + +fn provider_tool_result_input( + kind: &str, + call_id: &str, + result: &Value, +) -> Result, AiMuxError> { + if !matches!( + kind, + "local_shell" | "shell" | "apply_patch" | "computer" | "tool_search" | "custom" + ) { + return Ok(None); + } + validate_tool_result(kind, result)?; + let mut value = match kind { + "local_shell" => json!({"type":"local_shell_call_output", "output":result["output"]}), + "shell" => { + let mut output = result + .get("output") + .and_then(Value::as_array) + .cloned() + .ok_or_else(|| AiMuxError::InvalidArgument("shell output".into()))?; + for item in &mut output { + if let Some(outcome) = item.get_mut("outcome").and_then(Value::as_object_mut) + && let Some(code) = outcome.remove("exitCode") + { + outcome.insert("exit_code".into(), code); + } + } + json!({"type":"shell_call_output", "output":output}) + } + "apply_patch" => { + json!({"type":"apply_patch_call_output", "status":result["status"], "output":result["output"]}) + } + "computer" => { + json!({"type":"computer_call_output", "output":map_tool_fields(&result["output"], &[("imageUrl", "image_url"), ("fileId", "file_id"), ("detail", "detail")])}) + } + "tool_search" => { + json!({"type":"tool_search_output", "execution":"client", "status":"completed", "tools":result["tools"]}) + } + "custom" => { + json!({"type":"custom_tool_call_output", "output":result.as_str().map(str::to_owned).unwrap_or_else(|| result.to_string())}) + } + _ => return Ok(None), + }; + if kind == "computer" { + value["output"]["type"] = json!("computer_screenshot"); + if let Some(checks) = result.get("acknowledgedSafetyChecks") { + value["acknowledged_safety_checks"] = checks.clone(); + } + } + value["call_id"] = json!(call_id); + Ok(Some(value)) +} + +fn user_part_options(part: &UserPart) -> &Option { + match part { + UserPart::Text(part) => &part.provider_options, + UserPart::File(part) => &part.provider_options, + } +} + +fn system_content( + ns: ResponsesNamespace, + text: &str, + provider_options: &Option, +) -> Value { + match openai_sub_option(ns, provider_options, "promptCacheBreakpoint") { + Some(value) => { + json!([{ "type": "input_text", "text": text, "prompt_cache_breakpoint": value }]) + } + None => json!(text), + } +} + /// Convert a single user-message content part into the Responses input shape. -fn convert_user_part(ns: ResponsesNamespace, part: &UserPart) -> Result { +fn convert_user_part( + ns: ResponsesNamespace, + part: &UserPart, + index: usize, +) -> Result { Ok(match part { UserPart::Text(part) => json!({ "type": "input_text", "text": part.text }), - UserPart::File(file) => { - let mut item = match &file.data { - FileData::Data { data } => { - let b64 = match data { - FileBytes::Binary(bytes) => { - use base64::Engine; - base64::engine::general_purpose::STANDARD.encode(bytes) - } - FileBytes::Base64(data) => data.clone(), - }; - if get_top_level_media_type(&file.media_type) == "image" { - let mut img = json!({ - "type": "input_image", - "image_url": format!("data:{};base64,{}", resolve_full_media_type(file)?, b64), - }); - if let Some(detail) = - openai_sub_option(ns, &file.provider_options, "imageDetail") - { - img["detail"] = detail; - } - img - } else { - let fname = file - .filename - .clone() - .unwrap_or_else(|| "part.pdf".to_string()); - json!({ - "type": "input_file", - "filename": fname, - "file_data": format!("data:{};base64,{}", file.media_type, b64), - }) + UserPart::File(file) => match &file.data { + FileData::Data { data } => { + use base64::Engine; + let b64 = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) } + FileBytes::Base64(data) => data.clone(), + }; + inline_file( + ns, + &aimux_provider_utils::resolve_full_media_type(file)?, + file.filename.as_deref(), + &b64, + &file.provider_options, + index, + ) + } + FileData::Reference { reference } => { + let file_id = reference.get(ns.write_key()).ok_or_else(|| { + AiMuxError::InvalidArgument(format!("No file reference for {}", ns.write_key())) + })?; + let mut part = json!({ "type": if file.media_type.split('/').next() == Some("image") { "input_image" } else { "input_file" }, "file_id": file_id }); + if part["type"] == "input_image" + && let Some(detail) = + openai_sub_option(ns, &file.provider_options, "imageDetail") + { + part["detail"] = detail; } - FileData::Url { url, .. } => { - if get_top_level_media_type(&file.media_type) == "image" { - json!({ "type": "input_image", "image_url": url }) - } else { - json!({ "type": "input_file", "file_url": url }) + part + } + FileData::Url { url, .. } => { + if file.media_type.split('/').next() == Some("image") { + let mut part = json!({ "type": "input_image", "image_url": url }); + if let Some(detail) = + openai_sub_option(ns, &file.provider_options, "imageDetail") + { + part["detail"] = detail; } + part + } else { + json!({ "type": "input_file", "file_url": url }) } - FileData::Text { .. } => { - return Err(AiMuxError::UnsupportedFunctionality( - "text file parts".into(), - )); - } - FileData::Reference { reference } => { - let file_id = reference.get(ns.write_key()).ok_or_else(|| { - AiMuxError::InvalidArgument(format!( - "missing file reference for provider {}", - ns.write_key() - )) - })?; - json!({ "type": if get_top_level_media_type(&file.media_type) == "image" { "input_image" } else { "input_file" }, "file_id": file_id }) - } - }; - if item["type"] == "input_image" - && let Some(detail) = openai_sub_option(ns, &file.provider_options, "imageDetail") - { - item["detail"] = detail; } - item - } + FileData::Text { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "text file parts".into(), + )); + } + }, }) } +fn scalar_tool_result_cache_breakpoint( + ns: ResponsesNamespace, + output: &ToolResultOutput, +) -> Option { + let provider_options = match output { + ToolResultOutput::Text { + provider_options, .. + } + | ToolResultOutput::Json { + provider_options, .. + } + | ToolResultOutput::ErrorText { + provider_options, .. + } + | ToolResultOutput::ErrorJson { + provider_options, .. + } + | ToolResultOutput::ExecutionDenied { + provider_options, .. + } => provider_options, + ToolResultOutput::Content { .. } => return None, + }; + openai_sub_option(ns, provider_options, "promptCacheBreakpoint") +} + fn convert_tool_result_output( ns: ResponsesNamespace, output: &ToolResultOutput, + custom: bool, warnings: &mut Vec, ) -> Result { let ToolResultOutput::Content { value } = output else { @@ -491,14 +799,23 @@ fn convert_tool_result_output( &part.provider_options, ), ToolResultContent::File(part) => { - if matches!(part.data, FileData::Text { .. }) { + if matches!(part.data, FileData::Text { .. }) + || (custom && matches!(part.data, FileData::Reference { .. })) + { warnings.push(Warning::Other { - message: "unsupported tool content part type: file with data type: text" - .into(), + message: format!( + "unsupported {}tool content part type: file with data type: {}", + if custom { "custom " } else { "" }, + if matches!(part.data, FileData::Reference { .. }) { + "reference" + } else { + "text" + } + ), }); continue; } - let mut converted = convert_user_part(ns, &UserPart::File(part.clone()))?; + let mut converted = convert_user_part(ns, &UserPart::File(part.clone()), 0)?; if converted["type"] == "input_file" && converted.get("file_data").is_some() { converted["filename"] = json!(part.filename.as_deref().unwrap_or("data")); } @@ -512,7 +829,10 @@ fn convert_tool_result_output( } ToolResultContent::Custom { .. } => { warnings.push(Warning::Other { - message: "unsupported tool content part type: custom".into(), + message: format!( + "unsupported {}tool content part type: custom", + if custom { "custom " } else { "" } + ), }); continue; } @@ -525,7 +845,26 @@ fn convert_tool_result_output( Ok(json!(parts)) } -/// Serialize tool-call arguments as JSON, mirroring JSON.stringify. +fn inline_file( + ns: ResponsesNamespace, + media_type: &str, + filename: Option<&str>, + data: &str, + provider_options: &Option, + index: usize, +) -> Value { + if media_type.split('/').next() == Some("image") { + let mut part = json!({ "type": "input_image", "image_url": format!("data:{media_type};base64,{data}") }); + if let Some(detail) = openai_sub_option(ns, provider_options, "imageDetail") { + part["detail"] = detail; + } + part + } else { + json!({ "type": "input_file", "filename": filename.map(str::to_owned).unwrap_or_else(|| if media_type == "application/pdf" { format!("part-{index}.pdf") } else { format!("part-{index}") }), "file_data": format!("data:{media_type};base64,{data}") }) + } +} + +/// Serialize tool-call arguments as JSON, including null and string values. fn serialize_arguments(input: &Value) -> String { input.to_string() } @@ -592,62 +931,495 @@ pub struct PreparedResponsesTools { pub tool_warnings: Vec, } -/// Prepare `FunctionTool`s into the Responses `tools` / `tool_choice` JSON shape. +/// Prepare the upstream Responses function and provider tool family. /// -/// Mirrors the function-tool path of TS `prepareResponsesTools`. Built-in -/// provider tools (web_search, file_search, etc.) are out of scope for the -/// core implementation and emit an `unsupported` warning. -#[must_use] +/// # Errors +/// Returns an error for invalid provider tool arguments or conflicting namespaces. pub fn prepare_responses_tools( tools: &Option>, tool_choice: Option<&ToolChoice>, -) -> PreparedResponsesTools { - let non_empty = tools.as_ref().filter(|t| !t.is_empty()); - let mut tool_warnings: Vec = Vec::new(); - - let tools_opt = match non_empty { - None => None, - Some(tools) => { - let mut openai_tools: Vec = Vec::new(); - for t in tools { - match t { - Tool::Function(ft) => { - openai_tools.push(function_tool_to_json(ft)); +) -> Result { + prepare_responses_tools_for(tools, tool_choice, None, true) +} + +fn prepare_responses_tools_for( + tools: &Option>, + tool_choice: Option<&ToolChoice>, + allowed_tools: Option<&Value>, + supports_async: bool, +) -> Result { + let Some(tools) = tools.as_ref().filter(|tools| !tools.is_empty()) else { + return Ok(PreparedResponsesTools { + tools: None, + tool_choice: None, + tool_warnings: Vec::new(), + }); + }; + let mut tool_warnings = Vec::new(); + let mut prepared: Vec = Vec::new(); + for tool in tools { + match tool { + Tool::Function(ft) => { + let mut value = function_tool_to_json(ft); + value["parameters"] = crate::openai::convert::normalize_json_schema( + &ft.input_schema, + &mut tool_warnings, + )?; + let opts = ft + .provider_options + .as_ref() + .and_then(|options| options.get("openai")); + if let Some(opts) = opts { + for (key, wire) in [ + ("async", "async"), + ("deferLoading", "defer_loading"), + ("allowedCallers", "allowed_callers"), + ] { + if let Some(v) = opts.get(key) { + value[wire] = v.clone(); + } } - Tool::Provider(pt) => { - tool_warnings.push(Warning::Unsupported { - feature: format!("provider-defined tool {}", pt.id), - details: None, - }); + if let Some(schema) = opts.get("outputSchema") { + value["output_schema"] = crate::openai::convert::normalize_json_schema( + schema, + &mut tool_warnings, + )?; + } + resolve_async_tool_option( + &mut value, + supports_async, + &ft.name, + &mut tool_warnings, + ); + if let Some(namespace) = opts.get("namespace") { + let name = + namespace + .get("name") + .and_then(Value::as_str) + .ok_or_else(|| { + AiMuxError::InvalidArgument("tool namespace name".into()) + })?; + let description = namespace + .get("description") + .and_then(Value::as_str) + .ok_or_else(|| { + AiMuxError::InvalidArgument("tool namespace description".into()) + })?; + if let Some(group) = prepared + .iter_mut() + .find(|group| group["type"] == "namespace" && group["name"] == name) + { + if group["description"] != description { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "conflicting descriptions for OpenAI tool namespace {name}" + ))); + } + if let Some(items) = group["tools"].as_array_mut() { + items.push(value); + } + } else { + prepared.push(json!({"type":"namespace", "name":name, "description":description, "tools":[value]})); + } + continue; } } + prepared.push(value); } - if openai_tools.is_empty() { - None - } else { - Some(openai_tools) + Tool::Provider(tool) => { + if let Some(mut value) = provider_tool_to_json(tool)? { + resolve_async_tool_option( + &mut value, + supports_async, + &tool.name, + &mut tool_warnings, + ); + prepared.push(value); + } } } - }; + } + let mut tool_choice = tool_choice.map(|choice| match choice { + ToolChoice::Auto => json!("auto"), + ToolChoice::None => json!("none"), + ToolChoice::Required => json!("required"), + ToolChoice::Tool { tool_name } => { + let tool = tools.iter().find_map(|tool| match tool { + Tool::Provider(tool) if tool.name == *tool_name => Some(tool), + _ => None, + }); + match tool.map(|tool| tool.id.as_str()) { + Some("openai.custom") => json!({"type":"custom", "name":tool_name}), + Some(id) + if [ + "openai.code_interpreter", + "openai.file_search", + "openai.image_generation", + "openai.web_search_preview", + "openai.web_search", + "openai.mcp", + "openai.apply_patch", + "openai.computer", + "openai.programmatic_tool_calling", + ] + .contains(&id) => + { + json!({"type":id.trim_start_matches("openai.")}) + } + _ => json!({"type":"function", "name":tool_name}), + } + } + }); + if let Some(allowed) = allowed_tools { + tool_choice = Some(prepare_allowed_tools(allowed, tools, &mut tool_warnings)?); + } + Ok(PreparedResponsesTools { + tools: Some(prepared), + tool_choice, + tool_warnings, + }) +} + +fn resolve_async_tool_option( + value: &mut Value, + supported: bool, + name: &str, + warnings: &mut Vec, +) { + if !supported && value.get("async") == Some(&json!(true)) { + if let Some(object) = value.as_object_mut() { + object.remove("async"); + } + warnings.push(Warning::Unsupported { + feature: format!("async tool calling for {name}"), + details: Some("Async tool calling is not supported by this model.".into()), + }); + } +} - let tool_choice_opt = match (&tools_opt, tool_choice) { - (None, _) => None, - (Some(_), None) => None, - (Some(_), Some(tc)) => match tc { - ToolChoice::Auto => Some(json!("auto")), - ToolChoice::None => Some(json!("none")), - ToolChoice::Required => Some(json!("required")), - ToolChoice::Tool { tool_name } => { - Some(json!({ "type": "function", "name": tool_name })) +fn prepare_allowed_tools( + allowed: &Value, + tools: &[Tool], + warnings: &mut Vec, +) -> Result { + let invalid = || AiMuxError::InvalidArgument("allowedTools".into()); + let names = allowed + .get("toolNames") + .and_then(Value::as_array) + .filter(|names| !names.is_empty()) + .ok_or_else(invalid)?; + let mode = allowed + .get("mode") + .map(|mode| { + mode.as_str() + .filter(|mode| ["auto", "required"].contains(mode)) + .ok_or_else(invalid) + }) + .transpose()? + .unwrap_or("auto"); + let mut entries = Vec::new(); + for name in names { + let name = name.as_str().ok_or_else(invalid)?; + let direct = tools.iter().find(|tool| match tool { + Tool::Function(tool) => tool.name == name, + Tool::Provider(tool) => tool.name == name, + }); + let aliases: Vec<_> = tools.iter().filter(|tool| matches!(tool, Tool::Provider(tool) if tool.id.strip_prefix("openai.") == Some(name) && tool.name != name)).collect(); + let tool = if let Some(tool) = direct { + if !aliases.is_empty() { + warnings.push(Warning::Unsupported { feature:format!("allowedTools entry {name}"), details:Some("this name is both a tool name and the provider tool name of another tool in this request; the tool with this name is allowed".into()) }); } - }, + Some(tool) + } else { + if aliases.len() > 1 + && aliases.iter().any(|tool| match tool { + Tool::Provider(tool) => tool.id == "openai.custom" || tool.id == "openai.mcp", + _ => false, + }) + { + warnings.push(Warning::Unsupported { feature:format!("allowedTools entry {name}"), details:Some("several tools share this provider tool name; use the tool name from this request".into()) }); + continue; + } + aliases.first().copied() + }; + let entry = match tool { + Some(Tool::Function(tool)) => { + let opts = tool + .provider_options + .as_ref() + .and_then(|options| options.get("openai")); + if opts.is_some_and(|opts| { + opts.contains_key("namespace") || opts.get("deferLoading") == Some(&json!(true)) + }) { + warnings.push(Warning::Unsupported { feature:format!("allowedTools entry {name}"), details:Some("namespace and deferred tools are not visible to tool_choice.allowed_tools; the tool is removed from the allowed tools".into()) }); + continue; + } + json!({"type":"function", "name":tool.name}) + } + Some(Tool::Provider(tool)) => { + let Some(value) = provider_tool_to_json(tool)? else { + entries.push(json!({"type":"function", "name":name})); + continue; + }; + match value["type"].as_str() { + Some("custom") => json!({"type":"custom", "name":tool.name}), + Some("mcp") => json!({"type":"mcp", "server_label":value["server_label"]}), + Some("tool_search") => { + warnings.push(Warning::Unsupported { feature:format!("allowedTools entry {name}"), details:Some("tool_search is not visible to tool_choice.allowed_tools; the tool is removed from the allowed tools".into()) }); + continue; + } + Some(kind) => json!({"type":kind}), + None => continue, + } + } + None => { + warnings.push(Warning::Unsupported { feature:format!("allowedTools entry {name}"), details:Some("the tool is not part of the tools for this request and is sent as a function tool".into()) }); + json!({"type":"function", "name":name}) + } + }; + entries.push(entry); + } + if entries.is_empty() { + return Err(AiMuxError::UnsupportedFunctionality( + "allowedTools with only tools that cannot be allow-listed".into(), + )); + } + Ok(json!({"type":"allowed_tools", "mode":mode, "tools":entries})) +} + +// Only configured fields are sent: absent optional values are omitted, not null. +fn map_tool_fields(value: &Value, fields: &[(&str, &str)]) -> Value { + let mut result = json!({}); + for (name, wire) in fields { + if let Some(value) = value.get(*name) { + result[*wire] = value.clone(); + } + } + result +} + +fn provider_tool_to_json( + tool: &aimux_core::tool::ProviderTool, +) -> Result, AiMuxError> { + let kind = tool.id.strip_prefix("openai.").unwrap_or(""); + if ![ + "file_search", + "local_shell", + "shell", + "apply_patch", + "computer", + "web_search_preview", + "web_search", + "code_interpreter", + "image_generation", + "mcp", + "custom", + "programmatic_tool_calling", + "tool_search", + ] + .contains(&kind) + { + return Ok(None); + } + if [ + "local_shell", + "apply_patch", + "computer", + "programmatic_tool_calling", + ] + .contains(&kind) + { + return Ok(Some(json!({"type":kind}))); + } + let args = &Value::Object(tool.args.clone()); + validate_tool_args(kind, args)?; + let fields: &[(&str, &str)] = match kind { + "file_search" => &[ + ("vectorStoreIds", "vector_store_ids"), + ("maxNumResults", "max_num_results"), + ("filters", "filters"), + ], + "web_search_preview" | "web_search" => &[ + ("searchContextSize", "search_context_size"), + ("userLocation", "user_location"), + ], + "image_generation" => &[ + ("action", "action"), + ("background", "background"), + ("inputFidelity", "input_fidelity"), + ("model", "model"), + ("moderation", "moderation"), + ("partialImages", "partial_images"), + ("quality", "quality"), + ("outputCompression", "output_compression"), + ("outputFormat", "output_format"), + ("size", "size"), + ], + "mcp" => &[ + ("serverLabel", "server_label"), + ("authorization", "authorization"), + ("connectorId", "connector_id"), + ("headers", "headers"), + ("serverDescription", "server_description"), + ("serverUrl", "server_url"), + ], + "custom" => &[ + ("description", "description"), + ("format", "format"), + ("async", "async"), + ], + "tool_search" => &[ + ("execution", "execution"), + ("description", "description"), + ("parameters", "parameters"), + ], + _ => &[], }; + let mut value = map_tool_fields(args, fields); + value["type"] = json!(kind); + if let Some(location) = args.get("userLocation") { + value["user_location"] = map_tool_fields( + location, + &[ + ("type", "type"), + ("country", "country"), + ("city", "city"), + ("region", "region"), + ("timezone", "timezone"), + ], + ); + } + match kind { + "file_search" => { + if let Some(ranking) = args.get("ranking") { + value["ranking_options"] = map_tool_fields( + ranking, + &[("ranker", "ranker"), ("scoreThreshold", "score_threshold")], + ); + } + } + "web_search" => { + if let Some(filters) = args.get("filters") { + value["filters"] = map_tool_fields( + filters, + &[ + ("allowedDomains", "allowed_domains"), + ("blockedDomains", "blocked_domains"), + ], + ); + } + if let Some(access) = args.get("externalWebAccess") { + value["external_web_access"] = access.clone(); + } + } + "code_interpreter" => { + value["container"] = match args.get("container") { + Some(container) if container.is_string() => container.clone(), + container => { + let mut value = container + .map(|container| map_tool_fields(container, &[("fileIds", "file_ids")])) + .unwrap_or_else(|| json!({})); + value["type"] = json!("auto"); + value + } + } + } + "image_generation" => { + if let Some(mask) = args.get("inputImageMask") { + value["input_image_mask"] = + map_tool_fields(mask, &[("fileId", "file_id"), ("imageUrl", "image_url")]); + } + } + "custom" => { + value["name"] = json!(tool.name); + if let Some(format) = args.get("format") { + value["format"] = map_tool_fields( + format, + &[ + ("type", "type"), + ("syntax", "syntax"), + ("definition", "definition"), + ], + ); + } + } + "mcp" => { + if let Some(allowed) = args.get("allowedTools") { + value["allowed_tools"] = if allowed.is_array() { + allowed.clone() + } else { + map_tool_fields( + allowed, + &[("readOnly", "read_only"), ("toolNames", "tool_names")], + ) + }; + } + value["require_approval"] = args + .get("requireApproval") + .map(|approval| { + if approval.is_string() { + approval.clone() + } else if let Some(never) = approval.get("never") { + json!({"never":map_tool_fields(never, &[("toolNames", "tool_names")])}) + } else { + json!("never") + } + }) + .unwrap_or_else(|| json!("never")); + } + "shell" => { + if let Some(environment) = args.get("environment") { + value["environment"] = map_shell_environment(environment)?; + } + } + _ => {} + } + Ok(Some(value)) +} - PreparedResponsesTools { - tools: tools_opt, - tool_choice: tool_choice_opt, - tool_warnings, +fn map_shell_environment(environment: &Value) -> Result { + let kind = environment + .get("type") + .and_then(Value::as_str) + .unwrap_or("local"); + let mut value = match kind { + "containerReference" => map_tool_fields(environment, &[("containerId", "container_id")]), + "containerAuto" => map_tool_fields( + environment, + &[("fileIds", "file_ids"), ("memoryLimit", "memory_limit")], + ), + _ => map_tool_fields(environment, &[("skills", "skills")]), + }; + value["type"] = json!(match kind { + "containerReference" => "container_reference", + "containerAuto" => "container_auto", + _ => "local", + }); + if kind == "containerAuto" { + if let Some(policy) = environment.get("networkPolicy") { + value["network_policy"] = map_tool_fields( + policy, + &[ + ("type", "type"), + ("allowedDomains", "allowed_domains"), + ("domainSecrets", "domain_secrets"), + ], + ); + } + if let Some(skills) = environment.get("skills").and_then(Value::as_array) { + value["skills"] = Value::Array(skills.iter().map(|skill| { + if skill["type"] == "skillReference" { + let id = skill.get("providerReference").and_then(|reference| reference.get("openai")).and_then(Value::as_str).ok_or_else(|| AiMuxError::InvalidArgument("No skill reference for openai".into()))?; + Ok(json!({"type":"skill_reference", "skill_id":id, "version":skill.get("version").cloned().unwrap_or_else(|| json!("latest"))})) + } else { + let mut value = map_tool_fields(skill, &[("name", "name"), ("description", "description")]); + value["type"] = json!("inline"); + value["source"] = map_tool_fields(&skill["source"], &[("type", "type"), ("mediaType", "media_type"), ("data", "data")]); + Ok(value) + } + }).collect::, AiMuxError>>()?); + } } + Ok(value) } /// Convert a `FunctionTool` into the Responses `function` tool JSON shape. @@ -660,9 +1432,7 @@ fn function_tool_to_json(t: &FunctionTool) -> Value { if let Some(ref desc) = t.description { func["description"] = json!(desc); } - if let Some(strict) = t.strict { - func["strict"] = json!(strict); - } + func["strict"] = json!(t.strict.unwrap_or(false)); func } @@ -731,23 +1501,13 @@ fn resolve_responses_reasoning( } }); - let resolved_reasoning_summary: Option = - openai_option(ns, provider_opts, "reasoningSummary") - .map(|v| { - v.as_str() - .map(std::string::ToString::to_string) - .unwrap_or_else(|| v.to_string()) - }) - .or_else(|| { - if resolved_reasoning_effort - .as_deref() - .is_some_and(|e| e != "none") - { - Some("detailed".to_string()) - } else { - None - } - }); + let resolved_reasoning_summary = match openai_option(ns, provider_opts, "reasoningSummary") { + Some(value) => value.as_str().map(str::to_owned), + None => resolved_reasoning_effort + .as_deref() + .filter(|effort| *effort != "none") + .map(|_| "detailed".to_string()), + }; let is_reasoning_model = openai_option(ns, provider_opts, "forceReasoning") .map(|v| v.as_bool().unwrap_or(false)) @@ -861,6 +1621,7 @@ fn apply_responses_text_format( body: &mut Value, options: &CallOptions, provider_opts: &Option, + warnings: &mut Vec, ) { if let Some(ref rf) = options.response_format { match rf { @@ -880,9 +1641,11 @@ fn apply_responses_text_format( "type": "json_schema", "strict": strict_json, "name": name.clone().unwrap_or_else(|| "response".to_string()), - "description": description, - "schema": schema, + "schema": crate::openai::convert::normalize_json_schema(schema, warnings).unwrap_or_else(|_| schema.clone()), }); + if let Some(description) = description { + text["format"]["description"] = json!(description); + } } None => { text["format"] = json!({ "type": "json_object" }); @@ -912,6 +1675,7 @@ fn resolve_responses_include( ns: ResponsesNamespace, provider_opts: &Option, is_reasoning_model: bool, + tools: Option<&[Tool]>, ) -> Option> { let mut include: Option> = openai_option(ns, provider_opts, "include").and_then(|v| v.as_array().cloned()); @@ -934,6 +1698,27 @@ fn resolve_responses_include( add_include("reasoning.encrypted_content", &mut include); } + let top_logprobs = openai_option(ns, provider_opts, "logprobs") + .and_then(|v| if v == true { Some(20) } else { v.as_u64() }); + if top_logprobs.is_some_and(|value| value > 0) { + add_include("message.output_text.logprobs", &mut include); + } + for tool in tools.into_iter().flatten() { + if let Tool::Provider(tool) = tool { + match tool.id.as_str() { + "openai.web_search" | "openai.web_search_preview" + if openai_option(ns, provider_opts, "includeWebSearchSources") + != Some(json!(false)) => + { + add_include("web_search_call.action.sources", &mut include) + } + "openai.code_interpreter" => { + add_include("code_interpreter_call.outputs", &mut include) + } + _ => {} + } + } + } include } @@ -968,6 +1753,31 @@ fn apply_responses_provider_options( .collect(), ); } + if let Some(value) = openai_option(ns, provider_opts, "logprobs").and_then(|v| { + if v == true { + Some(json!(20)) + } else { + v.as_u64().map(|n| json!(n)) + } + }) { + body["top_logprobs"] = value; + } + if let Some(values) = + openai_option(ns, provider_opts, "contextManagement").and_then(|v| v.as_array().cloned()) + { + body["context_management"] = json!( + values + .into_iter() + .map(|v| { + let mut entry = json!({ "type": v["type"] }); + if let Some(threshold) = v.get("compactThreshold") { + entry["compact_threshold"] = threshold.clone(); + } + entry + }) + .collect::>() + ); + } } /// `service_tier` with model-capability validation. @@ -991,7 +1801,7 @@ fn apply_responses_service_tier( ), }); } - "priority" if !caps.supports_priority_processing => { + "priority" | "fast" if !caps.supports_priority_processing => { warnings.push(Warning::Unsupported { feature: "serviceTier".to_string(), details: Some("priority processing is only available for supported models (gpt-4, gpt-5, gpt-5-mini, o3, o4-mini) and requires Enterprise access. gpt-5-nano is not supported".to_string()), @@ -1041,13 +1851,10 @@ fn apply_responses_reasoning_block( } } -/// Build the OpenAI Responses API request body (without warnings). -/// -/// Splits the original ~380-line function into focused helpers (issue M11); -/// behavior is unchanged. +/// Build the OpenAI Responses API request body and warnings. /// /// # Errors -/// Returns an error for text file parts or an unresolved inline image media type. +/// Rejects unsupported prompt parts and invalid tool arguments. pub fn build_responses_request_body( model_id: &str, options: &CallOptions, @@ -1060,7 +1867,7 @@ pub fn build_responses_request_body( /// namespace. /// /// # Errors -/// Returns an error for text file parts or an unresolved inline image media type. +/// Rejects unsupported prompt parts and invalid tool arguments. pub fn build_responses_request_body_for( ns: ResponsesNamespace, model_id: &str, @@ -1075,8 +1882,24 @@ pub fn build_responses_request_body_for( push_unsupported_call_option_warnings(options, &mut warnings); // -- Reasoning resolution -- - let (resolved_reasoning_effort, resolved_reasoning_summary, is_reasoning_model) = + let (mut resolved_reasoning_effort, mut resolved_reasoning_summary, is_reasoning_model) = resolve_responses_reasoning(ns, provider_opts, options, &caps); + if let Some(efforts) = caps.supported_reasoning_efforts + && resolved_reasoning_effort + .as_deref() + .is_some_and(|effort| !efforts.contains(&effort)) + { + warnings.push(Warning::Unsupported { + feature: "reasoningEffort".to_string(), + details: Some(format!( + "{model_id} only supports the following reasoning efforts: {}", + efforts.join(", ") + )), + }); + resolved_reasoning_effort = None; + resolved_reasoning_summary = openai_option(ns, provider_opts, "reasoningSummary") + .and_then(|v| v.as_str().map(str::to_owned)); + } // -- conversation + previousResponseId conflict -- warn_conversation_conflict(ns, provider_opts, &mut warnings); @@ -1090,15 +1913,54 @@ pub fn build_responses_request_body_for( .and_then(|v| v.as_bool()) .unwrap_or(true); let has_previous_response_id = openai_option(ns, provider_opts, "previousResponseId").is_some(); - let has_conversation = openai_option(ns, provider_opts, "conversation").is_some(); - let input_result = convert_to_responses_input( + let mut input_result = convert_responses_input_with_conversation( ns, &options.prompt, system_message_mode, store_bool, has_previous_response_id, - has_conversation, + openai_option(ns, provider_opts, "conversation").is_some(), + options.tools.as_deref(), )?; + if let Some(effort) = openai_option(ns, provider_opts, "reasoningEffortUpdate") + .and_then(|v| v.as_str().map(str::to_owned)) + { + let reason = if !caps.supports_configuration_update { + Some("reasoningEffortUpdate is only supported by GPT-6 and later models".to_string()) + } else if openai_option(ns, provider_opts, "reasoningMode") == Some(json!("pro")) + || openai_option(ns, provider_opts, "contextManagement").is_some() + || openai_option(ns, provider_opts, "truncation") == Some(json!("auto")) + { + Some("reasoningEffortUpdate requires standard reasoning mode without automatic compaction or automatic truncation".to_string()) + } else if caps + .supported_reasoning_efforts + .is_some_and(|values| !values.contains(&effort.as_str())) + { + Some(format!( + "{model_id} does not support reasoning effort {effort}" + )) + } else { + None + }; + if let Some(details) = reason { + warnings.push(Warning::Unsupported { + feature: "reasoningEffortUpdate".to_string(), + details: Some(details), + }); + } else if input_result.input.first().is_none_or(|item| { + item["type"] != "configuration_update" || item["reasoning"]["effort"] != effort + }) { + input_result.input.insert( + 0, + json!({ "type": "configuration_update", "reasoning": { "effort": effort } }), + ); + } + } + if openai_option(ns, provider_opts, "compactionTrigger") == Some(json!(true)) { + input_result + .input + .push(json!({ "type": "compaction_trigger" })); + } warnings.extend(input_result.warnings); // -- Base body -- @@ -1127,10 +1989,15 @@ pub fn build_responses_request_body_for( ); // -- Response format (text.format) + verbosity -- - apply_responses_text_format(ns, &mut body, options, provider_opts); + apply_responses_text_format(ns, &mut body, options, provider_opts, &mut warnings); // -- include (computed) -- - let include = resolve_responses_include(ns, provider_opts, is_reasoning_model); + let include = resolve_responses_include( + ns, + provider_opts, + is_reasoning_model, + options.tools.as_deref(), + ); if let Some(inc) = include { body["include"] = json!(inc); } @@ -1142,6 +2009,47 @@ pub fn build_responses_request_body_for( // -- Other provider options (only sent when set) -- apply_responses_provider_options(ns, &mut body, provider_opts); + if caps.supports_configuration_update && body.get("prompt_cache_retention").is_some() { + body.as_object_mut() + .expect("request body is an object") + .remove("prompt_cache_retention"); + warnings.push(Warning::Unsupported { + feature: "promptCacheRetention".to_string(), + details: Some("promptCacheRetention is not supported by GPT-6 and later models; use promptCacheOptions instead".to_string()), + }); + } + + if is_reasoning_model + && caps.supported_reasoning_efforts.is_some() + && !(resolved_reasoning_effort.as_deref() == Some("none") + && caps.supports_non_reasoning_parameters) + { + let mut removed = body + .as_object_mut() + .expect("request body is an object") + .remove("top_logprobs") + .is_some(); + if let Some(include) = body.get_mut("include").and_then(Value::as_array_mut) { + let len = include.len(); + include.retain(|value| value.as_str() != Some("message.output_text.logprobs")); + removed |= include.len() != len; + } + if body + .get("include") + .and_then(Value::as_array) + .is_some_and(Vec::is_empty) + { + body.as_object_mut() + .expect("request body is an object") + .remove("include"); + } + if removed { + warnings.push(Warning::Unsupported { + feature: "logprobs".to_string(), + details: Some("logprobs is not supported for reasoning models".to_string()), + }); + } + } // -- service_tier (with capability validation) -- apply_responses_service_tier(ns, &mut body, provider_opts, &caps, &mut warnings); @@ -1157,7 +2065,13 @@ pub fn build_responses_request_body_for( ); // -- Tools -- - let prepared = prepare_responses_tools(&options.tools, options.tool_choice.as_ref()); + let allowed_tools = openai_option(ns, provider_opts, "allowedTools"); + let prepared = prepare_responses_tools_for( + &options.tools, + options.tool_choice.as_ref(), + allowed_tools.as_ref(), + caps.supports_configuration_update, + )?; if let Some(tools) = prepared.tools { body["tools"] = json!(tools); if let Some(tc) = prepared.tool_choice { @@ -1232,7 +2146,7 @@ fn apply_prefix_to_part(part: &mut Value, prefixes: &[&str]) { obj.remove("file_data"); obj.remove("filename"); obj.insert("file_id".to_string(), json!(data)); - if get_top_level_media_type(media_type) == "image" { + if media_type.starts_with("image/") { obj.insert("type".to_string(), json!("input_image")); } } diff --git a/aimux-providers/src/openai/responses/mod.rs b/aimux-providers/src/openai/responses/mod.rs index b9934f070..1bd543a18 100644 --- a/aimux-providers/src/openai/responses/mod.rs +++ b/aimux-providers/src/openai/responses/mod.rs @@ -23,7 +23,9 @@ //! [`responses_convert`] (RFC-0012 §3.5). pub mod convert; +mod provider_events; pub mod responses_convert; +pub(crate) mod tool_args; pub mod types; pub(crate) use convert::ResponsesProfile; @@ -84,9 +86,154 @@ impl OpenAIResponsesModel { stream: bool, ) -> Result { let profile = &self.config.responses; + let mut unused_warnings = Vec::new(); + if let Some(aimux_core::options::ResponseFormat::Json { + schema: Some(schema), + .. + }) = &options.response_format + { + super::convert::normalize_json_schema(schema, &mut unused_warnings)?; + } + for tool in options.tools.iter().flatten() { + if let aimux_core::tool::Tool::Function(tool) = tool { + super::convert::normalize_json_schema(&tool.input_schema, &mut unused_warnings)?; + } + } + let pass_through = options + .provider_options + .as_ref() + .and_then(|v| profile.namespace.find_in(v)) + .and_then(|v| v.get("passThroughUnsupportedFiles")) + .and_then(Value::as_bool) + == Some(true); + let caps = super::convert_common::get_model_capabilities(&self.model_id); + let provider_options = options + .provider_options + .as_ref() + .and_then(|v| profile.namespace.find_in(v)); + if provider_options + .and_then(|v| v.get("reasoningEffortUpdate")) + .is_some_and(|v| { + !matches!( + v.as_str(), + Some("none" | "low" | "medium" | "high" | "xhigh" | "max") + ) + }) + { + return Err(AiMuxError::InvalidArgument("reasoningEffortUpdate".into())); + } + let update_incompatible = provider_options.is_some_and(|v| { + v.get("reasoningMode") == Some(&serde_json::json!("pro")) + || v.get("contextManagement").is_some() + || v.get("truncation") == Some(&serde_json::json!("auto")) + }); + for message in &options.prompt { + use aimux_core::language_model_message::{ + AssistantPart, LanguageModelMessage, UserPart, + }; + use aimux_core::shared::FileData; + if let LanguageModelMessage::System { + content, + provider_options, + } = message + && let Some(effort) = provider_options + .as_ref() + .and_then(|v| profile.namespace.find_in(v)) + .and_then(|v| v.get("reasoningEffortUpdate")) + .and_then(Value::as_str) + { + if !matches!(effort, "none" | "low" | "medium" | "high" | "xhigh" | "max") { + return Err(AiMuxError::InvalidArgument("reasoningEffortUpdate".into())); + } + if !content.is_empty() + || !caps.supports_configuration_update + || update_incompatible + || caps + .supported_reasoning_efforts + .is_some_and(|values| !values.contains(&effort)) + { + return Err(AiMuxError::UnsupportedFunctionality( + "Message-level reasoningEffortUpdate".into(), + )); + } + } + let (files, user) = match message { + LanguageModelMessage::User { content, .. } => ( + content + .iter() + .filter_map(|part| match part { + UserPart::File(file) => Some(file), + UserPart::Text(_) => None, + }) + .collect::>(), + true, + ), + LanguageModelMessage::Assistant { content, .. } => ( + content + .iter() + .filter_map(|part| match part { + AssistantPart::File(file) => Some(file), + _ => None, + }) + .collect::>(), + false, + ), + _ => (Vec::new(), false), + }; + for file in files { + match &file.data { + FileData::Text { .. } if user => { + return Err(AiMuxError::UnsupportedFunctionality( + "text file parts".into(), + )); + } + FileData::Reference { reference } + if !reference.contains_key(profile.namespace.write_key()) => + { + return Err(AiMuxError::InvalidArgument(format!( + "No file reference for {}", + profile.namespace.write_key() + ))); + } + FileData::Data { .. } if user => { + aimux_provider_utils::resolve_full_media_type(file)?; + let media_type = &file.media_type; + if !pass_through + && !media_type.starts_with("image/") + && media_type != "image" + && media_type != "application/pdf" + { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "file part media type {media_type}" + ))); + } + } + _ => {} + } + } + } let mut result = build_responses_request_body_for(profile.namespace, &self.model_id, options, stream)?; convert::apply_file_id_prefixes(&mut result.body, &profile.file_id_prefixes); + if profile.explicit_message_item_type + && let Some(input) = result.body["input"].as_array_mut() + { + for item in input { + if item.get("role").is_some() { + item["type"] = serde_json::json!("message"); + } + } + } + if result.body["input"].as_array().is_some_and(|input| { + input.windows(2).any(|pair| { + pair[0]["type"] == "configuration_update" + && pair[1]["type"] == "configuration_update" + }) + }) { + return Err(AiMuxError::UnsupportedFunctionality( + "Adjacent reasoning effort configuration updates".into(), + )); + } result.body = self.config.transform_body(result.body); Ok(result) } @@ -134,7 +281,7 @@ impl LanguageModel for OpenAIResponsesModel { .unwrap_or_default(); let data = resp.value; - responses_convert::build_responses_generate_result( + responses_convert::build_responses_generate_result_with_tools( &data, &raw_body, request_result.warnings, @@ -142,6 +289,8 @@ impl LanguageModel for OpenAIResponsesModel { endpoint, body, response_headers, + tool_name_mapping(options), + prompt_approval_tool_call_ids(options), ) } @@ -180,7 +329,7 @@ impl LanguageModel for OpenAIResponsesModel { Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), first_event => first_event, }; - let stream = responses_convert::build_responses_event_stream( + let stream = responses_convert::build_responses_event_stream_with_tools( first_event, sse_stream, provider_key, @@ -189,6 +338,9 @@ impl LanguageModel for OpenAIResponsesModel { endpoint, body.clone(), response_headers.clone(), + options.include_raw_chunks == Some(true), + tool_name_mapping(options), + prompt_approval_tool_call_ids(options), )?; Ok(StreamResult { @@ -200,3 +352,51 @@ impl LanguageModel for OpenAIResponsesModel { }) } } + +fn tool_name_mapping(options: &CallOptions) -> HashMap { + let mut names = HashMap::new(); + for tool in options.tools.iter().flatten() { + if let aimux_core::tool::Tool::Provider(tool) = tool + && let Some(name) = tool.id.strip_prefix("openai.") + { + names.insert(name.to_owned(), tool.name.clone()); + } + } + if let Some(name) = options.tools.iter().flatten().find_map(|tool| match tool { + aimux_core::tool::Tool::Provider(tool) + if matches!( + tool.id.as_str(), + "openai.web_search" | "openai.web_search_preview" + ) => + { + Some(tool.name.clone()) + } + _ => None, + }) { + names.insert("web_search".into(), name); + } + names +} + +fn prompt_approval_tool_call_ids(options: &CallOptions) -> HashMap { + use aimux_core::language_model_message::{AssistantPart, LanguageModelMessage}; + + let mut mapping = HashMap::new(); + for message in &options.prompt { + if let LanguageModelMessage::Assistant { content, .. } = message { + for part in content { + if let AssistantPart::ToolCall(call) = part + && let Some(id) = call + .provider_options + .as_ref() + .and_then(|options| options.get("openai")) + .and_then(|options| options.get("approvalRequestId")) + .and_then(Value::as_str) + { + mapping.insert(id.to_owned(), call.tool_call_id.clone()); + } + } + } + } + mapping +} diff --git a/aimux-providers/src/openai/responses/provider_events.rs b/aimux-providers/src/openai/responses/provider_events.rs new file mode 100644 index 000000000..cbd7ce64c --- /dev/null +++ b/aimux-providers/src/openai/responses/provider_events.rs @@ -0,0 +1,314 @@ +//! Provider tool output mapping from the pinned Responses implementation. + +use std::collections::HashMap; + +use aimux_core::result::{GenerateContent, RawToolApprovalRequest}; +use aimux_core::shared::provider_namespace; +use aimux_core::tool::{RawToolCall, ToolResult}; +use serde_json::{Value, json}; + +pub(crate) fn provider_tool_content( + part: &Value, + provider_key: &str, + tool_names: &HashMap, + shell_provider_executed: bool, + mcp_tool_call_id: Option<&str>, +) -> Vec { + let kind = part["type"].as_str().unwrap_or_default(); + let item_id = part["id"].as_str().unwrap_or_default(); + let call_id = part["call_id"].as_str().unwrap_or(item_id); + let metadata = || { + Some( + provider_namespace(provider_key, json!({"itemId": item_id})) + .expect("provider metadata must be an object"), + ) + }; + let name = |canonical: &str| { + tool_names + .get(canonical) + .cloned() + .unwrap_or_else(|| canonical.to_string()) + }; + let call = |id: &str, tool_name: String, input: String, executed, dynamic, meta| { + GenerateContent::ToolCall(RawToolCall { + tool_call_id: id.to_string(), + tool_name, + input, + provider_executed: executed, + dynamic, + provider_metadata: meta, + }) + }; + let result = |id: &str, tool_name: String, value: Value, meta| { + GenerateContent::ToolResult(ToolResult { + tool_call_id: id.to_string(), + tool_name, + result: value, + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: meta, + }) + }; + let pair = |canonical: &str, input: String, output: Value| { + vec![ + call(item_id, name(canonical), input, Some(true), None, None), + result(item_id, name(canonical), output, None), + ] + }; + match kind { + "image_generation_call" => pair( + "image_generation", + "{}".into(), + json!({"result": part["result"]}), + ), + "web_search_call" => pair( + "web_search", + "{}".into(), + web_search_output(&part["action"]), + ), + "file_search_call" => pair( + "file_search", + "{}".into(), + json!({ + "queries": part["queries"], + "results": part["results"].as_array().map(|results| results.iter().map(|value| { + json!({"attributes": value["attributes"], "fileId": value["file_id"], + "filename": value["filename"], "score": value["score"], "text": value["text"]}) + }).collect::>()), + }), + ), + "code_interpreter_call" => pair( + "code_interpreter", + json!({"code": part["code"], "containerId": part["container_id"]}).to_string(), + json!({"outputs": part["outputs"]}), + ), + "computer_call" if part["call_id"].is_null() => pair( + "computer_use", + String::new(), + json!({"type": "computer_use_tool_result", "status": part["status"]}), + ), + "computer_call" => vec![call( + call_id, + name("computer"), + computer_input(part).to_string(), + None, + None, + metadata(), + )], + "local_shell_call" => vec![call( + call_id, + name("local_shell"), + json!({"action": part["action"]}).to_string(), + None, + None, + metadata(), + )], + "shell_call" => vec![call( + call_id, + name("shell"), + json!({"action": {"commands": part["action"]["commands"]}}).to_string(), + shell_provider_executed.then_some(true), + None, + metadata(), + )], + "shell_call_output" => vec![result( + call_id, + name("shell"), + json!({"output": part["output"].as_array().map(|items| items.iter().map(|item| { + let outcome = if item["outcome"]["type"] == "exit" { + json!({"type": "exit", "exitCode": item["outcome"]["exit_code"]}) + } else { + json!({"type": "timeout"}) + }; + json!({"stdout": item["stdout"], "stderr": item["stderr"], "outcome": outcome}) + }).collect::>())}), + None, + )], + "apply_patch_call" => vec![call( + call_id, + name("apply_patch"), + json!({"callId": part["call_id"], "operation": part["operation"]}).to_string(), + None, + None, + metadata(), + )], + "program" => vec![call( + call_id, + name("programmatic_tool_calling"), + json!({"code": part["code"], "fingerprint": part["fingerprint"]}).to_string(), + Some(true), + None, + metadata(), + )], + "program_output" => vec![result( + call_id, + name("programmatic_tool_calling"), + json!({"result": part["result"], "status": part["status"]}), + metadata(), + )], + "tool_search_call" => { + let mut input = json!({}); + for field in ["arguments", "call_id"] { + if let Some(value) = part.get(field) { + input[field] = value.clone(); + } + } + vec![call( + call_id, + name("tool_search"), + input.to_string(), + (part["execution"] == "server").then_some(true), + None, + metadata(), + )] + } + "tool_search_output" => vec![result( + call_id, + name("tool_search"), + json!({"tools": part["tools"]}), + metadata(), + )], + "compaction" => vec![GenerateContent::Custom { + kind: "openai.compaction".into(), + provider_metadata: Some( + provider_namespace( + provider_key, + json!({ + "type": "compaction", "itemId": part["id"], + "encryptedContent": part["encrypted_content"], + }), + ) + .expect("provider metadata must be an object"), + ), + }], + "mcp_approval_request" => { + let tool_call_id = aimux_provider_utils::generate_id(); + let approval_id = part + .get("approval_request_id") + .and_then(Value::as_str) + .unwrap_or(item_id) + .to_owned(); + vec![ + call( + &tool_call_id, + format!("mcp.{}", part["name"].as_str().unwrap_or_default()), + part["arguments"].as_str().unwrap_or_default().to_owned(), + Some(true), + Some(true), + None, + ), + GenerateContent::ToolApprovalRequest(RawToolApprovalRequest { + approval_id, + tool_call_id, + provider_metadata: None, + }), + ] + } + "mcp_call" => { + let tool_name = format!("mcp.{}", part["name"].as_str().unwrap_or_default()); + let mut output = json!({"type": "call", "serverLabel": part["server_label"], + "name": part["name"], "arguments": part["arguments"]}); + for field in ["output", "error"] { + if let Some(value) = part.get(field).filter(|value| !value.is_null()) { + output[field] = value.clone(); + } + } + vec![ + call( + mcp_tool_call_id.unwrap_or(item_id), + tool_name.clone(), + part["arguments"].as_str().unwrap_or_default().to_string(), + Some(true), + Some(true), + None, + ), + result( + mcp_tool_call_id.unwrap_or(item_id), + tool_name, + output, + metadata(), + ), + ] + } + _ => Vec::new(), + } +} + +pub(crate) fn web_search_output(action: &Value) -> Value { + match action["type"].as_str() { + Some("search") => { + let mut output = json!({"action": {"type": "search"}}); + for field in ["query", "queries"] { + if let Some(value) = action.get(field).filter(|value| !value.is_null()) { + output["action"][field] = value.clone(); + } + } + if let Some(sources) = action.get("sources").filter(|value| !value.is_null()) { + output["sources"] = sources.clone(); + } + output + } + Some("open_page" | "find_in_page") => { + let kind = if action["type"] == "open_page" { + "openPage" + } else { + "findInPage" + }; + let mut output = json!({"action": {"type": kind}}); + for field in ["url", "pattern"] { + if field == "pattern" && kind == "openPage" { + continue; + } + if let Some(value) = action.get(field) { + output["action"][field] = value.clone(); + } + } + output + } + _ => json!({}), + } +} + +fn computer_input(part: &Value) -> Value { + let actions = part["actions"].as_array().cloned().unwrap_or_else(|| { + part.get("action") + .filter(|action| !action.is_null()) + .cloned() + .into_iter() + .collect() + }); + let actions: Vec<_> = actions + .into_iter() + .map(|mut action| { + if action["type"] == "scroll" + && let Some(object) = action.as_object_mut() + { + for (from, to) in [("scroll_x", "scrollX"), ("scroll_y", "scrollY")] { + if let Some(value) = object.remove(from) { + object.insert(to.to_string(), value); + } + } + } + action + }) + .collect(); + let checks: Vec<_> = part["pending_safety_checks"] + .as_array() + .map(|checks| { + checks + .iter() + .map(|check| { + let mut result = json!({"id": check["id"]}); + for field in ["code", "message"] { + if let Some(value) = check.get(field).filter(|value| !value.is_null()) { + result[field] = value.clone(); + } + } + result + }) + .collect() + }) + .unwrap_or_default(); + json!({"actions": actions, "pendingSafetyChecks": checks, "status": part["status"]}) +} diff --git a/aimux-providers/src/openai/responses/responses_convert.rs b/aimux-providers/src/openai/responses/responses_convert.rs index 610823ed1..eef1f6eb1 100644 --- a/aimux-providers/src/openai/responses/responses_convert.rs +++ b/aimux-providers/src/openai/responses/responses_convert.rs @@ -14,7 +14,7 @@ //! function — only the shared framework is extracted. use aimux_core::tool::RawToolCall; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet, VecDeque}; use std::pin::Pin; use futures::{Stream, StreamExt}; @@ -26,7 +26,7 @@ use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Sourc use aimux_core::shared::provider_namespace; use aimux_core::stream_part::StreamPart; use aimux_core::types::{ - FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, Warning, + FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Warning, }; use super::convert::{convert_responses_usage, map_responses_finish_reason, parse_usage}; @@ -80,6 +80,31 @@ pub fn build_responses_generate_result( request_url: String, body: Value, response_headers: HashMap, +) -> Result { + build_responses_generate_result_with_tools( + data, + raw_body, + request_warnings, + provider_key, + request_url, + body, + response_headers, + HashMap::new(), + HashMap::new(), + ) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn build_responses_generate_result_with_tools( + data: &Value, + raw_body: &str, + request_warnings: Vec, + provider_key: String, + request_url: String, + body: Value, + response_headers: HashMap, + tool_names: HashMap, + prompt_approval_tool_call_ids: HashMap, ) -> Result { // Top-level error field. if let Some(err_obj) = data.get("error") @@ -126,8 +151,11 @@ pub fn build_responses_generate_result( }) })?; + let shell_provider_executed = hosted_shell(&body); let mut content: Vec = Vec::new(); let mut has_function_call = false; + let mut logprobs = Vec::new(); + let mut hosted_tool_search_ids = VecDeque::new(); for part in output { match part.get("type").and_then(|v| v.as_str()) { @@ -140,15 +168,27 @@ pub fn build_responses_generate_result( .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); - if !text.is_empty() { + if body + .get("top_logprobs") + .and_then(Value::as_u64) + .is_some_and(|n| n > 0) + && let Some(value) = cp.get("logprobs") + { + logprobs.push(value.clone()); + } + { content.push(GenerateContent::Text { text, - provider_metadata: Some(provider_namespace( - &provider_key, - json!({ - "itemId": part.get("id").cloned().unwrap_or(Value::Null), - }), - ).expect("provider metadata must be an object")), + provider_metadata: Some( + provider_namespace( + &provider_key, + text_metadata( + part, + cp.get("annotations").and_then(Value::as_array), + ), + ) + .expect("provider metadata must be an object"), + ), }); } } @@ -189,13 +229,8 @@ pub fn build_responses_generate_result( provider_executed: None, dynamic: None, provider_metadata: Some( - provider_namespace( - &provider_key, - json!({ - "itemId": part.get("id").cloned().unwrap_or(Value::Null), - }), - ) - .expect("provider metadata must be an object"), + provider_namespace(&provider_key, tool_metadata(part)) + .expect("provider metadata must be an object"), ), })); } @@ -221,13 +256,8 @@ pub fn build_responses_generate_result( provider_executed: None, dynamic: None, provider_metadata: Some( - provider_namespace( - &provider_key, - json!({ - "itemId": part.get("id").cloned().unwrap_or(Value::Null), - }), - ) - .expect("provider metadata must be an object"), + provider_namespace(&provider_key, tool_metadata(part)) + .expect("provider metadata must be an object"), ), })); } @@ -260,52 +290,23 @@ pub fn build_responses_generate_result( } } } - Some("compaction") => { - content.push(GenerateContent::Custom { - kind: "openai.compaction".to_string(), - provider_metadata: Some( - provider_namespace( - &provider_key, - json!({ - "type": "compaction", "itemId": part.get("id"), - "encryptedContent": part.get("encrypted_content"), - }), - ) - .expect("provider metadata must be an object"), - ), - }); - } - Some("mcp_approval_request") => { - let tool_call_id = generate_source_id(); - content.push(GenerateContent::ToolCall(RawToolCall { - tool_call_id: tool_call_id.clone(), - tool_name: format!( - "mcp.{}", - part.get("name").and_then(Value::as_str).unwrap_or("") - ), - input: part - .get("arguments") + _ => { + let mut part = part.clone(); + pair_tool_search_id(&mut part, &mut hosted_tool_search_ids); + let parts = super::provider_events::provider_tool_content( + &part, + &provider_key, + &tool_names, + shell_provider_executed, + part.get("approval_request_id") .and_then(Value::as_str) - .unwrap_or("") - .to_string(), - provider_executed: Some(true), - dynamic: Some(true), - provider_metadata: None, - })); - content.push(GenerateContent::ToolApprovalRequest( - aimux_core::result::RawToolApprovalRequest { - approval_id: part - .get("approval_request_id") - .or_else(|| part.get("id")) - .and_then(Value::as_str) - .unwrap_or("") - .to_string(), - tool_call_id, - provider_metadata: None, - }, - )); + .and_then(|id| prompt_approval_tool_call_ids.get(id)) + .map(String::as_str), + ); + has_function_call |= part["type"] == "apply_patch_call" + || (part["type"] == "computer_call" && !part["call_id"].is_null()); + content.extend(parts); } - _ => {} } } @@ -331,6 +332,9 @@ pub fn build_responses_generate_result( if let Some(st) = data.get("service_tier").and_then(|v| v.as_str()) { pm["serviceTier"] = json!(st); } + if !logprobs.is_empty() { + pm["logprobs"] = json!(logprobs); + } let provider_metadata = Some(provider_namespace(&provider_key, pm).expect("provider metadata must be an object")); @@ -396,13 +400,38 @@ fn reasoning_stream_metadata( item_id: &str, encrypted_content: Option<&str>, ) -> ProviderMetadata { - let mut inner = json!({ "itemId": item_id }); - if let Some(enc) = encrypted_content { - inner["reasoningEncryptedContent"] = json!(enc); - } + let inner = json!({ "itemId": item_id, "reasoningEncryptedContent": encrypted_content }); provider_namespace(provider_key, inner).expect("provider metadata must be an object") } +fn tool_metadata(item: &Value) -> Value { + let mut metadata = json!({ "itemId": item.get("id").cloned().unwrap_or(Value::Null) }); + for key in ["async", "namespace"] { + if let Some(value) = item.get(key).filter(|v| !v.is_null()) { + metadata[key] = value.clone(); + } + } + if let Some(caller) = item.get("caller").filter(|value| !value.is_null()) { + metadata["caller"] = if caller["type"] == "program" { + json!({"type": "program", "callerId": caller["caller_id"]}) + } else { + caller.clone() + }; + } + metadata +} + +fn text_metadata(item: &Value, annotations: Option<&Vec>) -> Value { + let mut metadata = json!({ "itemId": item.get("id").cloned().unwrap_or(Value::Null) }); + if let Some(phase) = item.get("phase").filter(|v| !v.is_null()) { + metadata["phase"] = phase.clone(); + } + if let Some(annotations) = annotations.filter(|v| !v.is_empty()) { + metadata["annotations"] = json!(annotations); + } + metadata +} + fn annotation_source(annotation: &Value, provider_key: &str) -> Result { let invalid = |field: &str| { AiMuxError::InvalidResponseData(format!("Invalid Responses annotation field: {field}")) @@ -470,7 +499,7 @@ fn annotation_source(annotation: &Value, provider_key: &str) -> Result( request_body: Value, response_headers: HashMap, ) -> Result +where + S: Stream> + Unpin + Send + 'static, +{ + build_responses_event_stream_with_raw( + first_event, + sse_stream, + provider_key, + warnings, + store_flag, + request_url, + request_body, + response_headers, + false, + ) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn build_responses_event_stream_with_raw( + first_event: Option>, + sse_stream: S, + provider_key: String, + warnings: Vec, + store_flag: bool, + request_url: String, + request_body: Value, + response_headers: HashMap, + include_raw_chunks: bool, +) -> Result +where + S: Stream> + Unpin + Send + 'static, +{ + build_responses_event_stream_with_tools( + first_event, + sse_stream, + provider_key, + warnings, + store_flag, + request_url, + request_body, + response_headers, + include_raw_chunks, + HashMap::new(), + HashMap::new(), + ) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn build_responses_event_stream_with_tools( + first_event: Option>, + sse_stream: S, + provider_key: String, + warnings: Vec, + store_flag: bool, + request_url: String, + request_body: Value, + response_headers: HashMap, + include_raw_chunks: bool, + tool_names: HashMap, + prompt_approval_tool_call_ids: HashMap, +) -> Result where S: Stream> + Unpin + Send + 'static, { @@ -519,42 +608,9 @@ where if let Some(Ok(ref val)) = first_event { let etype = val.get("type").and_then(|v| v.as_str()).unwrap_or(""); if etype == "error" || etype == "response.failed" { - let error_value = val - .get("response") - .and_then(|response| response.get("error")) - .or_else(|| val.get("error")); - let message = val - .get("response") - .and_then(|r| r.get("error")) - .and_then(|e| e.get("message")) - .and_then(|v| v.as_str()) - .or_else(|| { - val.get("error") - .and_then(|e| e.get("message")) - .and_then(|v| v.as_str()) - }) - .unwrap_or("Responses API stream error"); - let status_code = val - .get("status") - .or_else(|| val.get("status_code")) - .or_else(|| error_value.and_then(|error| error.get("status"))) - .or_else(|| error_value.and_then(|error| error.get("code"))) - .and_then(Value::as_u64) - .and_then(|status| u16::try_from(status).ok()) - .filter(|status| (400..=599).contains(status)); - let provider_code = val - .get("response") - .and_then(|r| r.get("error")) - .or_else(|| val.get("error")) - .and_then(|e| e.get("type").or_else(|| e.get("code"))) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - return Err(aimux_provider_utils::stream_error_api_call( - message, - provider_code, - status_code, + return Err(super::super::openai_stream_error( val, - request_url.clone(), + &request_url, request_body.clone(), response_headers.clone(), )); @@ -578,6 +634,17 @@ where let mut final_finish_reason: Option = None; let mut response_id: Option = None; let mut stream_errored = false; + let mut text_states: HashMap = HashMap::new(); + let mut output_ids: HashMap = HashMap::new(); + let mut text_annotations: HashMap> = HashMap::new(); + let mut logprobs = Vec::new(); + let shell_provider_executed = hosted_shell(&request_body); + let mut provider_items: HashMap = HashMap::new(); + let mut emitted_calls = HashSet::new(); + let mut ended_inputs = HashSet::new(); + let mut diff_received = HashSet::new(); + let mut hosted_tool_search_ids = VecDeque::new(); + let mut approval_tool_call_ids: HashMap = HashMap::new(); let mut event_iter = futures::stream::iter(first_event.into_iter()).chain(sse_stream); @@ -585,6 +652,7 @@ where while let Some(event) = event_iter.next().await { match event { Ok(parsed) => { + if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); } let etype = parsed .get("type") .and_then(|v| v.as_str()) @@ -638,7 +706,9 @@ where .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); - yield Ok(StreamPart::TextStart { id, provider_metadata: None}); + output_ids.insert(output_index, id.clone()); + text_states.insert(id.clone(), item.clone()); + yield Ok(StreamPart::TextStart { id, provider_metadata: Some(provider_namespace(&provider_key, text_metadata(item, None)).expect("provider metadata must be an object"))}); } "function_call" => { let call_id = item @@ -726,18 +796,64 @@ where provider_metadata: Some(meta), }); } - _ => {} + _ => { + provider_items.insert(output_index, item.clone()); + for part in provider_tool_start(item, &tool_names) { + if let StreamPart::ToolCall(call) = &part { emitted_calls.insert(call.tool_call_id.clone()); } + if let StreamPart::ToolInputEnd { id, .. } = &part { ended_inputs.insert(id.clone()); } + yield Ok(part); + } + } + } + } + } + + "response.image_generation_call.partial_image" => { + yield Ok(StreamPart::ToolResult(aimux_core::tool::ToolResult { + tool_call_id: parsed["item_id"].as_str().unwrap_or_default().to_owned(), + tool_name: tool_names.get("image_generation").map_or("image_generation", String::as_str).to_owned(), + result: json!({"result": parsed["partial_image_b64"]}), + preliminary: Some(true), is_error: None, dynamic: None, provider_metadata: None, + })); + } + "response.code_interpreter_call_code.delta" | "response.apply_patch_call_operation_diff.delta" => { + let index = parsed["output_index"].as_u64().unwrap_or(0) as usize; + if let Some(item) = provider_items.get(&index) { + let id = item.get("call_id").filter(|v| !v.is_null()).unwrap_or(&item["id"]).as_str().unwrap_or_default().to_owned(); + diff_received.insert(id.clone()); + yield Ok(StreamPart::ToolInputDelta { id, delta: escaped_delta(parsed["delta"].as_str().unwrap_or_default()), provider_metadata: None }); + } + } + "response.code_interpreter_call_code.done" | "response.apply_patch_call_operation_diff.done" => { + let index = parsed["output_index"].as_u64().unwrap_or(0) as usize; + if let Some(item) = provider_items.get(&index) { + let id = item.get("call_id").filter(|v| !v.is_null()).unwrap_or(&item["id"]).as_str().unwrap_or_default().to_owned(); + if ended_inputs.insert(id.clone()) { + let code = etype == "response.code_interpreter_call_code.done"; + if !code && !diff_received.contains(&id) { + yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: escaped_delta(parsed["diff"].as_str().unwrap_or_default()), provider_metadata: None }); + } + yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: if code { "\"}" } else { "\"}}" }.to_owned(), provider_metadata: None }); + yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: None }); + if code { + let mut item = item.clone(); item["code"] = parsed["code"].clone(); + for part in super::provider_events::provider_tool_content(&item, &provider_key, &tool_names, shell_provider_executed, None) { + if let GenerateContent::ToolCall(call) = part { emitted_calls.insert(id.clone()); yield Ok(StreamPart::ToolCall(call)); } + } + } } } } // ── response.output_text.delta → TextDelta ──────────────── "response.output_text.delta" => { - let id = parsed + if request_body.get("top_logprobs").and_then(Value::as_u64).is_some_and(|n| n > 0) && let Some(value) = parsed.get("logprobs") { logprobs.push(value.clone()); } + let mut id = parsed .get("item_id") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); + if let Some(index) = parsed.get("output_index").and_then(Value::as_u64) && let Some(original_id) = output_ids.get(&(index as usize)) { id = original_id.clone(); } let delta = parsed .get("delta") .and_then(|v| v.as_str()) @@ -878,37 +994,20 @@ where .and_then(|v| v.as_str()) .unwrap_or(""); match item_type { - "mcp_approval_request" => { - let tool_call_id = generate_source_id(); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: tool_call_id.clone(), - tool_name: format!("mcp.{}", item.get("name").and_then(Value::as_str).unwrap_or("")), - input: item.get("arguments").and_then(Value::as_str).unwrap_or("").to_string(), - provider_executed: Some(true), dynamic: Some(true), - provider_metadata: None, - })); - yield Ok(StreamPart::ToolApprovalRequest(aimux_core::result::RawToolApprovalRequest { - approval_id: item.get("approval_request_id").or_else(|| item.get("id")).and_then(Value::as_str).unwrap_or("").to_string(), - tool_call_id, provider_metadata: None, - })); - } - "compaction" => { - yield Ok(StreamPart::Custom { - kind: "openai.compaction".to_string(), - provider_metadata: Some(provider_namespace(&provider_key, json!({ - "type": "compaction", "itemId": item.get("id"), - "encryptedContent": item.get("encrypted_content"), - })).expect("provider metadata must be an object")), - }); - } - "message" => { - let id = item + let mut id = item .get("id") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); + if let Some(original_id) = output_ids.remove(&output_index) { id = original_id; } + let mut metadata_item = item.clone(); + metadata_item["id"] = json!(id); + if metadata_item.get("phase").is_none() && let Some(phase) = text_states.get(&id).and_then(|v| v.get("phase")) { + metadata_item["phase"] = phase.clone(); + } + let annotations = text_annotations.remove(&id); + yield Ok(StreamPart::TextEnd { id, provider_metadata: Some(provider_namespace(&provider_key, text_metadata(&metadata_item, annotations.as_ref())).expect("provider metadata must be an object"))}); } "function_call" => { has_function_call = true; @@ -939,7 +1038,7 @@ where input, provider_executed: None, dynamic: None, - provider_metadata: None, + provider_metadata: Some(provider_namespace(&provider_key, tool_metadata(item)).expect("provider metadata must be an object")), })); } "custom_tool_call" => { @@ -973,7 +1072,7 @@ where input, provider_executed: None, dynamic: None, - provider_metadata: None, + provider_metadata: Some(provider_namespace(&provider_key, tool_metadata(item)).expect("provider metadata must be an object")), })); } "reasoning" => { @@ -983,6 +1082,7 @@ where .unwrap_or("") .to_string(); if let Some(state) = active_reasoning.get_mut(&id) { + state.encrypted_content = item.get("encrypted_content").and_then(Value::as_str).map(str::to_owned); let meta = reasoning_stream_metadata( &provider_key, &id, @@ -1010,18 +1110,72 @@ where } active_reasoning.remove(&id); } - _ => {} + _ => { + let mut item = item.clone(); + if item_type == "local_shell_call" && let Some(action) = item["action"].as_object_mut() { + for (from, to) in [("timeout_ms", "timeoutMs"), ("working_directory", "workingDirectory")] { + if let Some(value) = action.remove(from) { action.insert(to.into(), value); } + } + } + if item_type == "apply_patch_call" { + let id = item["call_id"].as_str().unwrap_or_default().to_owned(); + if ended_inputs.insert(id.clone()) && item["operation"]["type"] != "delete_file" { + if !diff_received.contains(&id) { + yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: escaped_delta(item["operation"]["diff"].as_str().unwrap_or_default()), provider_metadata: None }); + } + yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: "\"}}".into(), provider_metadata: None }); + yield Ok(StreamPart::ToolInputEnd { id, provider_metadata: None }); + } + } + if item_type == "tool_search_call" && item["execution"] == "server" { + item["call_id"] = Value::Null; + } + pair_tool_search_id(&mut item, &mut hosted_tool_search_ids); + let mcp_tool_call_id = item.get("approval_request_id").and_then(Value::as_str) + .and_then(|id| approval_tool_call_ids.get(id).or_else(|| prompt_approval_tool_call_ids.get(id))).cloned(); + for part in super::provider_events::provider_tool_content(&item, &provider_key, &tool_names, shell_provider_executed, mcp_tool_call_id.as_deref()) { + match part { + GenerateContent::ToolCall(call) => { + if item_type == "code_interpreter_call" || (item_type == "apply_patch_call" && item["status"] != "completed") { continue; } + if emitted_calls.insert(call.tool_call_id.clone()) || item_type == "mcp_call" { + has_function_call |= call.provider_executed != Some(true) && matches!(item_type, "computer_call" | "apply_patch_call"); + if matches!(item_type, "computer_call" | "apply_patch_call" | "tool_search_call") { + if item_type == "tool_search_call" && item["execution"] != "server" { + yield Ok(StreamPart::ToolInputStart { id: call.tool_call_id.clone(), tool_name: call.tool_name.clone(), provider_executed: call.provider_executed, dynamic: None, title: None, provider_metadata: None }); + } + if !ended_inputs.contains(&call.tool_call_id) { + if item_type == "computer_call" && !call.input.is_empty() { yield Ok(StreamPart::ToolInputDelta { id: call.tool_call_id.clone(), delta: call.input.clone(), provider_metadata: None }); } + yield Ok(StreamPart::ToolInputEnd { id: call.tool_call_id.clone(), provider_metadata: None }); + } + } + yield Ok(StreamPart::ToolCall(call)); + } + } + GenerateContent::ToolResult(result) => { yield Ok(StreamPart::ToolResult(result)); } + GenerateContent::ToolApprovalRequest(approval) => { + approval_tool_call_ids.insert(approval.approval_id.clone(), approval.tool_call_id.clone()); + yield Ok(StreamPart::ToolApprovalRequest(approval)); + } + GenerateContent::Custom { kind, provider_metadata } => { + yield Ok(StreamPart::Custom { kind, provider_metadata }); + } + _ => {} + } + } + provider_items.remove(&output_index); + } } } } // ── response.output_text.annotation.added → Source ───────── "response.output_text.annotation.added" => { - match annotation_source( - parsed.get("annotation").unwrap_or(&Value::Null), - &provider_key, - ) { - Ok(source) => yield Ok(StreamPart::Source(source)), + match annotation_source(parsed.get("annotation").unwrap_or(&Value::Null), &provider_key) { + Ok(source) => { + let id = parsed.get("item_id").and_then(Value::as_str).unwrap_or_default().to_string(); + text_annotations.entry(id).or_default().push(parsed["annotation"].clone()); + yield Ok(StreamPart::Source(source)); + } Err(error) => { stream_errored = true; yield Ok(StreamPart::Error { error }); @@ -1036,10 +1190,9 @@ where .get("incomplete_details") .and_then(|d| d.get("reason")) .and_then(|v| v.as_str()); - final_finish_reason = Some(map_responses_finish_reason( - reason, - has_function_call, - )); + if !stream_errored { + final_finish_reason = Some(map_responses_finish_reason(reason, has_function_call)); + } final_usage = resp_obj.get("usage").and_then(parse_usage); // Keep the raw wire usage object for `usage.raw` @@ -1091,32 +1244,8 @@ where && resp_obj.get("error").is_some() { stream_errored = true; - let message = resp_obj - .get("error") - .and_then(|e| e.get("message")) - .and_then(|v| v.as_str()) - .unwrap_or("Responses API stream failed"); - // In-band failure inside the SSE stream: the shared - // helper redacts the raw request context. - yield Ok(StreamPart::Error { - error: aimux_provider_utils::stream_error_api_call( - message, - resp_obj - .get("error") - .and_then(|e| e.get("type").or_else(|| e.get("code"))) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - Some(200), - &parsed, - request_url.clone(), - request_body.clone(), - response_headers.clone(), - ), - }); - // A terminal error ends this stream; waiting - // for more events can hang on a source that - // keeps the connection open. - break; + yield Ok(StreamPart::Error { error: super::super::openai_stream_error(&parsed, &request_url, request_body.clone(), response_headers.clone()) }); + } } } @@ -1128,27 +1257,7 @@ where unified: FinishReasonUnified::Error, raw: Some("error".to_string()), }); - let message = parsed - .get("error") - .and_then(|e| e.get("message")) - .and_then(|v| v.as_str()) - .unwrap_or("Responses API stream error"); - yield Ok(StreamPart::Error { - error: aimux_provider_utils::stream_error_api_call( - message, - parsed - .get("error") - .and_then(|e| e.get("type").or_else(|| e.get("code"))) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - Some(200), - &parsed, - request_url.clone(), - request_body.clone(), - response_headers.clone(), - ), - }); - break; + yield Ok(StreamPart::Error { error: super::super::openai_stream_error(&parsed, &request_url, request_body.clone(), response_headers.clone()) }); } _ => { @@ -1160,11 +1269,9 @@ where } } Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; - } + stream_errored = true; + final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); + yield Ok(StreamPart::Error { error }); } } } @@ -1180,31 +1287,147 @@ where if let Some(ctx) = final_reasoning_context { pm["reasoningContext"] = ctx; } + if !logprobs.is_empty() { pm["logprobs"] = json!(logprobs); } let provider_metadata = Some(provider_namespace(&provider_key, pm).expect("provider metadata must be an object")); yield Ok(StreamPart::Finish { - finish_reason: if stream_errored { - FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - } - } else { - final_finish_reason.unwrap_or(FinishReason { - unified: if has_function_call { - FinishReasonUnified::ToolCalls - } else { - FinishReasonUnified::Stop - }, - raw: None, - }) - }, - usage: if stream_errored { - Usage::default() - } else { - convert_responses_usage(final_usage.as_ref(), final_raw_usage) - }, + finish_reason: final_finish_reason.unwrap_or(FinishReason { + unified: FinishReasonUnified::Other, raw: None, + }), + usage: convert_responses_usage(final_usage.as_ref(), final_raw_usage), provider_metadata, }); }; Ok(Box::pin(stream)) } + +fn pair_tool_search_id(item: &mut Value, hosted_ids: &mut VecDeque) { + match item["type"].as_str() { + Some("tool_search_call") if item["execution"] == "server" => { + hosted_ids.push_back( + item.get("call_id") + .filter(|v| !v.is_null()) + .unwrap_or(&item["id"]) + .clone(), + ); + } + Some("tool_search_output") if item.get("call_id").is_none_or(Value::is_null) => { + if let Some(id) = hosted_ids.pop_front() { + item["call_id"] = id; + } + } + _ => {} + } +} + +fn escaped_delta(value: &str) -> String { + let value = json!(value).to_string(); + value[1..value.len() - 1].to_owned() +} + +fn provider_tool_start(item: &Value, names: &HashMap) -> Vec { + let kind = item["type"].as_str().unwrap_or_default(); + let canonical = match kind { + "web_search_call" => "web_search", + "file_search_call" => "file_search", + "image_generation_call" => "image_generation", + "code_interpreter_call" => "code_interpreter", + "computer_call" => "computer", + "apply_patch_call" => "apply_patch", + "tool_search_call" if item["execution"] == "server" => "tool_search", + _ => return Vec::new(), + }; + let name = names + .get(canonical) + .map_or(canonical, String::as_str) + .to_owned(); + let id = item + .get("call_id") + .filter(|v| !v.is_null()) + .unwrap_or(&item["id"]) + .as_str() + .unwrap_or_default() + .to_owned(); + let hosted = matches!( + kind, + "web_search_call" + | "file_search_call" + | "image_generation_call" + | "code_interpreter_call" + | "tool_search_call" + ) + .then_some(true); + let call = || { + StreamPart::ToolCall(RawToolCall { + tool_call_id: id.clone(), + tool_name: name.clone(), + input: "{}".to_owned(), + provider_executed: hosted, + dynamic: None, + provider_metadata: None, + }) + }; + if matches!(kind, "file_search_call" | "image_generation_call") { + return vec![call()]; + } + let mut parts = vec![StreamPart::ToolInputStart { + id: id.clone(), + tool_name: name.clone(), + provider_executed: hosted, + dynamic: None, + title: None, + provider_metadata: None, + }]; + match kind { + "web_search_call" => { + parts.push(StreamPart::ToolInputEnd { + id: id.clone(), + provider_metadata: None, + }); + parts.push(call()); + } + "code_interpreter_call" => parts.push(StreamPart::ToolInputDelta { + id, + delta: format!("{{\"containerId\":{},\"code\":\"", item["container_id"]), + provider_metadata: None, + }), + "apply_patch_call" => { + let operation = &item["operation"]; + if operation["type"] == "delete_file" { + parts.push(StreamPart::ToolInputDelta { + id: id.clone(), + delta: json!({"callId":id,"operation":operation}).to_string(), + provider_metadata: None, + }); + parts.push(StreamPart::ToolInputEnd { + id, + provider_metadata: None, + }); + } else { + parts.push(StreamPart::ToolInputDelta { + id: id.clone(), + delta: format!( + "{{\"callId\":{},\"operation\":{{\"type\":{},\"path\":{},\"diff\":\"", + json!(id), + operation["type"], + operation["path"] + ), + provider_metadata: None, + }); + } + } + _ => {} + } + parts +} + +fn hosted_shell(body: &Value) -> bool { + body["tools"].as_array().is_some_and(|tools| { + tools.iter().any(|tool| { + tool["type"] == "shell" + && tool.get("environment").is_some_and(|environment| { + !matches!(environment["type"].as_str(), Some("local")) + }) + }) + }) +} diff --git a/aimux-providers/src/openai/responses/tool_args.rs b/aimux-providers/src/openai/responses/tool_args.rs new file mode 100644 index 000000000..337f7959d --- /dev/null +++ b/aimux-providers/src/openai/responses/tool_args.rs @@ -0,0 +1,437 @@ +//! Validation of the pinned provider tool argument schemas. + +use aimux_core::AiMuxError; +use serde_json::Value; + +fn strings(value: &Value) -> bool { + value + .as_array() + .is_some_and(|values| values.iter().all(Value::is_string)) +} + +fn optional(value: &Value, key: &str, valid: impl Fn(&Value) -> bool) -> bool { + value.get(key).is_none_or(valid) +} + +fn choices(value: &Value, options: &[&str]) -> bool { + value.as_str().is_some_and(|value| options.contains(&value)) +} + +fn string_fields(value: &Value, fields: &[&str]) -> bool { + fields + .iter() + .all(|key| optional(value, key, Value::is_string)) +} + +fn string_record(value: &Value) -> bool { + value + .as_object() + .is_some_and(|value| value.values().all(Value::is_string)) +} + +fn filter(value: &Value) -> bool { + if !value.is_object() { + return false; + } + if choices(&value["type"], &["and", "or"]) { + value["filters"] + .as_array() + .is_some_and(|items| items.iter().all(filter)) + } else { + value["key"].is_string() + && choices( + &value["type"], + &["eq", "ne", "gt", "gte", "lt", "lte", "in", "nin"], + ) + && (value["value"].is_string() + || value["value"].is_number() + || value["value"].is_boolean() + || strings(&value["value"])) + } +} + +fn location(value: &Value) -> bool { + value.is_object() + && value["type"] == "approximate" + && string_fields(value, &["country", "city", "region", "timezone"]) +} + +fn skills(value: &Value) -> bool { + value.as_array().is_some_and(|items| { + items.iter().all(|skill| { + skill.is_object() + && match skill["type"].as_str() { + Some("skillReference") => { + string_record(&skill["providerReference"]) + && optional(skill, "version", Value::is_string) + } + Some("inline") => { + skill["name"].is_string() + && skill["description"].is_string() + && skill["source"].is_object() + && skill["source"]["type"] == "base64" + && skill["source"]["mediaType"] == "application/zip" + && skill["source"]["data"].is_string() + } + _ => false, + } + }) + }) +} + +fn environment(value: &Value) -> bool { + if !value.is_object() { + return false; + } + match value["type"].as_str() { + Some("containerAuto") => { + optional(value, "fileIds", strings) + && optional(value, "memoryLimit", |v| { + choices(v, &["1g", "4g", "16g", "64g"]) + }) + && optional(value, "skills", skills) + && optional(value, "networkPolicy", |policy| { + policy.is_object() + && match policy["type"].as_str() { + Some("disabled") => true, + Some("allowlist") => { + strings(&policy["allowedDomains"]) + && optional(policy, "domainSecrets", |secrets| { + secrets.as_array().is_some_and(|secrets| { + secrets.iter().all(|secret| { + secret.is_object() + && secret["domain"].is_string() + && secret["name"].is_string() + && secret["value"].is_string() + }) + }) + }) + } + _ => false, + } + }) + } + Some("containerReference") => value["containerId"].is_string(), + Some("local") | None + if value + .get("type") + .is_none_or(|kind| matches!(kind.as_str(), Some("local"))) => + { + optional(value, "skills", |skills| { + skills.as_array().is_some_and(|skills| { + skills.iter().all(|skill| { + skill.is_object() + && skill["name"].is_string() + && skill["description"].is_string() + && skill["path"].is_string() + }) + }) + }) + } + _ => false, + } +} + +pub(crate) fn validate_tool_args(kind: &str, args: &Value) -> Result<(), AiMuxError> { + // These tools have no argument schema in the upstream preparation path. + if matches!( + kind, + "local_shell" | "apply_patch" | "computer" | "computer_use" | "programmatic_tool_calling" + ) { + return Ok(()); + } + let valid = + args.is_object() + && match kind { + "web_search" | "web_search_preview" => { + optional(args, "searchContextSize", |v| { + choices(v, &["low", "medium", "high"]) + }) && optional(args, "userLocation", location) + && (kind != "web_search" + || (optional(args, "externalWebAccess", Value::is_boolean) + && optional(args, "filters", |v| { + v.is_object() + && optional(v, "allowedDomains", strings) + && optional(v, "blockedDomains", strings) + }))) + } + "file_search" => { + strings(&args["vectorStoreIds"]) + && optional(args, "maxNumResults", Value::is_number) + && optional(args, "filters", filter) + && optional(args, "ranking", |ranking| { + ranking.is_object() + && optional(ranking, "ranker", Value::is_string) + && optional(ranking, "scoreThreshold", Value::is_number) + }) + } + "code_interpreter" => optional(args, "container", |container| { + container.is_string() + || (container.is_object() && optional(container, "fileIds", strings)) + }), + "image_generation" => { + let allowed = [ + "action", + "background", + "inputFidelity", + "inputImageMask", + "model", + "moderation", + "outputCompression", + "outputFormat", + "partialImages", + "quality", + "size", + ]; + args.as_object().is_some_and(|object| { + object.keys().all(|key| allowed.contains(&key.as_str())) + }) && optional(args, "action", |v| { + choices(v, &["generate", "edit", "auto"]) + }) && optional(args, "background", |v| { + choices(v, &["auto", "opaque", "transparent"]) + }) && optional(args, "inputFidelity", |v| choices(v, &["low", "high"])) + && optional(args, "inputImageMask", |mask| { + mask.is_object() && string_fields(mask, &["fileId", "imageUrl"]) + }) + && optional(args, "model", Value::is_string) + && optional(args, "moderation", |v| choices(v, &["auto", "low"])) + && optional(args, "outputFormat", |v| { + choices(v, &["png", "jpeg", "webp"]) + }) + && optional(args, "quality", |v| { + choices(v, &["auto", "low", "medium", "high", "xhigh", "max"]) + }) + && optional(args, "outputCompression", |v| { + v.as_f64() + .is_some_and(|n| n.fract() == 0.0 && (0.0..=100.0).contains(&n)) + }) + && optional(args, "partialImages", |v| { + v.as_f64() + .is_some_and(|n| n.fract() == 0.0 && (0.0..=3.0).contains(&n)) + }) + && optional(args, "size", |v| { + v.as_str().is_some_and(|size| { + size == "auto" + || size.split_once('x').is_some_and(|(w, h)| { + !w.is_empty() + && !h.is_empty() + && w.bytes().all(|b| b.is_ascii_digit()) + && h.bytes().all(|b| b.is_ascii_digit()) + }) + }) + }) + } + "mcp" => { + args["serverLabel"].is_string() + && string_fields( + args, + &[ + "authorization", + "connectorId", + "serverDescription", + "serverUrl", + ], + ) + && (args.get("serverUrl").is_some() || args.get("connectorId").is_some()) + && optional(args, "headers", string_record) + && optional(args, "allowedTools", |tools| { + strings(tools) + || (tools.is_object() + && optional(tools, "readOnly", Value::is_boolean) + && optional(tools, "toolNames", strings)) + }) + && optional(args, "requireApproval", |approval| { + choices(approval, &["always", "never"]) + || (approval.is_object() + && optional(approval, "never", |never| { + never.is_object() && optional(never, "toolNames", strings) + })) + }) + } + "shell" => optional(args, "environment", environment), + "tool_search" => { + optional(args, "execution", |v| choices(v, &["server", "client"])) + && optional(args, "description", Value::is_string) + && optional(args, "parameters", Value::is_object) + } + "custom" => { + optional(args, "description", Value::is_string) + && optional(args, "async", Value::is_boolean) + && optional(args, "format", |format| { + format.is_object() + && match format["type"].as_str() { + Some("text") => true, + Some("grammar") => { + choices(&format["syntax"], &["regex", "lark"]) + && format["definition"].is_string() + } + _ => false, + } + }) + } + _ => true, + }; + if valid { + Ok(()) + } else { + Err(AiMuxError::InvalidArgument(format!( + "Invalid arguments for provider tool openai.{kind}" + ))) + } +} + +fn safety_checks(value: &Value) -> bool { + value.as_array().is_some_and(|checks| { + checks.iter().all(|check| { + check.is_object() + && check["id"].is_string() + && string_fields(check, &["code", "message"]) + }) + }) +} + +fn computer_action(value: &Value) -> bool { + if !value.is_object() { + return false; + } + match value["type"].as_str() { + Some("click" | "double_click" | "move" | "scroll") => { + value["x"].is_number() + && value["y"].is_number() + && optional(value, "keys", strings) + && (value["type"] != "click" + || choices( + &value["button"], + &["left", "right", "wheel", "back", "forward"], + )) + && (value["type"] != "scroll" + || (value["scrollX"].is_number() && value["scrollY"].is_number())) + } + Some("drag") => { + optional(value, "keys", strings) + && value["path"].as_array().is_some_and(|points| { + points.iter().all(|point| { + point.is_object() && point["x"].is_number() && point["y"].is_number() + }) + }) + } + Some("keypress") => strings(&value["keys"]), + Some("type") => value["text"].is_string(), + Some("screenshot" | "wait") => true, + _ => false, + } +} + +pub(super) fn validate_tool_call(kind: &str, input: &Value) -> Result<(), AiMuxError> { + let valid = if kind == "custom" { + input.is_string() + } else { + input.is_object() + && match kind { + "local_shell" => { + let action = &input["action"]; + action.is_object() + && action["type"] == "exec" + && strings(&action["command"]) + && optional(action, "timeoutMs", Value::is_number) + && string_fields(action, &["user", "workingDirectory"]) + && optional(action, "env", string_record) + } + "shell" => { + let action = &input["action"]; + action.is_object() + && strings(&action["commands"]) + && optional(action, "timeoutMs", Value::is_number) + && optional(action, "maxOutputLength", Value::is_number) + } + "apply_patch" => { + let operation = &input["operation"]; + input["callId"].is_string() + && operation.is_object() + && operation["path"].is_string() + && match operation["type"].as_str() { + Some("create_file" | "update_file") => operation["diff"].is_string(), + Some("delete_file") => true, + _ => false, + } + } + "computer" => { + input["actions"] + .as_array() + .is_some_and(|actions| actions.iter().all(computer_action)) + && safety_checks(&input["pendingSafetyChecks"]) + && choices( + &input["status"], + &["in_progress", "completed", "incomplete"], + ) + } + "programmatic_tool_calling" => { + input["code"].is_string() && input["fingerprint"].is_string() + } + "tool_search" => optional(input, "call_id", |id| id.is_null() || id.is_string()), + _ => true, + } + }; + if valid { + Ok(()) + } else { + Err(AiMuxError::InvalidArgument(format!( + "Invalid input for provider tool openai.{kind}" + ))) + } +} + +pub(super) fn validate_tool_result(kind: &str, output: &Value) -> Result<(), AiMuxError> { + // Custom tools deliberately do not define an output schema upstream. + if kind == "custom" { + return Ok(()); + } + let valid = output.is_object() + && match kind { + "local_shell" => output["output"].is_string(), + "shell" => output["output"].as_array().is_some_and(|items| { + items.iter().all(|item| { + let outcome = &item["outcome"]; + item.is_object() + && item["stdout"].is_string() + && item["stderr"].is_string() + && outcome.is_object() + && match outcome["type"].as_str() { + Some("timeout") => true, + Some("exit") => outcome["exitCode"].is_number(), + _ => false, + } + }) + }), + "apply_patch" => { + choices(&output["status"], &["completed", "failed"]) + && optional(output, "output", Value::is_string) + } + "computer" => { + let screenshot = &output["output"]; + screenshot.is_object() + && screenshot["type"] == "computer_screenshot" + && (screenshot["imageUrl"].is_string() || screenshot["fileId"].is_string()) + && string_fields(screenshot, &["imageUrl", "fileId"]) + && optional(screenshot, "detail", |detail| { + choices(detail, &["auto", "low", "high", "original"]) + }) + && optional(output, "acknowledgedSafetyChecks", safety_checks) + } + "programmatic_tool_calling" => { + output["result"].is_string() + && choices(&output["status"], &["completed", "incomplete"]) + } + "tool_search" => output["tools"] + .as_array() + .is_some_and(|tools| tools.iter().all(Value::is_object)), + _ => true, + }; + if valid { + Ok(()) + } else { + Err(AiMuxError::InvalidArgument(format!( + "Invalid output for provider tool openai.{kind}" + ))) + } +} diff --git a/aimux-providers/src/openai/responses/types.rs b/aimux-providers/src/openai/responses/types.rs index 59a78070d..3e2ffc012 100644 --- a/aimux-providers/src/openai/responses/types.rs +++ b/aimux-providers/src/openai/responses/types.rs @@ -12,8 +12,6 @@ //! model code. Only [`ResponsesUsage`] — which has a stable, well-defined //! shape — gets a dedicated typed struct. -#![allow(dead_code)] - use serde::{Deserialize, Serialize}; // ── Usage ──────────────────────────────────────────────────────────────────── diff --git a/aimux-providers/src/openai/speech.rs b/aimux-providers/src/openai/speech.rs index 144c2b175..8f8ffc26c 100644 --- a/aimux-providers/src/openai/speech.rs +++ b/aimux-providers/src/openai/speech.rs @@ -123,7 +123,7 @@ fn build_request_body_and_warnings( if SUPPORTED_OUTPUT_FORMATS.contains(&output_format) { body.insert("response_format".to_string(), json!(output_format)); - } else { + } else if !output_format.is_empty() { warnings.push(Warning::Unsupported { feature: "outputFormat".to_string(), details: Some(format!( @@ -132,7 +132,37 @@ fn build_request_body_and_warnings( }); } - if let Some(ref language) = options.language { + if let Some(provider_options) = options + .provider_options + .as_ref() + .and_then(|options| options.get("openai")) + { + if let Some(value) = provider_options + .get("speed") + .filter(|value| !value.is_null()) + { + let speed = value + .as_f64() + .filter(|speed| (0.25..=4.0).contains(speed)) + .ok_or_else(|| { + AiMuxError::InvalidArgument("invalid openai speech option: speed".into()) + })?; + body.insert("speed".into(), json!(speed)); + } + if let Some(value) = provider_options + .get("instructions") + .filter(|value| !value.is_null()) + { + let instructions = value.as_str().ok_or_else(|| { + AiMuxError::InvalidArgument("invalid openai speech option: instructions".into()) + })?; + body.insert("instructions".into(), json!(instructions)); + } + } + + if let Some(ref language) = options.language + && !language.is_empty() + { warnings.push(Warning::Unsupported { feature: "language".to_string(), details: Some(format!( diff --git a/aimux-providers/src/openai/tools.rs b/aimux-providers/src/openai/tools.rs new file mode 100644 index 000000000..3e6713899 --- /dev/null +++ b/aimux-providers/src/openai/tools.rs @@ -0,0 +1,713 @@ +//! Typed factories for the native Responses provider tools. + +use std::collections::BTreeMap; + +use aimux_core::{AiMuxError, tool::ProviderTool}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SearchContextSize { + Low, + Medium, + High, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ImageAction { + Generate, + Edit, + Auto, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ImageBackground { + Auto, + Opaque, + Transparent, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum InputFidelity { + Low, + High, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ImageModeration { + Auto, + Low, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ImageOutputFormat { + Png, + Jpeg, + Webp, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ImageQuality { + Auto, + Low, + Medium, + High, + Xhigh, + Max, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ToolSearchExecution { + Server, + Client, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum GrammarSyntax { + Regex, + Lark, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ComparisonOperator { + Eq, + Ne, + Gt, + Gte, + Lt, + Lte, + In, + Nin, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CompoundOperator { + And, + Or, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub enum UserLocationType { + #[default] + #[serde(rename = "approximate")] + Approximate, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UserLocation { + pub r#type: UserLocationType, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub country: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub city: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub region: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub timezone: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WebSearchFilters { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub allowed_domains: Option>, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub blocked_domains: Option>, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WebSearchArgs { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub external_web_access: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub filters: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub search_context_size: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub user_location: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WebSearchPreviewArgs { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub search_context_size: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub user_location: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum FileSearchFilterValue { + String(String), + Number(f64), + Boolean(bool), + Strings(Vec), +} +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum FileSearchFilter { + Comparison { + key: String, + r#type: ComparisonOperator, + value: FileSearchFilterValue, + }, + Compound { + r#type: CompoundOperator, + filters: Vec, + }, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FileSearchRanking { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ranker: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub score_threshold: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FileSearchArgs { + pub vector_store_ids: Vec, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_num_results: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ranking: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub filters: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum CodeInterpreterContainer { + Id(String), + Files { + #[serde(rename = "fileIds", default, skip_serializing_if = "Option::is_none")] + file_ids: Option>, + }, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CodeInterpreterArgs { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub container: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct InputImageMask { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub file_id: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub image_url: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ImageGenerationArgs { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub action: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub background: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_fidelity: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_image_mask: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub moderation: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output_compression: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output_format: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub partial_images: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub quality: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub size: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum CustomFormat { + Text, + Grammar { + syntax: GrammarSyntax, + definition: String, + }, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CustomToolArgs { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub description: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub r#async: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub format: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum McpAllowedTools { + Names(Vec), + Filter { + #[serde(rename = "readOnly", default, skip_serializing_if = "Option::is_none")] + read_only: Option, + #[serde(rename = "toolNames", default, skip_serializing_if = "Option::is_none")] + tool_names: Option>, + }, +} +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum McpApprovalPolicy { + Always, + Never, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct McpApprovalFilter { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_names: Option>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum McpRequireApproval { + Policy(McpApprovalPolicy), + Filter { + #[serde(default, skip_serializing_if = "Option::is_none")] + never: Option, + }, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct McpArgs { + pub server_label: String, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub allowed_tools: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub authorization: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub connector_id: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub headers: Option>, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub require_approval: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub server_description: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub server_url: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum ShellMemoryLimit { + #[serde(rename = "1g")] + OneG, + #[serde(rename = "4g")] + FourG, + #[serde(rename = "16g")] + SixteenG, + #[serde(rename = "64g")] + SixtyFourG, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ShellDomainSecret { + pub domain: String, + + pub name: String, + + pub value: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum ShellNetworkPolicy { + Disabled, + Allowlist { + #[serde(rename = "allowedDomains")] + allowed_domains: Vec, + #[serde( + rename = "domainSecrets", + default, + skip_serializing_if = "Option::is_none" + )] + domain_secrets: Option>, + }, +} +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum ShellSkill { + SkillReference { + #[serde(rename = "providerReference")] + provider_reference: BTreeMap, + #[serde(default, skip_serializing_if = "Option::is_none")] + version: Option, + }, + Inline { + name: String, + description: String, + source: ShellSkillSource, + }, +} +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum ShellSkillSource { + Base64 { + #[serde(rename = "mediaType")] + media_type: ShellSkillMediaType, + data: String, + }, +} +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum ShellSkillMediaType { + #[serde(rename = "application/zip")] + Zip, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct LocalShellSkill { + pub name: String, + + pub description: String, + + pub path: String, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum ShellEnvironment { + ContainerAuto { + #[serde(rename = "fileIds", default, skip_serializing_if = "Option::is_none")] + file_ids: Option>, + #[serde( + rename = "memoryLimit", + default, + skip_serializing_if = "Option::is_none" + )] + memory_limit: Option, + #[serde( + rename = "networkPolicy", + default, + skip_serializing_if = "Option::is_none" + )] + network_policy: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + skills: Option>, + }, + ContainerReference { + #[serde(rename = "containerId")] + container_id: String, + }, + Local { + #[serde(default, skip_serializing_if = "Option::is_none")] + skills: Option>, + }, +} + +impl<'de> Deserialize<'de> for ShellEnvironment { + fn deserialize>(deserializer: D) -> Result { + #[derive(Deserialize)] + #[serde(remote = "ShellEnvironment", tag = "type", rename_all = "camelCase")] + enum TaggedEnvironment { + ContainerAuto { + #[serde(rename = "fileIds")] + file_ids: Option>, + #[serde(rename = "memoryLimit")] + memory_limit: Option, + #[serde(rename = "networkPolicy")] + network_policy: Option, + skills: Option>, + }, + ContainerReference { + #[serde(rename = "containerId")] + container_id: String, + }, + Local { + skills: Option>, + }, + } + + let mut value = Value::deserialize(deserializer)?; + if let Some(object) = value.as_object_mut() { + object + .entry("type") + .or_insert_with(|| Value::String("local".into())); + } + TaggedEnvironment::deserialize(value).map_err(serde::de::Error::custom) + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ShellArgs { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub environment: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ToolSearchArgs { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub execution: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub description: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub parameters: Option>, +} + +fn provider_tool(kind: &str, args: impl Serialize) -> Result { + let args = serde_json::to_value(args) + .map_err(|error| AiMuxError::InvalidArgument(error.to_string()))?; + super::responses::tool_args::validate_tool_args(kind, &args)?; + Ok(ProviderTool { + id: format!("openai.{kind}"), + name: kind.to_owned(), + args: args.as_object().cloned().ok_or_else(|| { + AiMuxError::InvalidArgument("provider tool arguments must be an object".into()) + })?, + }) +} + +/// Configure the `web_search` provider tool. +/// +/// # Errors +/// Returns an error when arguments do not match the provider tool schema. +pub fn web_search(args: WebSearchArgs) -> Result { + provider_tool("web_search", args) +} + +/// Configure the `web_search_preview` provider tool. +/// +/// # Errors +/// Returns an error when arguments do not match the provider tool schema. +pub fn web_search_preview(args: WebSearchPreviewArgs) -> Result { + provider_tool("web_search_preview", args) +} + +/// Configure the `file_search` provider tool. +/// +/// # Errors +/// Returns an error when arguments do not match the provider tool schema. +pub fn file_search(args: FileSearchArgs) -> Result { + provider_tool("file_search", args) +} + +/// Configure the `code_interpreter` provider tool. +/// +/// # Errors +/// Returns an error when arguments do not match the provider tool schema. +pub fn code_interpreter(args: CodeInterpreterArgs) -> Result { + provider_tool("code_interpreter", args) +} + +/// Configure the `image_generation` provider tool. +/// +/// # Errors +/// Returns an error when arguments do not match the provider tool schema. +pub fn image_generation(args: ImageGenerationArgs) -> Result { + provider_tool("image_generation", args) +} + +/// Configure the `custom_tool` provider tool. +/// +/// # Errors +/// Returns an error when arguments do not match the provider tool schema. +pub fn custom_tool(args: CustomToolArgs) -> Result { + provider_tool("custom", args) +} + +/// Configure the `mcp` provider tool. +/// +/// # Errors +/// Returns an error when arguments do not match the provider tool schema. +pub fn mcp(args: McpArgs) -> Result { + provider_tool("mcp", args) +} + +/// Configure the `shell` provider tool. +/// +/// # Errors +/// Returns an error when arguments do not match the provider tool schema. +pub fn shell(args: ShellArgs) -> Result { + provider_tool("shell", args) +} + +/// Configure the `tool_search` provider tool. +/// +/// # Errors +/// Returns an error when arguments do not match the provider tool schema. +pub fn tool_search(args: ToolSearchArgs) -> Result { + provider_tool("tool_search", args) +} + +/// Create the `apply_patch` provider tool. +#[must_use] +pub fn apply_patch() -> ProviderTool { + ProviderTool { + id: "openai.apply_patch".into(), + name: "apply_patch".into(), + args: serde_json::Map::new(), + } +} + +/// Create the `computer` provider tool. +#[must_use] +pub fn computer() -> ProviderTool { + ProviderTool { + id: "openai.computer".into(), + name: "computer".into(), + args: serde_json::Map::new(), + } +} + +/// Create the `local_shell` provider tool. +#[must_use] +pub fn local_shell() -> ProviderTool { + ProviderTool { + id: "openai.local_shell".into(), + name: "local_shell".into(), + args: serde_json::Map::new(), + } +} + +/// Create the `programmatic_tool_calling` provider tool. +#[must_use] +pub fn programmatic_tool_calling() -> ProviderTool { + ProviderTool { + id: "openai.programmatic_tool_calling".into(), + name: "programmatic_tool_calling".into(), + args: serde_json::Map::new(), + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct OpenAIComputerSafetyCheck { + pub id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub code: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub message: Option, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ComputerButton { + Left, + Right, + Wheel, + Back, + Forward, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ComputerPoint { + pub x: f64, + pub y: f64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum OpenAIComputerAction { + Click { + button: ComputerButton, + x: f64, + y: f64, + #[serde(default, skip_serializing_if = "Option::is_none")] + keys: Option>, + }, + DoubleClick { + x: f64, + y: f64, + #[serde(default, skip_serializing_if = "Option::is_none")] + keys: Option>, + }, + Drag { + path: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + keys: Option>, + }, + Keypress { + keys: Vec, + }, + Move { + x: f64, + y: f64, + #[serde(default, skip_serializing_if = "Option::is_none")] + keys: Option>, + }, + Screenshot, + Scroll { + x: f64, + y: f64, + #[serde(rename = "scrollX")] + scroll_x: f64, + #[serde(rename = "scrollY")] + scroll_y: f64, + #[serde(default, skip_serializing_if = "Option::is_none")] + keys: Option>, + }, + Type { + text: String, + }, + Wait, +} diff --git a/aimux-providers/src/openai/transcription.rs b/aimux-providers/src/openai/transcription.rs index cd302c7ba..7e56bbd65 100644 --- a/aimux-providers/src/openai/transcription.rs +++ b/aimux-providers/src/openai/transcription.rs @@ -281,8 +281,13 @@ impl TranscriptionModel for OpenAITranscriptionModel { let is_gpt4o_transcribe_model = self.model_id == "gpt-4o-transcribe" || self.model_id == "gpt-4o-mini-transcribe"; - // For non-whisper models, set response_format based on model type. - if self.model_id != "whisper-1" { + let has_openai_options = options + .provider_options + .as_ref() + .is_some_and(|options| options.contains_key("openai")); + + // Upstream only selects a non-default format when its options exist. + if self.model_id != "whisper-1" && has_openai_options { let response_format = if is_gpt4o_transcribe_model { "json" } else { @@ -313,7 +318,10 @@ impl TranscriptionModel for OpenAITranscriptionModel { // Temperature default is 0 when provider options are present (matching // the TS schema's `.default(0)`). - if options.provider_options.is_some() && openai_options.temperature.is_none() { + if has_openai_options && openai_options.timestamp_granularities.is_none() { + form.text("timestamp_granularities[]", "segment")?; + } + if has_openai_options && openai_options.temperature.is_none() { form.text("temperature", "0")?; } diff --git a/aimux-providers/src/openai/types.rs b/aimux-providers/src/openai/types.rs index c5d9d2e42..fa7ca10ea 100644 --- a/aimux-providers/src/openai/types.rs +++ b/aimux-providers/src/openai/types.rs @@ -9,12 +9,12 @@ use serde_json::Value; #[derive(Debug, Deserialize)] pub struct ChatCompletionResponse { - pub id: String, - pub model: String, + pub id: Option, + pub model: Option, #[serde(default)] pub created: Option, pub choices: Vec, - pub usage: UsageResponse, + pub usage: Option, } #[derive(Debug, Deserialize)] @@ -29,8 +29,10 @@ pub struct Choice { #[derive(Debug, Deserialize)] pub struct MessageResponse { - pub role: String, + pub role: Option, pub content: Option, + #[serde(default)] + pub audio: Option, /// Reasoning text (`reasoning` field of OpenAI-shaped servers). #[serde(default)] pub reasoning: Option, @@ -45,6 +47,11 @@ pub struct MessageResponse { pub annotations: Option>, } +#[derive(Debug, Deserialize)] +pub struct AudioResponse { + pub transcript: Option, +} + #[derive(Debug, Deserialize)] pub struct ToolCallResponse { pub id: String, @@ -58,7 +65,7 @@ pub struct FunctionCallResponse { pub arguments: String, } -#[derive(Debug, Deserialize, Clone, Serialize)] +#[derive(Debug, Default, Deserialize, Clone, Serialize)] pub struct UsageResponse { #[serde(default)] pub prompt_tokens: Option, @@ -146,12 +153,13 @@ pub struct Delta { #[derive(Debug, Deserialize)] pub struct DeltaToolCall { - /// Optional: some OpenAI-compatible providers omit it (the AI SDK schema - /// is `index: z.number().nullish()`); the tracker falls back to id/name. - #[serde(default)] - pub index: Option, + /// Required, as in `@ai-sdk/openai` (`index: z.number()`); a chunk without + /// it fails to parse. Servers that omit it go through `openai_compatible`. + pub index: usize, #[serde(default)] pub id: Option, + #[serde(default, rename = "type")] + pub r#type: Option, #[serde(default)] pub function: Option, } diff --git a/aimux-providers/src/openai_compatible/chat.rs b/aimux-providers/src/openai_compatible/chat.rs index 89b0e3da2..7b94fb9b3 100644 --- a/aimux-providers/src/openai_compatible/chat.rs +++ b/aimux-providers/src/openai_compatible/chat.rs @@ -15,7 +15,7 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::{LanguageModel, SupportedUrls}; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; -use aimux_core::shared::provider_namespace; +use aimux_core::shared::{RequestInfo, ResponseInfo, StreamResponseInfo, provider_namespace}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{ FinishReason, FinishReasonUnified, InputTokenUsage, OutputTokenUsage, ResponseMetadata, Usage, @@ -60,7 +60,6 @@ impl OpenAICompatibleChatModel { provider_options_name: self.config.provider_options_name(), include_usage: self.config.chat.include_usage, supports_structured_outputs: self.config.chat.supports_structured_outputs, - supports_multi_part_tool_content: self.config.chat.supports_multi_part_tool_content, dialect: &self.config.chat.dialect, }, ) @@ -72,27 +71,20 @@ impl OpenAICompatibleChatModel { /// Core usage from a parsed OpenAI-shaped `usage` object. /// /// `raw` is the provider's original object, kept verbatim in `Usage.raw`. -/// `cached_tokens` at the top level (Moonshot) wins over the nested value. +/// Cache reads come from `prompt_tokens_details`, as in the upstream converter. pub(crate) fn convert_usage(usage: &UsageResponse, raw: Option<&Value>) -> Usage { let prompt_tokens = usage.prompt_tokens.unwrap_or(0); let completion_tokens = usage.completion_tokens.unwrap_or(0); - let nested_cached = usage + let cached = usage .prompt_tokens_details .as_ref() .and_then(|d| d.cached_tokens) .unwrap_or(0); - let cached = usage.cached_tokens.unwrap_or(nested_cached); - let cache_write = usage - .prompt_tokens_details - .as_ref() - .and_then(|d| d.cache_write_tokens); // Saturating: some servers report more cached or reasoning tokens than // prompt or completion tokens. - let no_cache = prompt_tokens - .saturating_sub(cached) - .saturating_sub(cache_write.unwrap_or(0)); + let no_cache = prompt_tokens.saturating_sub(cached); let reasoning_tokens = usage .completion_tokens_details .as_ref() @@ -104,14 +96,14 @@ pub(crate) fn convert_usage(usage: &UsageResponse, raw: Option<&Value>) -> Usage total: Some(prompt_tokens), no_cache: Some(no_cache), cache_read: Some(cached), - cache_write, + cache_write: None, }, output_tokens: OutputTokenUsage { total: Some(completion_tokens), text: Some(completion_tokens.saturating_sub(reasoning_tokens)), reasoning: Some(reasoning_tokens), }, - raw: raw.and_then(|value| value.as_object().cloned()), + raw: raw.and_then(Value::as_object).cloned(), } } @@ -199,10 +191,13 @@ impl LanguageModel for OpenAICompatibleChatModel { fn supported_urls(&self) -> SupportedUrls { self.config .chat + .dialect .supported_urls .as_ref() - .map(|urls| urls()) - .unwrap_or_default() + .map_or_else( + || self.config.chat.supported_urls.clone(), + |supported_urls| supported_urls(), + ) } async fn do_generate(&self, options: &CallOptions) -> Result { @@ -234,12 +229,7 @@ impl LanguageModel for OpenAICompatibleChatModel { })?; let mut content = Vec::new(); - for text in content_texts(choice.message.content.as_ref()) { - content.push(GenerateContent::Text { - text, - provider_metadata: None, - }); - } + content.extend(content_parts(choice.message.content.as_ref())); let reasoning = choice .message .reasoning_content @@ -298,18 +288,13 @@ impl LanguageModel for OpenAICompatibleChatModel { }); let dialect = &self.config.chat.dialect; - let usage = self - .config - .chat - .convert_usage - .as_ref() - .map(|convert| convert(raw.get("usage").unwrap_or(&Value::Null))) - .unwrap_or_else(|| usage_from_raw(raw.get("usage"))); + let usage = dialect.convert_usage.as_ref().map_or_else( + || usage_from_raw(raw.get("usage")), + |convert| convert(raw.get("usage")), + ); - let mut namespace = Map::new(); - prediction_tokens(data.usage.as_ref(), &mut namespace); let mut metadata = HashMap::new(); - metadata.insert(metadata_key, namespace); + metadata.insert(metadata_key.clone(), Map::new()); if let Some(extractor) = &dialect.metadata_extractor && let Some(extra) = extractor.extract_metadata(&raw) { @@ -318,23 +303,29 @@ impl LanguageModel for OpenAICompatibleChatModel { } } + prediction_tokens( + data.usage.as_ref(), + metadata.entry(metadata_key).or_default(), + ); + Ok(GenerateResult { content, finish_reason, usage, warnings: built.warnings, provider_metadata: Some(metadata), - response: Some(aimux_core::shared::ResponseInfo { + response: Some(ResponseInfo { id: data.id, timestamp: data .created + .filter(|created| *created != 0) .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) .map(|dt| dt.to_rfc3339()), model_id: data.model, headers: Some(response_headers), body: Some(raw), }), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + request: Some(RequestInfo { body: Some(body) }), }) } @@ -362,31 +353,11 @@ impl LanguageModel for OpenAICompatibleChatModel { let response_headers = resp.response_headers; let mut sse_stream = resp.value; - // Only the first event is checked before returning the stream: an - // error reported immediately stays inside Core's operation-retry - // boundary. A normal event is chained back and never consumed. - let first_event = match sse_stream.next().await { - Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), - first => first, - }; - if let Some(Ok(event)) = &first_event - && let Some(error) = event.get("error") - { - return Err(stream_error( - &metadata_key, - error, - &endpoint, - body.clone(), - response_headers.clone(), - )); - } - let dialect = self.config.chat.dialect.clone(); let mut metadata_extractor = dialect .metadata_extractor .as_ref() .map(|extractor| extractor.create_stream_extractor()); - let convert_usage = self.config.chat.convert_usage.clone(); let emit_raw_chunks = options.include_raw_chunks == Some(true); let stream_error_url = endpoint; let stream_error_body = body.clone(); @@ -395,11 +366,11 @@ impl LanguageModel for OpenAICompatibleChatModel { let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings }); - let text_id = 0usize; + let text_id = "txt-0"; let mut text_started = false; let reasoning_id = "reasoning-0".to_string(); let mut reasoning_started = false; - let mut final_usage_raw: Option = None; + let mut final_usage_raw = None; let mut final_usage_parsed: Option = None; let mut final_finish_reason: Option = None; let mut response_metadata_emitted = false; @@ -418,14 +389,7 @@ impl LanguageModel for OpenAICompatibleChatModel { let mut pending: HashMap = HashMap::new(); let mut forwarded: HashSet = HashSet::new(); - let mut event_iter = - futures::stream::iter(first_event.into_iter()).chain(sse_stream); - let mut stream_errored = false; - - while let Some(event) = event_iter.next().await { - if stream_errored { - break; - } + while let Some(event) = sse_stream.next().await { match event { Ok(parsed) => { if emit_raw_chunks { @@ -444,8 +408,8 @@ impl LanguageModel for OpenAICompatibleChatModel { stream_response_headers.clone(), ), }); - stream_errored = true; - break; + final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); + continue; } let chunk_usage_raw: Option = match &dialect.stream_usage_key { @@ -457,19 +421,23 @@ impl LanguageModel for OpenAICompatibleChatModel { let chunk: StreamChunk = match serde_json::from_value(parsed) { Ok(chunk) => chunk, Err(e) => { - yield Err(e.into()); + final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); + yield Ok(StreamPart::Error { error: e.into() }); continue; } }; if !response_metadata_emitted - && (chunk.id.is_some() || chunk.model.is_some()) + && (chunk.id.as_ref().is_some_and(|id| !id.is_empty()) + || chunk.model.as_ref().is_some_and(|model| !model.is_empty()) + || chunk.created.is_some_and(|created| created != 0)) { response_metadata_emitted = true; yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { id: chunk.id.clone(), timestamp: chunk .created + .filter(|created| *created != 0) .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) .map(|dt| dt.to_rfc3339()), model_id: chunk.model.clone(), @@ -481,52 +449,49 @@ impl LanguageModel for OpenAICompatibleChatModel { final_usage_parsed = serde_json::from_value(raw_usage.clone()).ok(); } - for choice in chunk.choices { + for choice in chunk.choices.into_iter().take(1) { let reasoning_delta = choice .delta .reasoning_content .clone() .or_else(|| choice.delta.reasoning.clone()); - if let Some(reasoning) = reasoning_delta - && !reasoning.is_empty() - { - if !reasoning_started { - reasoning_started = true; - yield Ok(StreamPart::ReasoningStart { - id: reasoning_id.clone(), - provider_metadata: None, - }); - } - yield Ok(StreamPart::ReasoningDelta { - id: reasoning_id.clone(), - delta: reasoning, + let mut delta_content = Vec::new(); + if let Some(text) = reasoning_delta.filter(|text| !text.is_empty()) { + delta_content.push(GenerateContent::Reasoning(ReasoningOutput { + text, provider_metadata: None, - }); + })); } - - if let Some(content) = choice.delta.content { - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, - }); - reasoning_started = false; - } - if !text_started { - text_started = true; - yield Ok(StreamPart::TextStart { - id: format!("{text_id}"), - provider_metadata: None, - }); + delta_content.extend(content_parts(choice.delta.content.as_ref())); + for part in delta_content { + match part { + GenerateContent::Reasoning(ReasoningOutput { text, .. }) => { + if text_started { + yield Ok(StreamPart::TextEnd { id: text_id.to_string(), provider_metadata: None }); + text_started = false; + } + if !reasoning_started { + reasoning_started = true; + yield Ok(StreamPart::ReasoningStart { id: reasoning_id.clone(), provider_metadata: None }); + } + yield Ok(StreamPart::ReasoningDelta { id: reasoning_id.clone(), delta: text, provider_metadata: None }); + } + GenerateContent::Text { text, .. } => { + if reasoning_started { + yield Ok(StreamPart::ReasoningEnd { id: reasoning_id.clone(), provider_metadata: None }); + reasoning_started = false; + } + if !text_started { + text_started = true; + yield Ok(StreamPart::TextStart { id: text_id.to_string(), provider_metadata: None }); + } + yield Ok(StreamPart::TextDelta { id: text_id.to_string(), delta: text, provider_metadata: None }); + } + _ => unreachable!(), } - yield Ok(StreamPart::TextDelta { - id: format!("{text_id}"), - delta: content, - provider_metadata: None, - }); } - if let Some(deltas) = choice.delta.tool_calls { + if let Some(deltas) = choice.delta.tool_calls.filter(|deltas| !deltas.is_empty()) { if reasoning_started { yield Ok(StreamPart::ReasoningEnd { id: reasoning_id.clone(), @@ -610,30 +575,13 @@ impl LanguageModel for OpenAICompatibleChatModel { } if let Some(reason) = choice.finish_reason { - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, - }); - reasoning_started = false; - } - if text_started { - yield Ok(StreamPart::TextEnd { - id: format!("{text_id}"), - provider_metadata: None, - }); - text_started = false; - } final_finish_reason = Some(parse_finish_reason(&reason)); } } } Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; - } + final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); + yield Ok(StreamPart::Error { error }); } } } @@ -646,10 +594,26 @@ impl LanguageModel for OpenAICompatibleChatModel { } if text_started { yield Ok(StreamPart::TextEnd { - id: format!("{text_id}"), + id: text_id.to_string(), provider_metadata: None, }); } + for (index, entry) in pending { + let delta = StreamingToolCallDelta { + index: Some(index), + id: entry.id.as_deref(), + arguments: Some(&entry.arguments), + provider_metadata: signature_metadata(Some(&entry.extra)), + ..Default::default() + }; + if let Err(error) = tool_calls.process(delta, &mut tool_parts) { + yield Err(error.into()); + return; + } + for part in tool_parts.drain(..) { + yield Ok(part); + } + } // A parsable argument buffer can still be a prefix of a longer // input: finalize only when the stream flushes. tool_calls.finish(&mut tool_parts); @@ -657,10 +621,8 @@ impl LanguageModel for OpenAICompatibleChatModel { yield Ok(part); } - let mut namespace = Map::new(); - prediction_tokens(final_usage_parsed.as_ref(), &mut namespace); let mut metadata = HashMap::new(); - metadata.insert(metadata_key.clone(), namespace); + metadata.insert(metadata_key.clone(), Map::new()); if let Some(extractor) = &metadata_extractor && let Some(extra) = extractor.build_metadata() { @@ -669,27 +631,30 @@ impl LanguageModel for OpenAICompatibleChatModel { } } + prediction_tokens(final_usage_parsed.as_ref(), metadata.entry(metadata_key).or_default()); + + if final_finish_reason.is_none() { + yield Ok(StreamPart::Error { + error: AiMuxError::InvalidResponseData("Response stream ended without a finish reason.".to_string()), + }); + } yield Ok(StreamPart::Finish { - finish_reason: if stream_errored { - FinishReason { unified: FinishReasonUnified::Error, raw: None } - } else { - final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }) - }, - usage: if stream_errored { Usage::default() } else { - convert_usage.as_ref().map(|convert| convert(final_usage_raw.as_ref().unwrap_or(&Value::Null))) - .unwrap_or_else(|| usage_from_raw(final_usage_raw.as_ref())) - }, + finish_reason: final_finish_reason.unwrap_or(FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }), + usage: dialect.convert_usage.as_ref().map_or_else( + || usage_from_raw(final_usage_raw.as_ref()), + |convert| convert(final_usage_raw.as_ref()), + ), provider_metadata: Some(metadata), }); }; Ok(StreamResult { stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { + request: Some(RequestInfo { body: Some(body) }), + response: Some(StreamResponseInfo { headers: Some(response_headers), }), }) @@ -703,20 +668,54 @@ struct PendingToolCall { extra: Value, } -/// Text segments of a message `content`: a string, or an array of -/// `{ type: "text", text }` parts. -fn content_texts(content: Option<&Value>) -> Vec { - match content { - Some(Value::String(text)) if !text.is_empty() => vec![text.clone()], +/// Normalize text and thinking parts, preserving their order. +fn content_parts(content: Option<&Value>) -> Vec { + let parts: Vec<(bool, String)> = match content { + Some(Value::String(text)) if !text.is_empty() => vec![(false, text.clone())], Some(Value::Array(parts)) => parts .iter() - .filter(|part| part.get("type").and_then(Value::as_str) == Some("text")) - .filter_map(|part| part.get("text").and_then(Value::as_str)) - .filter(|text| !text.is_empty()) - .map(str::to_string) + .filter_map(|part| match part.get("type").and_then(Value::as_str) { + Some("text") => part + .get("text") + .and_then(Value::as_str) + .map(|text| (false, text.to_string())), + Some("thinking") => part + .get("thinking") + .and_then(Value::as_array) + .map(|chunks| { + ( + true, + chunks + .iter() + .filter(|chunk| { + chunk.get("type").and_then(Value::as_str) == Some("text") + }) + .filter_map(|chunk| chunk.get("text").and_then(Value::as_str)) + .collect::(), + ) + }), + _ => None, + }) + .filter(|(_, text)| !text.is_empty()) .collect(), _ => Vec::new(), - } + }; + parts + .into_iter() + .map(|(reasoning, text)| { + if reasoning { + GenerateContent::Reasoning(ReasoningOutput { + text, + provider_metadata: None, + }) + } else { + GenerateContent::Text { + text, + provider_metadata: None, + } + } + }) + .collect() } /// One `GET {base_url}/models` exchange for discovery: no retry, no @@ -764,120 +763,3 @@ pub(crate) async fn list_models_once( }) .collect()) } - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use serde_json::{Value, json}; - use wiremock::matchers::{method, path}; - use wiremock::{Mock, MockServer, ResponseTemplate}; - - use aimux_core::error::AiMuxError; - use aimux_core::language_model::LanguageModel; - use aimux_core::language_model_message::LanguageModelMessage; - use aimux_core::options::CallOptions; - use aimux_provider_utils::ProviderErrorParts; - - use crate::openai_compatible::config::{BaseUrl, ChatDialect}; - use crate::openai_compatible::{Assembly, ChatProfile, OpenAICompatibleProvider}; - use crate::shared::Credential; - - fn hello() -> CallOptions { - CallOptions::new(vec![LanguageModelMessage::user_text("hi")]) - } - - fn provider(base_url: BaseUrl, dialect: ChatDialect) -> OpenAICompatibleProvider { - OpenAICompatibleProvider::assemble(Assembly { - name: "acme".to_string(), - base_url, - credential: Credential::None, - fixed_headers: Vec::new(), - headers: None, - query_params: None, - fetch: None, - transform_request_body: None, - profile: ChatProfile { - include_usage: false, - supports_structured_outputs: false, - supports_multi_part_tool_content: false, - dialect, - supported_urls: None, - convert_usage: None, - }, - }) - .unwrap() - } - - async fn serve(server: &MockServer, response: ResponseTemplate) { - Mock::given(method("POST")) - .and(path("/chat/completions")) - .respond_with(response) - .mount(server) - .await; - } - - #[tokio::test] - async fn an_error_structure_decides_the_message_and_code() { - let server = MockServer::start().await; - serve( - &server, - ResponseTemplate::new(400).set_body_json(json!({"detail": "bad prompt", "kind": 42})), - ) - .await; - let mut dialect = ChatDialect::baseline(); - dialect.error_structure = Arc::new(|data: &Value| ProviderErrorParts { - message: data["detail"].as_str().unwrap_or_default().to_string(), - provider_code: data["kind"].as_u64().map(|k| k.to_string()), - }); - let error = provider(BaseUrl::Fixed(server.uri()), dialect) - .chat("m") - .do_generate(&hello()) - .await - .unwrap_err(); - match error { - AiMuxError::ApiCall(detail) => { - assert_eq!(detail.message, "bad prompt"); - assert_eq!(detail.provider_code.as_deref(), Some("42")); - } - other => panic!("expected ApiCall, got {other:?}"), - } - } - - #[tokio::test] - async fn a_lazy_base_url_is_evaluated_per_request_and_its_error_fails_the_request() { - let server = MockServer::start().await; - serve( - &server, - ResponseTemplate::new(200).set_body_json(json!({ - "id": "c", "model": "m", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}] - })), - ) - .await; - let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let counter = calls.clone(); - let uri = server.uri(); - let lazy = BaseUrl::Lazy(Arc::new(move || { - if counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { - Err(AiMuxError::InvalidArgument("not ready".into())) - } else { - Ok(uri.clone()) - } - })); - let provider = provider(lazy, ChatDialect::baseline()); - let model = provider.chat("m"); - assert_eq!( - calls.load(std::sync::atomic::Ordering::SeqCst), - 0, - "creating a provider or a model resolves nothing" - ); - let error = model.do_generate(&hello()).await.unwrap_err(); - assert!(matches!(error, AiMuxError::InvalidArgument(_)), "{error:?}"); - model - .do_generate(&hello()) - .await - .expect("resolved on the next request"); - assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2); - } -} diff --git a/aimux-providers/src/openai_compatible/config.rs b/aimux-providers/src/openai_compatible/config.rs index 630faa0fe..8a367057c 100644 --- a/aimux-providers/src/openai_compatible/config.rs +++ b/aimux-providers/src/openai_compatible/config.rs @@ -21,6 +21,9 @@ use aimux_provider_utils::{ pub use crate::shared::TransformRequestBody; +/// Converts the raw optional token usage of a chat response. +pub type ConvertUsage = Arc) -> Usage + Send + Sync>; + /// Where the base URL comes from. /// /// A plain provider has a [`Fixed`](Self::Fixed) URL. A preset whose URL is @@ -82,6 +85,8 @@ pub(crate) struct ChatDialect { /// Streaming usage rides in `chunk[key].usage` instead of `chunk.usage`. pub stream_usage_key: Option, pub metadata_extractor: Option>, + pub supported_urls: Option SupportedUrls + Send + Sync>>, + pub convert_usage: Option, pub error_structure: ErrorStructure, } @@ -95,6 +100,8 @@ impl ChatDialect { max_tokens_key: None, stream_usage_key: None, metadata_extractor: None, + supported_urls: None, + convert_usage: None, error_structure: Arc::new(default_error_structure), } } @@ -120,8 +127,6 @@ pub(crate) fn default_error_structure(data: &Value) -> ProviderErrorParts { } } -/// Converts the upstream token usage object into language-model usage. -pub type ConvertUsage = Arc Usage + Send + Sync>; /// Returns the URLs supported by chat models. pub type SupportedUrlsFn = Arc SupportedUrls + Send + Sync>; @@ -130,9 +135,7 @@ pub type SupportedUrlsFn = Arc SupportedUrls + Send + Sync>; pub(crate) struct ChatSettings { pub include_usage: bool, pub supports_structured_outputs: bool, - pub supports_multi_part_tool_content: bool, - pub supported_urls: Option, - pub convert_usage: Option, + pub supported_urls: SupportedUrls, pub dialect: Arc, } diff --git a/aimux-providers/src/openai_compatible/convert.rs b/aimux-providers/src/openai_compatible/convert.rs index b61494981..af80c6a36 100644 --- a/aimux-providers/src/openai_compatible/convert.rs +++ b/aimux-providers/src/openai_compatible/convert.rs @@ -7,7 +7,7 @@ //! `reasoning_effort`, `verbosity`, `messages`, `tools`, `tool_choice`). //! Compatible endpoint capabilities arrive as [`ChatDialect`] data. -use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; +use aimux_core::shared::SharedProviderOptions; use base64::Engine; use serde_json::{Map, Value, json}; @@ -17,6 +17,7 @@ use aimux_core::language_model_message::{ ToolCallPart, ToolPart, ToolResultPart, UserPart, }; use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; +use aimux_core::shared::{FileBytes, FileData}; use aimux_core::tool::Tool; use aimux_core::types::{FinishReason, FinishReasonUnified, Warning}; use aimux_provider_utils::{get_top_level_media_type, resolve_full_media_type}; @@ -85,10 +86,8 @@ pub(crate) fn prepare_function_tools(options: &CallOptions) -> PreparedTools { }), } } - if !wire.is_empty() { - prepared.tools = Some(wire); - prepared.tool_choice = options.tool_choice.as_ref().map(tool_choice_value); - } + prepared.tools = Some(wire); + prepared.tool_choice = options.tool_choice.as_ref().map(tool_choice_value); prepared } @@ -116,7 +115,6 @@ pub(crate) struct ChatBodySpec<'a> { pub provider_options_name: &'a str, pub include_usage: bool, pub supports_structured_outputs: bool, - pub supports_multi_part_tool_content: bool, pub dialect: &'a ChatDialect, } @@ -139,10 +137,10 @@ fn namespace<'a>(options: &'a CallOptions, key: &str) -> Option<&'a Map String { let camel = to_camel_case(name); - if camel != name && namespace(options, name).is_none() && namespace(options, &camel).is_some() { + if camel != name && namespace(options, &camel).is_some() { camel } else { name.to_string() @@ -155,7 +153,7 @@ fn option_string( key: &str, ) -> Result, AiMuxError> { match merged.get(key) { - None | Some(Value::Null) => Ok(None), + None => Ok(None), Some(Value::String(value)) => Ok(Some(value.clone())), Some(_) => Err(AiMuxError::InvalidArgument(format!( "invalid provider options for \"{namespace_name}\": `{key}` must be a string" @@ -163,6 +161,21 @@ fn option_string( } } +fn validate_options(options: &Map, namespace_name: &str) -> Result<(), AiMuxError> { + for key in ["user", "reasoningEffort", "textVerbosity"] { + option_string(options, namespace_name, key)?; + } + if options + .get("strictJsonSchema") + .is_some_and(|value| !value.is_boolean()) + { + return Err(AiMuxError::InvalidArgument(format!( + "invalid provider options for \"{namespace_name}\": `strictJsonSchema` must be a boolean" + ))); + } + Ok(()) +} + /// Build the request body (`getArgs`). /// /// # Errors @@ -184,6 +197,7 @@ pub(crate) fn build_request_body( // own, then its camelCase form; later wins per key. let mut merged = Map::new(); if let Some(deprecated) = namespace(options, "openai-compatible") { + validate_options(deprecated, "openai-compatible")?; warnings.push(Warning::Deprecated { setting: "providerOptions key 'openai-compatible'".to_string(), message: "Use 'openaiCompatible' instead.".to_string(), @@ -192,6 +206,7 @@ pub(crate) fn build_request_body( } for key in ["openaiCompatible", name, camel.as_str()] { if let Some(object) = namespace(options, key) { + validate_options(object, key)?; merged.extend(object.iter().map(|(k, v)| (k.clone(), v.clone()))); } } @@ -322,7 +337,6 @@ pub(crate) fn build_request_body( &options.prompt, &MessageSpec { metadata_key: &metadata_key, - supports_multi_part_tool_content: spec.supports_multi_part_tool_content, }, )?; body.insert("messages".into(), Value::Array(messages)); @@ -365,22 +379,15 @@ pub(crate) fn build_request_body( pub(crate) struct MessageSpec<'a> { pub metadata_key: &'a str, - pub supports_multi_part_tool_content: bool, } -/// The provider options of a message or part that belong on the wire object -/// (`getOpenAIMetadata`): the generic namespace, then the provider's own. -fn wire_metadata( - provider_options: Option<&SharedProviderOptions>, - key: &str, -) -> Map { - let mut out = Map::new(); - for ns in ["openaiCompatible", key] { - if let Some(object) = provider_options.and_then(|options| options.get(ns)) { - out.extend(object.iter().map(|(k, v)| (k.clone(), v.clone()))); - } - } - out +/// The generic `openaiCompatible` metadata of a message or part +/// (`getOpenAIMetadata`). +fn wire_metadata(provider_options: Option<&SharedProviderOptions>) -> Map { + provider_options + .and_then(|options| options.get("openaiCompatible")) + .cloned() + .unwrap_or_default() } fn with_metadata(mut object: Value, metadata: Map) -> Value { @@ -411,39 +418,48 @@ fn convert_message( message: &LanguageModelMessage, spec: &MessageSpec<'_>, ) -> Result, AiMuxError> { - let key = spec.metadata_key; match message { - LanguageModelMessage::System { content, provider_options } => Ok(vec![with_metadata( + LanguageModelMessage::System { + content, + provider_options, + } => Ok(vec![with_metadata( json!({ "role": "system", "content": content }), - wire_metadata(provider_options.as_ref(), key), + wire_metadata(provider_options.as_ref()), )]), - LanguageModelMessage::User { content, provider_options } => Ok(vec![convert_user_message( - content, spec, wire_metadata(provider_options.as_ref(), key), + LanguageModelMessage::User { + content, + provider_options, + } => Ok(vec![convert_user_message( + content, + wire_metadata(provider_options.as_ref()), )?]), - LanguageModelMessage::Assistant { content, provider_options } => Ok(vec![convert_assistant_message( - content, spec, wire_metadata(provider_options.as_ref(), key), + LanguageModelMessage::Assistant { + content, + provider_options, + } => Ok(vec![convert_assistant_message( + content, + spec, + wire_metadata(provider_options.as_ref()), )]), LanguageModelMessage::Tool { content, .. } => Ok(content .iter() .filter_map(|part| { let ToolPart::ToolResult(ToolResultPart { - tool_call_id, output, provider_options, .. - }) = part else { return None; }; + tool_call_id, + output, + provider_options, + .. + }) = part + else { + return None; + }; Some(with_metadata( json!({ "role": "tool", "tool_call_id": tool_call_id, - "content": if spec.supports_multi_part_tool_content { - if let aimux_core::language_model_message::ToolResultOutput::Content { value } = output { - crate::openai::convert::tool_result_content_value(value) - } else { - crate::openai::convert::tool_result_to_content(output) - } - } else { - crate::openai::convert::tool_result_to_content(output) - }, + "content": crate::openai::convert::tool_result_to_content(output), }), - wire_metadata(provider_options.as_ref(), key), + wire_metadata(provider_options.as_ref()), )) }) .collect()), @@ -452,10 +468,8 @@ fn convert_message( fn convert_user_message( content: &[UserPart], - spec: &MessageSpec<'_>, message_metadata: Map, ) -> Result { - let key = spec.metadata_key; if let [ UserPart::Text(TextPart { text, @@ -463,8 +477,7 @@ fn convert_user_message( }), ] = content { - let mut metadata = message_metadata; - metadata.extend(wire_metadata(provider_options.as_ref(), key)); + let metadata = wire_metadata(provider_options.as_ref()); return Ok(with_metadata( json!({ "role": "user", "content": text }), metadata, @@ -472,71 +485,75 @@ fn convert_user_message( } let parts = content .iter() - .map(|part| { - let (wire, provider_options) = match part { - UserPart::Text(TextPart { - text, - provider_options, - }) => (json!({ "type": "text", "text": text }), provider_options), - UserPart::File(file) => (file_part(file, key)?, &file.provider_options), - }; - Ok(with_metadata( - wire, - wire_metadata(provider_options.as_ref(), key), - )) - }) - .collect::, AiMuxError>>()?; + .map(convert_user_part) + .collect::, _>>()?; Ok(with_metadata( json!({ "role": "user", "content": parts }), message_metadata, )) } -fn file_part(part: &FilePart, key: &str) -> Result { - let unsupported = |what: String| AiMuxError::InvalidArgument(what); - if let FileData::Reference { reference } = &part.data { - let file_id = reference - .get(key) - .or_else(|| reference.get("openaiCompatible")) - .ok_or_else(|| { - let mut available: Vec<&str> = reference.keys().map(String::as_str).collect(); - available.sort_unstable(); - unsupported(format!( - "No provider reference found for provider '{key}'. Available providers: {}", - available.join(", ") - )) - })?; - return Ok(json!({ "type": "file", "file": { "file_id": file_id } })); - } - let base64_data = match &part.data { - FileData::Data { - data: FileBytes::Binary(bytes), - } => Some(std::borrow::Cow::Owned( - base64::engine::general_purpose::STANDARD.encode(bytes), - )), - FileData::Data { - data: FileBytes::Base64(data), - } => Some(std::borrow::Cow::Borrowed(data.as_str())), - FileData::Url { .. } => None, - FileData::Text { .. } => return Err(unsupported("text file parts".to_string())), - FileData::Reference { .. } => unreachable!("handled above"), +fn convert_user_part(part: &UserPart) -> Result { + let (wire, provider_options) = match part { + UserPart::Text(TextPart { + text, + provider_options, + }) => (json!({ "type": "text", "text": text }), provider_options), + UserPart::File(file) => { + let wire = match &file.data { + FileData::Data { + data: FileBytes::Binary(data), + } => file_part( + file, + FileSource::Base64(&base64::engine::general_purpose::STANDARD.encode(data)), + )?, + FileData::Data { + data: FileBytes::Base64(data), + } => file_part(file, FileSource::Base64(data))?, + FileData::Url { url, .. } => file_part(file, FileSource::Url(url))?, + FileData::Reference { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "file parts with provider references".to_string(), + )); + } + FileData::Text { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "text file parts".to_string(), + )); + } + }; + (wire, &file.provider_options) + } }; + Ok(with_metadata( + wire, + wire_metadata(provider_options.as_ref()), + )) +} + +enum FileSource<'a> { + Base64(&'a str), + Url(&'a str), +} + +fn file_part(part: &FilePart, source: FileSource<'_>) -> Result { let media_type = part.media_type.as_str(); + let filename = part.filename.as_deref(); + let unsupported = |what: String| AiMuxError::UnsupportedFunctionality(what); + match get_top_level_media_type(media_type) { kind @ ("image" | "video") => { - let url = match &part.data { - FileData::Url { url, .. } => url.to_string(), - _ => format!( - "data:{};base64,{}", - resolve_full_media_type(part)?, - base64_data.as_deref().unwrap() - ), + let url = match source { + FileSource::Url(url) => url.to_string(), + FileSource::Base64(b64) => { + format!("data:{};base64,{}", resolve_full_media_type(part)?, b64) + } }; let kind = format!("{kind}_url"); Ok(json!({ "type": kind, kind: { "url": url } })) } "audio" => { - let Some(b64) = base64_data.as_deref() else { + let FileSource::Base64(b64) = source else { return Err(unsupported("audio file parts with URLs".to_string())); }; let full_media_type = resolve_full_media_type(part)?; @@ -544,15 +561,13 @@ fn file_part(part: &FilePart, key: &str) -> Result { "audio/wav" => "wav", "audio/mp3" | "audio/mpeg" => "mp3", other => { - return Err(unsupported(format!( - "audio content parts with media type {other}" - ))); + return Err(unsupported(format!("audio media type {other}"))); } }; Ok(json!({ "type": "input_audio", "input_audio": { "data": b64, "format": format } })) } "application" => { - let Some(b64) = base64_data.as_deref() else { + let FileSource::Base64(b64) = source else { return Err(unsupported("PDF file parts with URLs".to_string())); }; let full_media_type = resolve_full_media_type(part)?; @@ -561,7 +576,7 @@ fn file_part(part: &FilePart, key: &str) -> Result { "file part media type {full_media_type}" ))); } - let filename = part.filename.as_deref().unwrap_or("document.pdf"); + let filename = filename.unwrap_or("document.pdf"); Ok(json!({ "type": "file", "file": { @@ -571,14 +586,14 @@ fn file_part(part: &FilePart, key: &str) -> Result { })) } "text" => { - let Some(b64) = base64_data.as_deref() else { - return Err(unsupported("text file parts with URLs".to_string())); + let b64 = match source { + FileSource::Url(url) => return Ok(json!({ "type": "text", "text": url })), + FileSource::Base64(b64) => b64, }; let bytes = base64::engine::general_purpose::STANDARD .decode(b64) .map_err(|e| unsupported(format!("text file part is not valid base64: {e}")))?; - let text = String::from_utf8(bytes) - .map_err(|_| unsupported("text file part is not valid UTF-8".to_string()))?; + let text = String::from_utf8_lossy(&bytes); Ok(json!({ "type": "text", "text": text })) } _ => Err(unsupported(format!("file part media type {media_type}"))), @@ -594,16 +609,10 @@ fn convert_assistant_message( let mut text = String::new(); let mut reasoning = String::new(); let mut tool_calls = Vec::new(); - let mut metadata = message_metadata; + let metadata = message_metadata; for part in content { match part { - AssistantPart::Text(TextPart { - text: t, - provider_options, - }) => { - text.push_str(t); - metadata.extend(wire_metadata(provider_options.as_ref(), key)); - } + AssistantPart::Text(TextPart { text: t, .. }) => text.push_str(t), AssistantPart::Reasoning(ReasoningPart { text: t, .. }) => reasoning.push_str(t), AssistantPart::ToolCall(ToolCallPart { tool_call_id, @@ -618,34 +627,29 @@ fn convert_assistant_message( "type": "function", "function": { "name": tool_name, "arguments": arguments }, }); - if let Some(signature) = provider_options - .as_ref() - .and_then(|options| { - options - .get(key) - .and_then(|values| values.get("thoughtSignature")) - .filter(|value| !value.is_null()) - .or_else(|| { - options - .get("google") - .and_then(|values| values.get("thoughtSignature")) - }) - }) - .filter(|value| match value { - Value::Null => false, - Value::Bool(v) => *v, - Value::String(v) => !v.is_empty(), - Value::Number(v) => v.as_f64() != Some(0.0), - _ => true, - }) - { - let signature = signature_string(signature); + call = with_metadata(call, wire_metadata(provider_options.as_ref())); + let signature = provider_options.as_ref().and_then(|options| { + options + .get(key) + .and_then(|options| options.get("thoughtSignature")) + .filter(|value| !value.is_null()) + .or_else(|| { + options + .get("google") + .and_then(|options| options.get("thoughtSignature")) + }) + }); + let signature = signature.and_then(|value| match value { + Value::Null | Value::Bool(false) => None, + Value::String(text) if text.is_empty() => None, + Value::Number(number) if number.as_f64() == Some(0.0) => None, + Value::String(text) => Some(text.clone()), + other => Some(signature_string(other)), + }); + if let Some(signature) = signature { call["extra_content"] = json!({ "google": { "thought_signature": signature } }); } - tool_calls.push(with_metadata( - call, - wire_metadata(provider_options.as_ref(), key), - )); + tool_calls.push(call); } // Provider-executed results and media have no assistant-content // form on this wire. diff --git a/aimux-providers/src/openai_compatible/embedding.rs b/aimux-providers/src/openai_compatible/embedding.rs index 767a674d5..baec1c17e 100644 --- a/aimux-providers/src/openai_compatible/embedding.rs +++ b/aimux-providers/src/openai_compatible/embedding.rs @@ -2,7 +2,7 @@ //! //! The Rust form of `OpenAICompatibleEmbeddingModel`: `POST {base_url}/embeddings` //! with `encoding_format: "float"`, the `dimensions` and `user` options read -//! from the `openaiCompatible` namespace and then the provider's own. +//! from the compatible namespaces and then the provider's own. use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -12,6 +12,8 @@ use aimux_core::embedding_model::{ }; use aimux_core::error::AiMuxError; +use aimux_core::shared::Warning; + use super::config::CompatModelConfig; use super::convert::to_camel_case; @@ -27,13 +29,14 @@ impl OpenAICompatibleEmbeddingModel { } } -/// `dimensions` and `user` from the generic namespace, then the provider's -/// own (later wins). -fn embedding_options(options: &EmbeddingCallOptions, name: &str) -> (Option, Option) { - let camel = to_camel_case(name); - let mut dimensions = None; - let mut user = None; - for key in ["openaiCompatible", name, camel.as_str()] { +/// Validate each namespace before merging, as upstream's provider-options schema does. +fn embedding_options( + options: &EmbeddingCallOptions, + name: &str, + warnings: &mut Vec, +) -> Result, AiMuxError> { + let mut merged = Map::new(); + for key in ["openai-compatible", "openaiCompatible", name] { let Some(object) = options .provider_options .as_ref() @@ -41,14 +44,43 @@ fn embedding_options(options: &EmbeddingCallOptions, name: &str) -> (Option else { continue; }; - if let Some(value) = object.get("dimensions").and_then(Value::as_u64) { - dimensions = Some(value); - } - if let Some(value) = object.get("user").and_then(Value::as_str) { - user = Some(value.to_string()); + for (setting, valid) in [ + ("dimensions", Value::is_number as fn(&Value) -> bool), + ("user", Value::is_string), + ] { + if let Some(value) = object.get(setting) { + if !valid(value) { + return Err(AiMuxError::InvalidArgument(format!( + "invalid providerOptions.{key}.{setting}" + ))); + } + merged.insert(setting.to_string(), value.clone()); + } } } - (dimensions, user) + if options + .provider_options + .as_ref() + .is_some_and(|all| all.contains_key("openai-compatible")) + { + warnings.push(Warning::Deprecated { + setting: "providerOptions key 'openai-compatible'".into(), + message: "Use 'openaiCompatible' instead.".into(), + }); + } + let camel = to_camel_case(name); + if camel != name + && options + .provider_options + .as_ref() + .is_some_and(|all| all.contains_key(name)) + { + warnings.push(Warning::Deprecated { + setting: format!("providerOptions key '{name}'"), + message: format!("Use '{camel}' instead."), + }); + } + Ok(merged) } #[async_trait] @@ -73,18 +105,23 @@ impl EmbeddingModel for OpenAICompatibleEmbeddingModel { &self, options: &EmbeddingCallOptions, ) -> Result { - let (dimensions, user) = embedding_options(options, self.config.provider_options_name()); + let mut warnings = Vec::new(); + let compatible_options = + embedding_options(options, self.config.provider_options_name(), &mut warnings)?; + if options.values.len() > self.max_embeddings_per_call().unwrap_or(2048) as usize { + return Err(AiMuxError::InvalidArgument(format!( + "Too many embedding values for {} / {}: maximum 2048, received {}", + self.provider(), + self.model_id(), + options.values.len() + ))); + } let mut body = Map::new(); body.insert("model".to_string(), json!(self.model_id)); body.insert("input".to_string(), json!(options.values)); body.insert("encoding_format".to_string(), json!("float")); - if let Some(dimensions) = dimensions { - body.insert("dimensions".to_string(), json!(dimensions)); - } - if let Some(user) = user { - body.insert("user".to_string(), json!(user)); - } + body.extend(compatible_options); let body = Value::Object(body); let headers = self @@ -131,12 +168,16 @@ impl EmbeddingModel for OpenAICompatibleEmbeddingModel { Ok(EmbeddingResult { embeddings, usage, - provider_metadata: None, + provider_metadata: raw_value + .get("providerMetadata") + .map(|metadata| serde_json::from_value(metadata.clone())) + .transpose() + .map_err(|error| AiMuxError::InvalidResponseData(error.to_string()))?, response: Some(EmbeddingResponse { headers: Some(response_headers), body: Some(raw_value), }), - warnings: Vec::new(), + warnings, }) } } diff --git a/aimux-providers/src/openai_compatible/image.rs b/aimux-providers/src/openai_compatible/image.rs index 008b99cb6..4b80e33bb 100644 --- a/aimux-providers/src/openai_compatible/image.rs +++ b/aimux-providers/src/openai_compatible/image.rs @@ -2,12 +2,10 @@ //! //! The Rust form of `OpenAICompatibleImageModel`: //! -//! - generation: `POST {base_url}/images/generations`, JSON body with -//! `response_format: "b64_json"`; -//! - editing (`files` given): `POST {base_url}/images/edits`, multipart form. +//! - generation: `POST {base_url}/images/generations`, JSON body; +//! - editing (nonempty `files`): `POST {base_url}/images/edits`, multipart form. //! -//! The provider options of the `openaiCompatible` namespace, then the -//! provider's own, are forwarded to the endpoint as given. +//! Provider-name options are forwarded to the endpoint as given. use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -15,7 +13,7 @@ use serde_json::{Map, Value, json}; use aimux_core::error::AiMuxError; use aimux_core::image_model::{ ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, ImageResponse, - ImageResult, + ImageResult, ImageUsage, }; use aimux_core::shared::Warning; use aimux_provider_utils::{HttpBody, MultipartForm}; @@ -34,13 +32,22 @@ impl OpenAICompatibleImageModel { Self { model_id, config } } - /// The provider options that go to the endpoint: generic namespace, then - /// the provider's own, then its camelCase form (later wins). - fn forwarded_options(&self, options: &ImageCallOptions) -> Map { + /// The provider's own namespace, then its camelCase form (later wins). + fn forwarded_options( + &self, + options: &ImageCallOptions, + warnings: &mut Vec, + ) -> Map { let name = self.config.provider_options_name(); let camel = to_camel_case(name); let mut out = Map::new(); - for key in ["openaiCompatible", name, camel.as_str()] { + if camel != name && options.provider_options.contains_key(name) { + warnings.push(Warning::Deprecated { + setting: format!("providerOptions key '{name}'"), + message: format!("Use '{camel}' instead."), + }); + } + for key in [name, camel.as_str()] { if let Some(object) = options.provider_options.get(key) { out.extend(object.iter().map(|(k, v)| (k.clone(), v.clone()))); } @@ -49,7 +56,7 @@ impl OpenAICompatibleImageModel { } } -fn file_bytes(file: &ImageFile) -> Result<(Vec, String), AiMuxError> { +async fn file_bytes(file: &ImageFile) -> Result<(Vec, String), AiMuxError> { match file { ImageFile::File { media_type, data } => { let bytes = match data { @@ -61,9 +68,26 @@ fn file_bytes(file: &ImageFile) -> Result<(Vec, String), AiMuxError> { }; Ok((bytes, media_type.clone())) } - ImageFile::Url { .. } => Err(AiMuxError::InvalidArgument( - "image edit files must be inline data, not URLs".to_string(), - )), + ImageFile::Url { url } => { + // Upstream downloads file URLs with the default transport, without provider headers. + let response = aimux_provider_utils::get_from_api( + aimux_provider_utils::HttpRequest { + url: url.clone(), + validate_url: true, + ..Default::default() + }, + aimux_provider_utils::create_binary_response_handler(), + aimux_provider_utils::create_status_code_error_response_handler(), + ) + .await?; + let media_type = response + .response_headers + .get("content-type") + .cloned() + .filter(|media_type| !media_type.is_empty()) + .unwrap_or_else(|| "application/octet-stream".into()); + Ok((response.value.to_vec(), media_type)) + } } } @@ -110,63 +134,63 @@ impl ImageModel for OpenAICompatibleImageModel { .config .request_headers(options.headers.as_ref()) .await?; - let forwarded = self.forwarded_options(options); - - let (value, response_headers): (Value, _) = if let Some(files) = &options.files { - let mut form = MultipartForm::new(); - form.text("model", &self.model_id)?; - if let Some(prompt) = &options.prompt { - form.text("prompt", prompt)?; - } - form.text("n", &options.n.to_string())?; - if let Some(size) = options.size { - form.text("size", &size.to_string())?; - } - for (name, value) in &forwarded { - form.text(name, &form_value(value))?; - } - let field = if files.len() == 1 { "image" } else { "image[]" }; - for file in files { - let (bytes, media_type) = file_bytes(file)?; - form.file(field, "image", &media_type, &bytes)?; - } - if let Some(mask) = &options.mask { - let (bytes, media_type) = file_bytes(mask)?; - form.file("mask", "mask", &media_type, &bytes)?; - } - let (bytes, content_type) = form.finish(); - let resp = aimux_provider_utils::post_to_api( - self.config - .http_request("/images/edits", headers, options)?, - HttpBody::Bytes(bytes, content_type), - aimux_provider_utils::create_json_response_handler(), - self.config.failed_response_handler(), - ) - .await?; - (resp.value, resp.response_headers) - } else { - let mut body = Map::new(); - body.insert("model".into(), json!(self.model_id)); - if let Some(prompt) = &options.prompt { - body.insert("prompt".into(), json!(prompt)); - } - body.insert("n".into(), json!(options.n)); - if let Some(size) = options.size { - body.insert("size".into(), json!(size.to_string())); - } - body.extend(forwarded); - body.insert("response_format".into(), json!("b64_json")); - let body = Value::Object(body); - let resp = aimux_provider_utils::post_json_to_api( - self.config - .http_request("/images/generations", headers, options)?, - body, - aimux_provider_utils::create_json_response_handler(), - self.config.failed_response_handler(), - ) - .await?; - (resp.value, resp.response_headers) - }; + let forwarded = self.forwarded_options(options, &mut warnings); + + let (value, response_headers): (Value, _) = + if let Some(files) = options.files.as_ref().filter(|files| !files.is_empty()) { + let mut form = MultipartForm::new(); + form.text("model", &self.model_id)?; + if let Some(prompt) = &options.prompt { + form.text("prompt", prompt)?; + } + form.text("n", &options.n.to_string())?; + if let Some(size) = options.size { + form.text("size", &size.to_string())?; + } + for (name, value) in &forwarded { + form.text(name, &form_value(value))?; + } + let field = if files.len() == 1 { "image" } else { "image[]" }; + for file in files { + let (bytes, media_type) = file_bytes(file).await?; + form.file(field, "image", &media_type, &bytes)?; + } + if let Some(mask) = &options.mask { + let (bytes, media_type) = file_bytes(mask).await?; + form.file("mask", "mask", &media_type, &bytes)?; + } + let (bytes, content_type) = form.finish(); + let resp = aimux_provider_utils::post_to_api( + self.config + .http_request("/images/edits", headers, options)?, + HttpBody::Bytes(bytes, content_type), + aimux_provider_utils::create_json_response_handler(), + self.config.failed_response_handler(), + ) + .await?; + (resp.value, resp.response_headers) + } else { + let mut body = Map::new(); + body.insert("model".into(), json!(self.model_id)); + if let Some(prompt) = &options.prompt { + body.insert("prompt".into(), json!(prompt)); + } + body.insert("n".into(), json!(options.n)); + if let Some(size) = options.size { + body.insert("size".into(), json!(size.to_string())); + } + body.extend(forwarded); + let body = Value::Object(body); + let resp = aimux_provider_utils::post_json_to_api( + self.config + .http_request("/images/generations", headers, options)?, + body, + aimux_provider_utils::create_json_response_handler(), + self.config.failed_response_handler(), + ) + .await?; + (resp.value, resp.response_headers) + }; let images: Vec = value .get("data") @@ -189,7 +213,23 @@ impl ImageModel for OpenAICompatibleImageModel { model_id: Some(self.model_id.clone()), headers: Some(response_headers), }, - usage: None, + usage: value + .get("usage") + .filter(|usage| !usage.is_null()) + .map(|usage| ImageUsage { + input_tokens: usage + .get("input_tokens") + .and_then(Value::as_u64) + .map(|tokens| tokens as u32), + output_tokens: usage + .get("output_tokens") + .and_then(Value::as_u64) + .map(|tokens| tokens as u32), + total_tokens: usage + .get("total_tokens") + .and_then(Value::as_u64) + .map(|tokens| tokens as u32), + }), }) } } diff --git a/aimux-providers/src/openai_compatible/mod.rs b/aimux-providers/src/openai_compatible/mod.rs index 285f3f16b..cf6ec7a52 100644 --- a/aimux-providers/src/openai_compatible/mod.rs +++ b/aimux-providers/src/openai_compatible/mod.rs @@ -2,7 +2,7 @@ //! //! [`create_openai_compatible`] is the Rust form of the AI SDK's //! `createOpenAICompatible`: it takes [`OpenAICompatibleProviderSettings`], -//! validates the name and base URL and returns an [`OpenAICompatibleProvider`] +//! validates the base URL and returns an [`OpenAICompatibleProvider`] //! whose models speak the chat-completions, embeddings and image endpoints of //! any OpenAI-compatible server. //! @@ -11,8 +11,8 @@ //! - The explicit settings (`name`, `base_url`, `headers`, `query_params`, //! `fetch`, the capability flags, `transform_request_body`) are fixed in the //! factory. -//! - A non-empty `api_key` is used as given. `None` or `""` sends no -//! `Authorization` header, which is what a local server wants. +//! - A non-empty `api_key` produces a bearer header. `None` or `""` sends +//! no `Authorization` header, which is what a local server wants. //! //! Identity: every model reports `"{name}.{method}"` (`groq.chat`, //! `local.embedding`, ...). The providerOptions namespace is the same name: @@ -49,10 +49,12 @@ use futures::future::BoxFuture; use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; use aimux_core::image_model::ImageModel; -use aimux_core::language_model::LanguageModel; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; use aimux_core::model_catalogue::RuntimeModel; use aimux_core::provider::{Provider, ProviderDiscovery}; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{ + FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, combine_headers, validate_base_url, +}; use crate::shared::{Credential, provider_headers}; use config::{BaseUrl, ChatDialect, ChatSettings, CompatModelConfig}; @@ -62,7 +64,7 @@ use config::{BaseUrl, ChatDialect, ChatSettings, CompatModelConfig}; #[derive(Clone, Default)] pub struct OpenAICompatibleProviderSettings { /// The provider name: the first segment of every model's `provider()` - /// string and the providerOptions namespace. Required, non-empty, no `.`. + /// string and the providerOptions namespace. Preserved as supplied. pub name: String, /// Base URL for the API calls. Required, `http(s)` with a host; a trailing /// slash is removed. @@ -83,14 +85,14 @@ pub struct OpenAICompatibleProviderSettings { /// Whether the chat endpoint accepts `json_schema` response formats; when /// not, a schema degrades to `json_object` with a warning. pub supports_structured_outputs: Option, - /// Rewrites chat request bodies once, after they are serialized and + /// Rewrites each chat JSON request body once, after it is serialized and /// before it is sent. pub transform_request_body: Option, - /// Extract provider-specific metadata from generated and streamed responses. + /// Extracts metadata from chat responses and streaming chunks. pub metadata_extractor: Option>, - /// Returns the supported URLs for chat models. + /// URL patterns supported by chat models, evaluated when requested. pub supported_urls: Option, - /// Converts chat response token usage. + /// Custom chat token accounting. pub convert_usage: Option, } @@ -127,22 +129,25 @@ impl std::fmt::Debug for OpenAICompatibleProviderSettings { /// /// # Errors /// -/// Returns `AiMuxError::InvalidArgument` when `name` is empty or contains `.`, -/// or when `base_url` is not an `http(s)` URL with a host. The key is not read +/// Returns `AiMuxError::InvalidArgument` when `base_url` is not an `http(s)` +/// URL with a host. The key is not read /// here. pub fn create_openai_compatible( settings: OpenAICompatibleProviderSettings, ) -> Result { let base_url = validate_base_url(&settings.base_url)?; - let mut dialect = ChatDialect::baseline(); - dialect.metadata_extractor = settings.metadata_extractor; - OpenAICompatibleProvider::assemble(Assembly { + let credential = settings + .api_key + .clone() + .filter(|key| !key.is_empty()) + .map_or(Credential::None, |key| { + Credential::Explicit(Resolvable::Value(key)) + }); + let user_headers = settings.headers.clone(); + let mut provider = OpenAICompatibleProvider::assemble(Assembly { name: settings.name, base_url: BaseUrl::Fixed(base_url), - credential: match settings.api_key { - Some(key) => Credential::Optional(Resolvable::Value(key)), - None => Credential::None, - }, + credential: credential.clone(), fixed_headers: Vec::new(), headers: settings.headers, query_params: settings.query_params, @@ -151,12 +156,20 @@ pub fn create_openai_compatible( profile: ChatProfile { include_usage: settings.include_usage.unwrap_or(false), supports_structured_outputs: settings.supports_structured_outputs.unwrap_or(false), - supports_multi_part_tool_content: false, - dialect, - supported_urls: settings.supported_urls, - convert_usage: settings.convert_usage, + dialect: ChatDialect::baseline(), }, - }) + })?; + let dialect = Arc::make_mut(&mut provider.chat.dialect); + dialect.metadata_extractor = settings.metadata_extractor; + dialect.supported_urls = settings.supported_urls; + dialect.convert_usage = settings.convert_usage; + let headers = compatible_headers(credential, user_headers); + provider.headers = aimux_provider_utils::headers::with_user_agent_suffix_fn( + headers, + "openai-compatible", + "3.0.59", + ); + Ok(provider) } /// The chat-endpoint behavior of a compatible vendor: the capability flags the @@ -167,10 +180,7 @@ pub fn create_openai_compatible( pub(crate) struct ChatProfile { pub include_usage: bool, pub supports_structured_outputs: bool, - pub supports_multi_part_tool_content: bool, pub dialect: ChatDialect, - pub supported_urls: Option, - pub convert_usage: Option, } /// Everything a compatible provider is made of. The public factory fills it @@ -200,20 +210,26 @@ pub struct OpenAICompatibleProvider { chat: ChatSettings, } +fn compatible_headers(credential: Credential, user: Option) -> HeadersFn { + Resolvable::from_async_fn(move || { + let credential = credential.clone(); + let user = user.clone(); + async move { + let mut headers = HeaderMapOpt::new(); + if let Some(key) = credential.secret().await?.filter(|key| !key.is_empty()) { + headers.insert("Authorization".into(), Some(format!("Bearer {key}"))); + } + Ok(match user { + Some(user) => combine_headers(&[&headers, &user]), + None => headers, + }) + } + }) +} + impl OpenAICompatibleProvider { pub(crate) fn assemble(assembly: Assembly) -> Result { - let name = assembly.name.trim().to_string(); - if name.is_empty() { - return Err(AiMuxError::InvalidArgument( - "OpenAI-compatible provider `name` must not be empty".to_string(), - )); - } - if name.contains('.') { - return Err(AiMuxError::InvalidArgument(format!( - "OpenAI-compatible provider name {name:?} must not contain `.`: it is the first \ - segment of the provider string and the providerOptions namespace" - ))); - } + let name = assembly.name; let query_params = assembly.query_params.map(|params| { let mut pairs: Vec<(String, String)> = params.into_iter().collect(); pairs.sort(); @@ -237,9 +253,7 @@ impl OpenAICompatibleProvider { chat: ChatSettings { include_usage: assembly.profile.include_usage, supports_structured_outputs: assembly.profile.supports_structured_outputs, - supports_multi_part_tool_content: assembly.profile.supports_multi_part_tool_content, - supported_urls: assembly.profile.supported_urls, - convert_usage: assembly.profile.convert_usage, + supported_urls: SupportedUrls::default(), dialect: Arc::new(assembly.profile.dialect), }, }) diff --git a/aimux-providers/src/openai_compatible/types.rs b/aimux-providers/src/openai_compatible/types.rs index 4a6bdfa3c..6e182cebf 100644 --- a/aimux-providers/src/openai_compatible/types.rs +++ b/aimux-providers/src/openai_compatible/types.rs @@ -73,9 +73,6 @@ pub(crate) struct UsageResponse { pub completion_tokens: Option, #[serde(default)] pub total_tokens: Option, - /// Top-level `cached_tokens` (Moonshot); wins over the nested value. - #[serde(default)] - pub cached_tokens: Option, #[serde(default)] pub prompt_tokens_details: Option, #[serde(default)] @@ -86,9 +83,6 @@ pub(crate) struct UsageResponse { pub(crate) struct PromptTokensDetails { #[serde(default)] pub cached_tokens: Option, - /// Alibaba (DashScope) reports cache writes as `cache_creation_input_tokens`. - #[serde(default, alias = "cache_creation_input_tokens")] - pub cache_write_tokens: Option, } #[derive(Debug, Deserialize, Clone, Serialize)] @@ -126,7 +120,7 @@ pub(crate) struct StreamChoice { #[derive(Debug, Deserialize, Default)] pub(crate) struct Delta { #[serde(default)] - pub content: Option, + pub content: Option, #[serde(default)] pub reasoning: Option, #[serde(default)] diff --git a/aimux-providers/src/preset.rs b/aimux-providers/src/preset.rs index e138519b4..2ae87dcc5 100644 --- a/aimux-providers/src/preset.rs +++ b/aimux-providers/src/preset.rs @@ -381,13 +381,39 @@ fn assemble( let mut dialect = ChatDialect::baseline(); dialect.supports_top_k = true; dialect.max_tokens_key = descriptor.max_tokens_key; + if matches!(descriptor.name, "alibaba" | "moonshotai") { + dialect.convert_usage = Some(Arc::new(move |raw| { + let mut usage = crate::openai_compatible::chat::usage_from_raw(raw); + let Some(raw) = raw.filter(|raw| !raw.is_null()) else { + return usage; + }; + let tokens = |value: &serde_json::Value| { + value.as_u64().and_then(|value| u32::try_from(value).ok()) + }; + if descriptor.name == "alibaba" { + let details = &raw["prompt_tokens_details"]; + let cache_write = tokens(&details["cache_creation_input_tokens"]) + .or_else(|| tokens(&details["cache_write_tokens"])) + .unwrap_or(0); + usage.input_tokens.cache_write = Some(cache_write); + usage.input_tokens.no_cache = usage + .input_tokens + .no_cache + .map(|tokens| tokens.saturating_sub(cache_write)); + } else if let Some(cached) = tokens(&raw["cached_tokens"]) { + usage.input_tokens.cache_read = Some(cached); + usage.input_tokens.no_cache = usage + .input_tokens + .total + .map(|tokens| tokens.saturating_sub(cached)); + } + usage + })); + } let profile = ChatProfile { include_usage: true, supports_structured_outputs: true, - supports_multi_part_tool_content: false, dialect, - supported_urls: None, - convert_usage: None, }; OpenAICompatibleProvider::assemble(Assembly { diff --git a/aimux-providers/src/provider.rs b/aimux-providers/src/provider.rs index bcbe94fe0..1a870a102 100644 --- a/aimux-providers/src/provider.rs +++ b/aimux-providers/src/provider.rs @@ -84,10 +84,7 @@ impl ProviderProfile { Ok(ChatProfile { include_usage: true, supports_structured_outputs: true, - supports_multi_part_tool_content: false, dialect, - supported_urls: None, - convert_usage: None, }) } } diff --git a/aimux-providers/src/shared/mod.rs b/aimux-providers/src/shared/mod.rs index 15939030a..3f4917139 100644 --- a/aimux-providers/src/shared/mod.rs +++ b/aimux-providers/src/shared/mod.rs @@ -48,8 +48,6 @@ pub(crate) enum Credential { /// back to the environment. A [`Resolvable::Future`] is awaited once, an /// `AsyncFn` on every request. Explicit(Resolvable), - /// An optional key: an empty resolved value sends no credential. - Optional(Resolvable), /// Read the environment variable on every request; unset fails that /// request with `AiMuxError::LoadApiKey`. Env { var: String, description: String }, @@ -81,7 +79,6 @@ impl Credential { match self { Self::None => Ok(None), Self::Explicit(key) => Ok(Some(key.resolve().await?)), - Self::Optional(key) => Ok(Some(key.resolve().await?).filter(|key| !key.is_empty())), Self::Env { var, description } => Ok(Some(load_api_key(None, var, description)?)), Self::OptionalEnv(var) => Ok(std::env::var(var).ok().filter(|key| !key.is_empty())), } diff --git a/aimux-providers/src/vertex/anthropic_model.rs b/aimux-providers/src/vertex/anthropic_model.rs index 136ba0c0e..5f5511e10 100644 --- a/aimux-providers/src/vertex/anthropic_model.rs +++ b/aimux-providers/src/vertex/anthropic_model.rs @@ -11,7 +11,7 @@ //! - the body drops `model` and gains `anthropic_version: "vertex-2023-10-16"` //! (the version is not a header), //! - credentials and the base URL are the Vertex provider's, resolved on every -//! request (OAuth bearer token or Express API key; project and location), +//! request (OAuth bearer token; project and location), //! - errors are Google-shaped, //! - no URL sources, no structured-output beta and no `strict` tool //! definitions. @@ -29,7 +29,7 @@ use aimux_provider_utils::{FetchFunction, Resolvable}; use crate::anthropic::AnthropicMessagesModel; use crate::anthropic::config::{AnthropicModelConfig, AnthropicModelHooks, TransformRequestBody}; -use crate::anthropic::options::CANONICAL; +use crate::anthropic::options::options_name_of; use crate::shared::Endpoint; /// `anthropic_version` envelope value required by the Vertex AI `rawPredict` / @@ -78,7 +78,7 @@ fn request_config( supported_urls: SupportedUrls::default(), transform_request_body, base_url, - provider_options_name: CANONICAL.to_string(), + provider_options_name: options_name_of(PROVIDER), hooks: AnthropicModelHooks { request_url: Some(Arc::new(move |model_id, stream| { let method = if stream { diff --git a/aimux-providers/src/vertex/anthropic_provider.rs b/aimux-providers/src/vertex/anthropic_provider.rs new file mode 100644 index 000000000..4e731bcd9 --- /dev/null +++ b/aimux-providers/src/vertex/anthropic_provider.rs @@ -0,0 +1,116 @@ +//! Anthropic provider on Vertex AI with application default credentials. + +use std::sync::{Arc, OnceLock}; + +use aimux_core::{ + AiMuxError, embedding_model::EmbeddingModel, image_model::ImageModel, + language_model::LanguageModel, provider::Provider, +}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable}; + +use super::{GoogleAuthOptions, VertexProvider, VertexProviderSettings, create_google_vertex}; + +/// Settings for the Vertex Anthropic provider. +#[derive(Clone, Default)] +pub struct VertexAnthropicProviderSettings { + pub project: Option, + pub location: Option, + pub base_url: Option, + pub headers: Option>, + pub fetch: Option, + pub google_auth_options: Option, + /// Overrides the OAuth token generator. A missing token becomes `Bearer null`. + pub generate_auth_token: Option>>, +} + +impl std::fmt::Debug for VertexAnthropicProviderSettings { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VertexAnthropicProviderSettings") + .field("project", &self.project) + .field("location", &self.location) + .field("base_url", &self.base_url) + .field("headers", &self.headers.is_some()) + .field("fetch", &self.fetch.is_some()) + .field("google_auth_options", &self.google_auth_options) + .field("generate_auth_token", &self.generate_auth_token.is_some()) + .finish() + } +} + +/// Create a Vertex Anthropic provider. +/// +/// # Errors +/// +/// Credentials are resolved when a model makes a request. +pub fn create_google_vertex_anthropic( + settings: VertexAnthropicProviderSettings, +) -> Result { + let mut auth = settings.google_auth_options.unwrap_or_default(); + if let Some(generator) = settings.generate_auth_token { + auth = GoogleAuthOptions { + auth_client: Some(generator), + ..Default::default() + }; + } + let mut provider = create_google_vertex(VertexProviderSettings { + project: settings.project, + location: settings.location, + base_url: settings.base_url, + headers: settings.headers, + fetch: settings.fetch, + google_auth_options: Some(auth.clone()), + ..Default::default() + })?; + let resolver = Arc::get_mut(&mut provider.resolver).expect("new provider resolver is unique"); + resolver.use_access_token_env = false; + resolver.auth = super::auth::GoogleAuth::new(auth, provider.fetch.clone()); + Ok(VertexAnthropicProvider(provider)) +} + +/// Default Vertex Anthropic provider. +pub fn google_vertex_anthropic() -> &'static VertexAnthropicProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_google_vertex_anthropic(VertexAnthropicProviderSettings::default()) + .expect("default Vertex Anthropic settings are valid") + }) +} + +/// The Vertex Anthropic model factory. +pub struct VertexAnthropicProvider(VertexProvider); + +impl VertexAnthropicProvider { + #[must_use] + pub fn call(&self, model_id: &str) -> Arc { + let resolver = self.0.resolver.clone(); + Arc::new(super::anthropic_model::model( + model_id, + Arc::new(move || { + let resolver = resolver.clone(); + Box::pin(async move { resolver.endpoint(super::Publisher::Anthropic).await }) + }), + self.0.fetch.clone(), + self.0.transform_request_body.clone(), + )) + } +} + +impl Provider for VertexAnthropicProvider { + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) + } + + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::NoSuchModel { + model_id: model_id.into(), + model_type: "embeddingModel".into(), + }) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::NoSuchModel { + model_id: model_id.into(), + model_type: "imageModel".into(), + }) + } +} diff --git a/aimux-providers/src/vertex/auth.rs b/aimux-providers/src/vertex/auth.rs new file mode 100644 index 000000000..787c638a6 --- /dev/null +++ b/aimux-providers/src/vertex/auth.rs @@ -0,0 +1,1090 @@ +//! Application default Google credentials and OAuth access-token generation. + +use std::path::PathBuf; +use std::time::{Duration, Instant}; + +use aimux_core::error::AiMuxError; +use aimux_provider_utils::{ + FetchFunction, HttpBody, HttpRequest, Resolvable, create_binary_response_handler, + create_json_response_handler, get_from_api, post_to_api, +}; +use base64::Engine; +use ring::{ + rand::SystemRandom, + signature::{RSA_PKCS1_SHA256, RsaKeyPair}, +}; +use serde::Deserialize; +use serde_json::{Value, json}; +use tokio::sync::Mutex; + +const TOKEN_URL: &str = "https://oauth2.googleapis.com/token"; +const CLOUD_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; + +/// OAuth scopes, as a space-delimited string or an array. +#[derive(Clone, Debug, serde::Serialize, Deserialize)] +#[serde(untagged)] +pub enum GoogleAuthScopes { + String(String), + Array(Vec), +} + +impl GoogleAuthScopes { + fn array(&self) -> Vec { + match self { + Self::String(value) if value.is_empty() => Vec::new(), + Self::String(value) => vec![value.clone()], + Self::Array(values) => values.clone(), + } + } + + fn joined(&self) -> String { + match self { + Self::String(value) => value.clone(), + Self::Array(values) => values.join(" "), + } + } +} + +/// Options passed to the Google credential resolver. +#[derive(Clone, Default)] +pub struct GoogleAuthOptions { + /// Inline service-account or authorized-user credentials. + pub credentials: Option, + /// Credential JSON filename (`keyFile` in the upstream library). + pub key_file: Option, + /// Alias of `key_file` (`keyFilename` upstream). + pub key_filename: Option, + /// OAuth scopes. Defaults to the Cloud Platform scope. + pub scopes: Option, + /// Project associated with the credentials, independent of the API endpoint project. + pub project_id: Option, + /// Overrides the access-token client. + pub auth_client: Option>>, + /// Options passed to the OAuth client. Unsupported options fail explicitly. + pub client_options: Option, + /// Service domain of the Google Cloud universe. + pub universe_domain: Option, + /// Google auth-library API-key mode (not a Vertex Express-mode key). + pub api_key: Option, +} + +impl std::fmt::Debug for GoogleAuthOptions { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GoogleAuthOptions") + .field("credentials", &self.credentials.is_some()) + .field("key_file", &self.key_file) + .field("key_filename", &self.key_filename) + .field("scopes", &self.scopes) + .field("project_id", &self.project_id) + .field("auth_client", &self.auth_client.is_some()) + .field("client_options", &self.client_options.is_some()) + .field("universe_domain", &self.universe_domain) + .field("api_key", &self.api_key.is_some()) + .finish() + } +} + +#[derive(Deserialize)] +struct TokenResponse { + access_token: String, + expires_in: Option, + refresh_token: Option, +} + +struct CachedToken { + token: String, + expiry: Option, + credentials: Option, +} + +pub(super) struct GoogleAuth { + options: GoogleAuthOptions, + token: Mutex>, + fetch: Option, +} + +impl GoogleAuth { + pub(super) fn new(options: GoogleAuthOptions, fetch: Option) -> Self { + Self { + options, + token: Mutex::new(None), + fetch, + } + } + + fn credentials(&self) -> Result, AiMuxError> { + if let Some(credentials) = &self.options.credentials { + return Ok(Some(credentials.clone())); + } + let explicit = self + .options + .key_filename + .clone() + .filter(|path| !path.as_os_str().is_empty()) + .or_else(|| { + self.options + .key_file + .clone() + .filter(|path| !path.as_os_str().is_empty()) + }) + .or_else(|| { + std::env::var_os("GOOGLE_APPLICATION_CREDENTIALS") + .filter(|value| !value.is_empty()) + .or_else(|| { + std::env::var_os("google_application_credentials") + .filter(|value| !value.is_empty()) + }) + .map(PathBuf::from) + }); + let path = explicit.or_else(|| { + let directory = std::env::var_os("CLOUDSDK_CONFIG") + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + .or_else(|| { + if cfg!(windows) { + std::env::var_os("APPDATA") + .filter(|value| !value.is_empty()) + .map(|directory| PathBuf::from(directory).join("gcloud")) + } else { + std::env::var_os("HOME") + .filter(|value| !value.is_empty()) + .map(|home| PathBuf::from(home).join(".config/gcloud")) + } + })?; + let path = directory.join("application_default_credentials.json"); + path.exists().then_some(path) + }); + let Some(path) = path else { + return Ok(None); + }; + let data = std::fs::read(path).map_err(|error| { + AiMuxError::InvalidArgument(format!( + "Could not read Google application credentials: {error}" + )) + })?; + serde_json::from_slice(&data).map(Some).map_err(|error| { + AiMuxError::InvalidArgument(format!("Invalid Google application credentials: {error}")) + }) + } + + pub(super) async fn access_token(&self) -> Result { + if let Some(client) = &self.options.auth_client { + return Ok(client.resolve().await?.unwrap_or_else(|| "null".into())); + } + self.validate_options()?; + let mut cache = self.token.lock().await; + if let Some(cached) = cache.as_ref() + && cached.expiry.is_none_or(|expiry| expiry > Instant::now()) + { + return Ok(cached.token.clone()); + } + let mut credentials = match cache.as_ref() { + Some(cached) => cached.credentials.clone(), + None => self.credentials()?, + }; + if cache.is_none() + && let Some(credential) = credentials.as_ref() + && credential.get("type").and_then(Value::as_str) == Some("authorized_user") + { + for field in ["client_id", "client_secret", "refresh_token"] { + required(credential, field)?; + } + if let Some(initial) = self + .options + .client_options + .as_ref() + .and_then(|options| options.get("credentials")) + && let Some(token) = initial + .get("access_token") + .and_then(Value::as_str) + .filter(|token| !token.is_empty()) + { + let threshold = self + .options + .client_options + .as_ref() + .and_then(|options| options.get("eagerRefreshThresholdMillis")) + .and_then(Value::as_i64) + .filter(|value| *value != 0) + .unwrap_or(300_000); + let remaining = initial + .get("expiry_date") + .and_then(Value::as_i64) + .map(|expiry| { + expiry + .saturating_sub(chrono::Utc::now().timestamp_millis()) + .saturating_sub(threshold) + }); + if remaining.is_none_or(|remaining| remaining > 0) { + let expiry = remaining + .map(|remaining| { + Instant::now() + .checked_add(Duration::from_millis(remaining as u64)) + .ok_or_else(|| { + AiMuxError::InvalidResponseData( + "Google OAuth token expiry is out of range".into(), + ) + }) + }) + .transpose()?; + let token = token.to_string(); + *cache = Some(CachedToken { + token: token.clone(), + expiry, + credentials, + }); + return Ok(token); + } + } + } + let response = self.credential_token(credentials.clone()).await?; + if let Some(token) = &response.refresh_token + && let Some(credentials) = credentials.as_mut() + && credentials.get("type").and_then(Value::as_str) + == Some("external_account_authorized_user") + { + credentials["refresh_token"] = json!(token); + } + let external = credentials + .as_ref() + .and_then(|value| value.get("type")) + .and_then(Value::as_str) + .is_some_and(|kind| { + matches!( + kind, + "external_account" | "external_account_authorized_user" + ) + }); + if !external && response.access_token.is_empty() { + return Err(AiMuxError::InvalidResponseData( + "Could not refresh Google access token".into(), + )); + } + let threshold = self + .options + .client_options + .as_ref() + .and_then(|options| options.get("eagerRefreshThresholdMillis")) + .and_then(Value::as_u64) + .filter(|threshold| external || *threshold != 0) + .unwrap_or(300_000); + let expiry = response + .expires_in + .map(|seconds| { + Instant::now() + .checked_add( + Duration::from_secs(seconds) + .saturating_sub(Duration::from_millis(threshold)), + ) + .ok_or_else(|| { + AiMuxError::InvalidResponseData( + "Google OAuth token expiry is out of range".into(), + ) + }) + }) + .transpose()?; + *cache = Some(CachedToken { + token: response.access_token.clone(), + expiry, + credentials, + }); + Ok(response.access_token) + } + + fn credential_token( + &self, + credentials: Option, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { + Ok(match credentials { + Some(credentials) + if credentials.get("type").and_then(Value::as_str) + == Some("authorized_user") => + { + self.exchange( + vec![ + ("grant_type", "refresh_token".into()), + ("client_id", required(&credentials, "client_id")?.into()), + ( + "client_secret", + required(&credentials, "client_secret")?.into(), + ), + ( + "refresh_token", + required(&credentials, "refresh_token")?.into(), + ), + ], + TOKEN_URL, + ) + .await? + } + Some(credentials) + if credentials + .get("type") + .and_then(Value::as_str) + .is_none_or(|kind| kind == "service_account") => + { + let email = required(&credentials, "client_email")?; + let key = required(&credentials, "private_key")?; + let audience = TOKEN_URL; + let scopes = self + .options + .scopes + .as_ref() + .map_or_else(|| CLOUD_SCOPE.to_string(), GoogleAuthScopes::joined); + let timestamp = chrono::Utc::now().timestamp(); + let claims = json!({"iss": email, "scope": scopes, "aud": audience, "iat": timestamp, "exp": timestamp + 3600}); + let mut header = json!({"alg": "RS256", "typ": "JWT"}); + if let Some(id) = credentials.get("private_key_id").and_then(Value::as_str) { + header["kid"] = json!(id); + } + let encoding = &base64::engine::general_purpose::URL_SAFE_NO_PAD; + let unsigned = format!( + "{}.{}", + encoding.encode(header.to_string()), + encoding.encode(claims.to_string()) + ); + let der = base64::engine::general_purpose::STANDARD + .decode( + key.lines() + .filter(|line| !line.starts_with("-----")) + .collect::(), + ) + .map_err(|_| { + AiMuxError::InvalidArgument( + "Invalid Google service-account private key encoding".into(), + ) + })?; + let key_pair = RsaKeyPair::from_pkcs8(&der) + .or_else(|_| RsaKeyPair::from_der(&der)) + .map_err(|_| { + AiMuxError::InvalidArgument( + "Invalid Google service-account RSA private key".into(), + ) + })?; + let mut signature = vec![0; key_pair.public().modulus_len()]; + key_pair + .sign( + &RSA_PKCS1_SHA256, + &SystemRandom::new(), + unsigned.as_bytes(), + &mut signature, + ) + .map_err(|_| { + AiMuxError::InvalidArgument( + "Could not sign Google service-account assertion".into(), + ) + })?; + let assertion = format!("{unsigned}.{}", encoding.encode(signature)); + self.exchange( + vec![ + ( + "grant_type", + "urn:ietf:params:oauth:grant-type:jwt-bearer".into(), + ), + ("assertion", assertion), + ], + audience, + ) + .await? + } + Some(credentials) + if credentials.get("type").and_then(Value::as_str) + == Some("external_account_authorized_user") => + { + self.exchange_headers( + vec![ + ("grant_type", "refresh_token".into()), + ( + "refresh_token", + required(&credentials, "refresh_token")?.into(), + ), + ], + credentials + .get("token_url") + .and_then(Value::as_str) + .unwrap_or(&format!( + "https://sts.{}/v1/oauthtoken", + self.universe(&credentials) + )), + basic_auth(&credentials)?, + ) + .await? + } + Some(credentials) + if credentials.get("type").and_then(Value::as_str) + == Some("external_account") => + { + self.external_token(&normalize_external_credentials(credentials)) + .await? + } + Some(credentials) + if credentials.get("type").and_then(Value::as_str) + == Some("impersonated_service_account") => + { + let source = credentials.get("source_credentials").ok_or_else(|| { + AiMuxError::InvalidArgument( + "Google impersonation credentials require source_credentials".into(), + ) + })?; + let token = self.credential_token(Some(source.clone())).await?; + self.impersonate( + &credentials, + required(&credentials, "service_account_impersonation_url")?, + &token.access_token, + ) + .await? + } + Some(credentials) + if credentials.get("type").and_then(Value::as_str) + == Some("gdch_service_account") => + { + self.gdch_token(&credentials).await? + } + Some(credentials) => { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "Google credential type '{}' is not supported", + credentials + .get("type") + .and_then(Value::as_str) + .unwrap_or("unknown") + ))); + } + None => { + let host = std::env::var("GCE_METADATA_IP") + .ok() + .filter(|value| !value.is_empty()) + .or_else(|| { + std::env::var("GCE_METADATA_HOST") + .ok() + .filter(|value| !value.is_empty()) + }) + .unwrap_or_else(|| "http://169.254.169.254".into()); + let host = if host.starts_with("http://") || host.starts_with("https://") { + host + } else { + format!("http://{host}") + }; + let mut url = url::Url::parse(&host).map_err(|error| { + AiMuxError::InvalidArgument(format!( + "Invalid Google metadata host: {error}" + )) + })?; + url.set_path("/computeMetadata/v1/instance/service-accounts/default/token"); + url.set_query(None); + let scopes = self + .options + .scopes + .as_ref() + .map_or_else(|| vec![CLOUD_SCOPE.into()], GoogleAuthScopes::array) + .join(","); + if !scopes.is_empty() { + url.query_pairs_mut().append_pair("scopes", &scopes); + } + let response = get_from_api( + HttpRequest { + url: url.into(), + headers: vec![("Metadata-Flavor".into(), "Google".into())], + response_timeout: Some(Duration::from_secs(3)), + fetch: self.fetch.clone(), + ..Default::default() + }, + create_json_response_handler::(), + crate::google::google_failed_response_handler(), + ) + .await?; + if !response.response_headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case("Metadata-Flavor") && value == "Google" + }) { + return Err(AiMuxError::InvalidResponseData( + "Invalid Google metadata response Metadata-Flavor header".into(), + )); + } + response.value + } + }) + }) + } + + async fn gdch_token(&self, credentials: &Value) -> Result { + if required(credentials, "format_version")? != "1" { + return Err(AiMuxError::InvalidArgument( + "Google GDCH credentials require format_version 1".into(), + )); + } + let client = self.options.client_options.as_ref(); + let audience = client.and_then(|options| options.get("apiAudience")).and_then(Value::as_str).filter(|value| !value.is_empty()).ok_or_else(|| AiMuxError::InvalidArgument("Audience cannot be null or empty for GDCH service account credentials; provide client_options.apiAudience".into()))?; + let fetch = match ( + self.fetch.clone(), + credentials + .get("ca_cert_path") + .and_then(Value::as_str) + .filter(|path| !path.is_empty()), + ) { + (None, Some(path)) => Some(super::auth_certificate::ca_fetch(path)?), + (fetch, _) => fetch, + }; + let token_url = required(credentials, "token_uri")?; + let identity = format!( + "system:serviceaccount:{}:{}", + required(credentials, "project")?, + required(credentials, "name")? + ); + let now = chrono::Utc::now().timestamp(); + let lifetime = client + .and_then(|options| options.get("lifetime")) + .and_then(Value::as_i64) + .filter(|value| *value != 0) + .unwrap_or(3600); + let expiry = now.checked_add(lifetime).ok_or_else(|| { + AiMuxError::InvalidArgument("Invalid Google GDCH token lifetime".into()) + })?; + let header = + json!({"alg": "ES256", "typ": "JWT", "kid": required(credentials, "private_key_id")?}); + let claims = + json!({"iss": identity, "sub": identity, "iat": now, "exp": expiry, "aud": token_url}); + let encoding = &base64::engine::general_purpose::URL_SAFE_NO_PAD; + let unsigned = format!( + "{}.{}", + encoding.encode(header.to_string()), + encoding.encode(claims.to_string()) + ); + let key = required(credentials, "private_key")?; + let mut der = base64::engine::general_purpose::STANDARD + .decode( + key.lines() + .filter(|line| !line.starts_with("-----")) + .collect::(), + ) + .map_err(|_| { + AiMuxError::InvalidArgument("Invalid Google GDCH private key encoding".into()) + })?; + if key.contains("BEGIN EC PRIVATE KEY") { + let mut wrapped = vec![ + 0x02, 0x01, 0x00, 0x30, 0x13, 0x06, 0x07, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02, 0x01, + 0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07, + ]; + wrapped.extend(der_wrap(0x04, der)); + der = der_wrap(0x30, wrapped); + } + let random = SystemRandom::new(); + let pair = ring::signature::EcdsaKeyPair::from_pkcs8( + &ring::signature::ECDSA_P256_SHA256_FIXED_SIGNING, + &der, + &random, + ) + .map_err(|_| AiMuxError::InvalidArgument("Invalid Google GDCH P-256 private key".into()))?; + let signature = pair.sign(&random, unsigned.as_bytes()).map_err(|_| { + AiMuxError::InvalidArgument("Could not sign Google GDCH assertion".into()) + })?; + let assertion = format!("{unsigned}.{}", encoding.encode(signature.as_ref())); + let response = post_to_api(HttpRequest { url: token_url.into(), response_timeout: Some(Duration::from_secs(10)), fetch, ..Default::default() }, HttpBody::Json(json!({"audience": audience, "grant_type": "urn:ietf:params:oauth:token-type:token-exchange", "requested_token_type": "urn:ietf:params:oauth:token-type:access_token", "subject_token": assertion, "subject_token_type": "urn:k8s:params:oauth:token-type:serviceaccount"})), create_json_response_handler::(), crate::google::google_failed_response_handler()).await?.value; + if response.access_token.is_empty() + || response.expires_in.is_none_or(|seconds| seconds == 0) + { + return Err(AiMuxError::InvalidResponseData( + "Google GDCH token response requires access_token and expires_in".into(), + )); + } + Ok(response) + } + + async fn external_token(&self, credentials: &Value) -> Result { + let source = credentials.get("credential_source").ok_or_else(|| { + AiMuxError::InvalidArgument( + "Google external credentials require credential_source".into(), + ) + })?; + let file = source + .get("file") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()); + let url = source + .get("url") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()); + let certificate = source.get("certificate").filter(|value| !value.is_null()); + if (file.is_some() && url.is_some()) + || (url.is_some() && certificate.is_some()) + || (file.is_some() && certificate.is_some()) + { + return Err(AiMuxError::InvalidArgument( + "Google identity pool credentials require exactly one of file, URL or certificate" + .into(), + )); + } + let parse_format = source.get("environment_id").is_none() + && source.get("executable").is_none() + && certificate.is_none(); + let mut fetch = self.fetch.clone(); + let token = if let Some(certificate) = certificate { + let (token, certificate_fetch) = + super::auth_certificate::certificate(certificate, fetch)?; + fetch = certificate_fetch; + token + } else if source.get("environment_id").is_some() { + super::auth_external::aws(credentials, self.fetch.clone()).await? + } else if source.get("executable").is_some() { + super::auth_external::executable(credentials).await? + } else if let Some(file) = file { + std::fs::read_to_string(file).map_err(|error| { + AiMuxError::InvalidArgument(format!("Could not read Google subject token: {error}")) + })? + } else if let Some(url) = url { + let headers = source + .get("headers") + .and_then(Value::as_object) + .map(|values| { + values + .iter() + .map(|(name, value)| { + value + .as_str() + .map(|value| (name.clone(), value.to_string())) + .ok_or_else(|| { + AiMuxError::InvalidArgument( + "Google credential_source headers must be strings".into(), + ) + }) + }) + .collect::, _>>() + }) + .transpose()? + .unwrap_or_default(); + let bytes = get_from_api( + HttpRequest { + url: url.into(), + headers, + fetch: self.fetch.clone(), + ..Default::default() + }, + create_binary_response_handler(), + crate::google::google_failed_response_handler(), + ) + .await? + .value; + String::from_utf8(bytes.to_vec()).map_err(|_| { + AiMuxError::InvalidResponseData("Google subject token is not UTF-8".into()) + })? + } else { + return Err(AiMuxError::UnsupportedFunctionality( + "Google external credential sources other than file or URL are not supported" + .into(), + )); + }; + let token = if !parse_format { + token + } else { + match source + .get("format") + .and_then(|format| format.get("type")) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .unwrap_or("text") + { + "text" => token, + "json" => { + let value: Value = serde_json::from_str(&token).map_err(|error| { + AiMuxError::InvalidResponseData(format!( + "Invalid Google subject-token JSON: {error}" + )) + })?; + required( + &value, + required(&source["format"], "subject_token_field_name")?, + )? + .into() + } + format => { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "Google subject-token format {format} is not supported" + ))); + } + } + }; + if token.is_empty() { + return Err(AiMuxError::InvalidResponseData( + "Unable to parse Google subject token".into(), + )); + } + let token_auth = GoogleAuth::new(self.options.clone(), fetch); + let scope = self + .options + .scopes + .as_ref() + .map_or_else(|| CLOUD_SCOPE.into(), GoogleAuthScopes::joined); + let mut fields = vec![ + ( + "grant_type", + "urn:ietf:params:oauth:grant-type:token-exchange".into(), + ), + ( + "requested_token_type", + "urn:ietf:params:oauth:token-type:access_token".into(), + ), + ("audience", required(credentials, "audience")?.into()), + ( + "subject_token_type", + required(credentials, "subject_token_type")?.into(), + ), + ("subject_token", token), + ( + "scope", + if credentials + .get("service_account_impersonation_url") + .is_some() + { + CLOUD_SCOPE.into() + } else { + scope + }, + ), + ]; + if let Some(project) = credentials + .get("workforce_pool_user_project") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + { + let audience = required(credentials, "audience")?; + if !regex::Regex::new( + r"//iam\.googleapis\.com/locations/[^/]+/workforcePools/[^/]+/providers/.+", + ) + .expect("static workforce audience pattern") + .is_match(audience) + { + return Err(AiMuxError::InvalidArgument( + "workforce_pool_user_project must not be set for non-workforce credentials" + .into(), + )); + } + if credentials + .get("client_id") + .and_then(Value::as_str) + .is_none_or(str::is_empty) + { + fields.push(("options", json!({"userProject": project}).to_string())); + } + } + let response = token_auth + .exchange_headers( + fields, + credentials + .get("token_url") + .and_then(Value::as_str) + .unwrap_or(&format!( + "https://sts.{}/v1/token", + self.universe(credentials) + )), + basic_auth(credentials)?, + ) + .await?; + if let Some(url) = credentials + .get("service_account_impersonation_url") + .and_then(Value::as_str) + { + return token_auth + .impersonate(credentials, url, &response.access_token) + .await; + } + Ok(response) + } + + async fn impersonate( + &self, + credentials: &Value, + url: &str, + token: &str, + ) -> Result { + #[derive(Deserialize)] + struct ImpersonationToken { + #[serde(rename = "accessToken")] + access_token: String, + #[serde(rename = "expireTime")] + expire_time: String, + } + let scopes = self + .options + .scopes + .as_ref() + .map_or_else(|| vec![CLOUD_SCOPE.into()], GoogleAuthScopes::array); + let lifetime = credentials + .get("lifetime") + .or_else(|| { + credentials + .get("service_account_impersonation") + .and_then(|options| options.get("token_lifetime_seconds")) + }) + .and_then(Value::as_u64) + .unwrap_or(3600); + let impersonated_url; + let url = if credentials.get("type").and_then(Value::as_str) + == Some("impersonated_service_account") + { + if url.len() > 256 { + return Err(AiMuxError::InvalidArgument( + "Google impersonation target principal is too long".into(), + )); + } + let target = url + .rsplit('/') + .next() + .and_then(|target| { + target + .strip_suffix(":generateAccessToken") + .or_else(|| target.strip_suffix(":generateIdToken")) + }) + .filter(|target| !target.is_empty()) + .ok_or_else(|| { + AiMuxError::InvalidArgument( + "Invalid Google impersonation target principal".into(), + ) + })?; + let endpoint = credentials + .get("endpoint") + .and_then(Value::as_str) + .unwrap_or("https://iamcredentials.googleapis.com"); + impersonated_url = + format!("{endpoint}/v1/projects/-/serviceAccounts/{target}:generateAccessToken"); + &impersonated_url + } else { + url + }; + let mut body = json!({"scope": scopes, "lifetime": format!("{lifetime}s")}); + if credentials.get("type").and_then(Value::as_str) == Some("impersonated_service_account") { + body["delegates"] = credentials.get("delegates").cloned().unwrap_or(json!([])); + } + let response = post_to_api( + HttpRequest { + url: url.into(), + headers: vec![("Authorization".into(), format!("Bearer {token}"))], + fetch: self.fetch.clone(), + ..Default::default() + }, + HttpBody::Json(body), + create_json_response_handler::(), + crate::google::google_failed_response_handler(), + ) + .await? + .value; + let expiry = + chrono::DateTime::parse_from_rfc3339(&response.expire_time).map_err(|error| { + AiMuxError::InvalidResponseData(format!( + "Invalid Google impersonation expiry: {error}" + )) + })?; + Ok(TokenResponse { + access_token: response.access_token, + expires_in: Some((expiry.timestamp() - chrono::Utc::now().timestamp()).max(0) as u64), + refresh_token: None, + }) + } + + fn universe<'a>(&'a self, credentials: &'a Value) -> &'a str { + self.options + .universe_domain + .as_deref() + .or_else(|| { + self.options + .client_options + .as_ref() + .and_then(|options| { + options + .get("universe_domain") + .or_else(|| options.get("universeDomain")) + }) + .and_then(Value::as_str) + }) + .or_else(|| credentials.get("universe_domain").and_then(Value::as_str)) + .unwrap_or("googleapis.com") + } + + fn validate_options(&self) -> Result<(), AiMuxError> { + let client = self.options.client_options.as_ref(); + if let Some(options) = client { + let Some(options) = options.as_object() else { + return Err(AiMuxError::InvalidArgument( + "Google auth client_options must be an object".into(), + )); + }; + for (name, value) in options { + match name.as_str() { + "projectId" | "project_id" | "quotaProjectId" | "quota_project_id" + | "universeDomain" | "universe_domain" | "apiKey" => { + if !value.is_string() && !value.is_null() { + return Err(AiMuxError::InvalidArgument(format!( + "Google auth client option {name} must be a string" + ))); + } + } + "credentials" if value.is_object() || value.is_null() => {} + "apiAudience" if value.is_string() => {} + "lifetime" if value.as_i64().is_some() => {} + "eagerRefreshThresholdMillis" if value.as_u64().is_some() => {} + "forceRefreshOnFailure" | "useAuthRequestParameters" if value.is_boolean() => {} + _ => { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "Google auth client option {name} is not supported" + ))); + } + } + } + } + let api_key = self + .options + .api_key + .as_deref() + .or_else(|| { + client + .and_then(|options| options.get("apiKey")) + .and_then(Value::as_str) + }) + .filter(|key| !key.is_empty()); + if api_key.is_some() { + return Err( + if self.options.credentials.is_some() + || client + .and_then(|options| options.get("credentials")) + .is_some_and(|value| !value.is_null()) + { + AiMuxError::InvalidArgument("API Keys and Credentials are mutually exclusive authentication methods and cannot be used together.".into()) + } else { + AiMuxError::UnsupportedFunctionality("Google auth API-key clients cannot generate OAuth access tokens; use the Vertex provider api_key setting".into()) + }, + ); + } + Ok(()) + } + + async fn exchange( + &self, + fields: Vec<(&str, String)>, + url: &str, + ) -> Result { + self.exchange_headers(fields, url, Vec::new()).await + } + + async fn exchange_headers( + &self, + fields: Vec<(&str, String)>, + url: &str, + headers: Vec<(String, String)>, + ) -> Result { + let body = url::form_urlencoded::Serializer::new(String::new()) + .extend_pairs(fields) + .finish(); + post_to_api( + HttpRequest { + url: url.into(), + headers, + fetch: self.fetch.clone(), + ..Default::default() + }, + HttpBody::Bytes( + body.into_bytes(), + "application/x-www-form-urlencoded".into(), + ), + create_json_response_handler::(), + crate::google::google_failed_response_handler(), + ) + .await + .map(|response| response.value) + } +} + +pub(super) fn required<'a>(credentials: &'a Value, name: &str) -> Result<&'a str, AiMuxError> { + credentials + .get(name) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .ok_or_else(|| { + AiMuxError::InvalidArgument(format!("Google application credentials require {name}")) + }) +} + +fn basic_auth(credentials: &Value) -> Result, AiMuxError> { + match credentials.get("client_id").and_then(Value::as_str) { + Some(id) => Ok(vec![( + "Authorization".into(), + format!( + "Basic {}", + base64::engine::general_purpose::STANDARD + .encode(format!("{id}:{}", required(credentials, "client_secret")?)) + ), + )]), + None => Ok(Vec::new()), + } +} + +fn der_wrap(tag: u8, bytes: Vec) -> Vec { + let mut result = vec![tag]; + if bytes.len() < 128 { + result.push(bytes.len() as u8); + } else { + let length = bytes.len().to_be_bytes(); + let length = &length[length + .iter() + .position(|byte| *byte != 0) + .unwrap_or(length.len() - 1)..]; + result.push(0x80 | length.len() as u8); + result.extend_from_slice(length); + } + result.extend(bytes); + result +} + +fn normalize_external_credentials(mut credentials: Value) -> Value { + aliases( + &mut credentials, + &[ + "token_url", + "subject_token_type", + "client_id", + "client_secret", + "workforce_pool_user_project", + "service_account_impersonation_url", + "service_account_impersonation", + "universe_domain", + "credential_source", + ], + ); + if let Some(impersonation) = credentials.get_mut("service_account_impersonation") { + aliases(impersonation, &["token_lifetime_seconds"]); + } + if let Some(source) = credentials.get_mut("credential_source") { + aliases( + source, + &[ + "region_url", + "regional_cred_verification_url", + "imdsv2_session_token_url", + ], + ); + if let Some(format) = source.get_mut("format") { + aliases(format, &["subject_token_field_name"]); + } + } + credentials +} + +fn aliases(value: &mut Value, fields: &[&str]) { + let Some(value) = value.as_object_mut() else { + return; + }; + for field in fields { + if value.get(*field).is_some_and(|value| !value.is_null()) { + continue; + } + let mut parts = field.split('_'); + let mut camel = parts.next().unwrap_or_default().to_string(); + for part in parts { + let mut chars = part.chars(); + if let Some(first) = chars.next() { + camel.extend(first.to_uppercase()); + camel.extend(chars); + } + } + if let Some(alias) = value.get(&camel).cloned() { + value.insert((*field).into(), alias); + } + } +} diff --git a/aimux-providers/src/vertex/auth_certificate.rs b/aimux-providers/src/vertex/auth_certificate.rs new file mode 100644 index 000000000..6c6368ee0 --- /dev/null +++ b/aimux-providers/src/vertex/auth_certificate.rs @@ -0,0 +1,194 @@ +//! Certificate subject tokens and their native TLS transport. + +use std::path::PathBuf; +use std::sync::Arc; + +use aimux_core::AiMuxError; +use aimux_provider_utils::{Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse}; +use base64::Engine; +use futures::StreamExt; +use serde_json::Value; + +fn invalid(message: impl std::fmt::Display) -> AiMuxError { + AiMuxError::InvalidArgument(message.to_string()) +} + +fn read(path: &str) -> Result, AiMuxError> { + std::fs::read(path) + .map_err(|error| invalid(format!("Cannot read certificate source {path}: {error}"))) +} + +fn nonempty(value: Option<&str>) -> Option<&str> { + value.filter(|value| !value.is_empty()) +} + +fn env(name: &str) -> Option { + std::env::var(name).ok().filter(|value| !value.is_empty()) +} + +pub(super) fn certificate( + options: &Value, + fetch: Option, +) -> Result<(String, Option), AiMuxError> { + let use_default = options["use_default_certificate_config"].as_bool() == Some(true); + let location = nonempty(options["certificate_config_location"].as_str()); + if use_default == location.is_some() { + return Err(invalid( + "Provide either use_default_certificate_config or certificate_config_location", + )); + } + let path = location + .map(PathBuf::from) + .or_else(|| env("GOOGLE_API_CERTIFICATE_CONFIG").map(PathBuf::from)) + .unwrap_or_else(|| { + let directory = env("CLOUDSDK_CONFIG") + .map(PathBuf::from) + .unwrap_or_else(|| { + if cfg!(windows) { + PathBuf::from(env("APPDATA").unwrap_or_default()).join("gcloud") + } else { + PathBuf::from(env("HOME").unwrap_or_default()).join(".config/gcloud") + } + }); + directory.join("certificate_config.json") + }); + let config: Value = serde_json::from_slice(&read(&path.to_string_lossy())?).map_err(invalid)?; + let workload = &config["cert_configs"]["workload"]; + let cert_path = nonempty(workload["cert_path"].as_str()) + .ok_or_else(|| invalid("Missing workload cert_path"))?; + let key_path = nonempty(workload["key_path"].as_str()) + .ok_or_else(|| invalid("Missing workload key_path"))?; + let cert = read(cert_path)?; + let key = read(key_path)?; + let certs = pem_certificates(&cert)?; + let leaf = certs.first().cloned().unwrap_or_else(|| cert.clone()); + let mut identity_pem = if !certs.is_empty() { + cert + } else { + format!( + "-----BEGIN CERTIFICATE-----\n{}\n-----END CERTIFICATE-----\n", + base64::engine::general_purpose::STANDARD.encode(&leaf) + ) + .into_bytes() + }; + identity_pem.push(b'\n'); + identity_pem.extend(key); + let identity = reqwest::Identity::from_pem(&identity_pem).map_err(invalid)?; + let mut chain = match nonempty(options["trust_chain_path"].as_str()) { + Some(path) => pem_certificates(&read(path)?)?, + None => Vec::new(), + }; + match chain.iter().position(|cert| cert == &leaf) { + Some(0) => {} + Some(index) => { + return Err(invalid(format!( + "Leaf certificate exists in trust chain at index {index}, rather than first" + ))); + } + None => chain.insert(0, leaf), + } + // Root-store parsing validates DER without changing the mTLS server trust. + let mut validation = reqwest::Client::builder(); + for der in &chain { + validation = + validation.add_root_certificate(reqwest::Certificate::from_der(der).map_err(invalid)?); + } + validation.build().map_err(invalid)?; + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .identity(identity) + .build() + .map_err(invalid)?; + let subject_token = serde_json::to_string( + &chain + .iter() + .map(|der| base64::engine::general_purpose::STANDARD.encode(der)) + .collect::>(), + ) + .map_err(invalid)?; + Ok(( + subject_token, + Some(fetch.unwrap_or_else(|| Arc::new(CertificateFetch(client)))), + )) +} + +fn pem_certificates(bytes: &[u8]) -> Result>, AiMuxError> { + let text = String::from_utf8_lossy(bytes); + let pattern = regex::Regex::new(r"-----BEGIN CERTIFICATE-----[^-]+-----END CERTIFICATE-----") + .map_err(invalid)?; + pattern + .find_iter(&text) + .map(|block| { + let encoded: String = block + .as_str() + .trim_start_matches("-----BEGIN CERTIFICATE-----") + .trim_end_matches("-----END CERTIFICATE-----") + .chars() + .filter(|character| !character.is_whitespace()) + .collect(); + base64::engine::general_purpose::STANDARD + .decode(encoded) + .map_err(invalid) + }) + .collect() +} + +pub(super) fn ca_fetch(path: &str) -> Result { + let certificates = reqwest::Certificate::from_pem_bundle(&read(path)?).map_err(invalid)?; + let mut builder = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .tls_built_in_root_certs(false); + for certificate in certificates { + builder = builder.add_root_certificate(certificate); + } + Ok(Arc::new(CertificateFetch( + builder.build().map_err(invalid)?, + ))) +} + +struct CertificateFetch(reqwest::Client); + +#[async_trait::async_trait] +impl Fetch for CertificateFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + let signal = request.signal; + let mut builder = self + .0 + .request(request.method, request.url) + .headers(request.headers) + .body(request.body); + if let Some(timeout) = request.timeout { + builder = builder.timeout(timeout); + } + let response = match signal { + Some(signal) => tokio::select! { + biased; + () = signal.cancelled() => return Err(FetchError::Aborted), + response = builder.send() => response, + }, + None => builder.send().await, + } + .map_err(fetch_error)?; + Ok(FetchResponse { + status: response.status(), + headers: response.headers().clone(), + url: response.url().clone(), + body: response + .bytes_stream() + .map(|chunk| chunk.map_err(fetch_error)) + .boxed(), + }) + } +} + +fn fetch_error(error: reqwest::Error) -> FetchError { + if error.is_builder() { + FetchError::Other(error.to_string()) + } else if error.is_timeout() { + FetchError::Timeout + } else if error.is_connect() { + FetchError::Connect(error.to_string()) + } else { + FetchError::Io(error.to_string()) + } +} diff --git a/aimux-providers/src/vertex/auth_external.rs b/aimux-providers/src/vertex/auth_external.rs new file mode 100644 index 000000000..26f533e73 --- /dev/null +++ b/aimux-providers/src/vertex/auth_external.rs @@ -0,0 +1,383 @@ +//! Subject-token sources for Google workload identity federation. + +use std::collections::BTreeMap; + +use aimux_core::AiMuxError; +use aimux_provider_utils::{FetchFunction, FetchRequest, default_fetch}; +use futures::TryStreamExt; +use serde_json::{Value, json}; + +use super::auth::required; + +pub(super) async fn text( + fetch: Option, + url: &str, + method: reqwest::Method, + headers: &[(String, String)], +) -> Result { + let mut request = FetchRequest::new( + method, + url.parse().map_err(|error| { + AiMuxError::InvalidArgument(format!("Invalid Google credential URL: {error}")) + })?, + ); + for (name, value) in headers { + request.headers.insert( + reqwest::header::HeaderName::from_bytes(name.as_bytes()) + .map_err(|error| AiMuxError::InvalidArgument(error.to_string()))?, + value + .parse() + .map_err(|error: reqwest::header::InvalidHeaderValue| { + AiMuxError::InvalidArgument(error.to_string()) + })?, + ); + } + let response = fetch + .unwrap_or_else(default_fetch) + .fetch(request) + .await + .map_err(|error| { + AiMuxError::InvalidResponseData(format!( + "Google credential source request failed: {error}" + )) + })?; + if !response.status.is_success() { + return Err(AiMuxError::InvalidResponseData(format!( + "Google credential source returned HTTP {}", + response.status + ))); + } + let bytes = response + .body + .try_fold(Vec::new(), |mut bytes, chunk| async move { + bytes.extend_from_slice(&chunk); + Ok(bytes) + }) + .await + .map_err(|error| AiMuxError::InvalidResponseData(error.to_string()))?; + String::from_utf8(bytes).map_err(|error| AiMuxError::InvalidResponseData(error.to_string())) +} + +pub(super) async fn aws( + credentials: &Value, + fetch: Option, +) -> Result { + let source = &credentials["credential_source"]; + if required(source, "environment_id")? != "aws1" { + return Err(AiMuxError::InvalidArgument( + "Unsupported Google external AWS environment version".into(), + )); + } + let region = std::env::var("AWS_REGION") + .ok() + .filter(|value| !value.is_empty()) + .or_else(|| { + std::env::var("AWS_DEFAULT_REGION") + .ok() + .filter(|value| !value.is_empty()) + }); + let access = std::env::var("AWS_ACCESS_KEY_ID") + .ok() + .filter(|value| !value.is_empty()); + let secret = std::env::var("AWS_SECRET_ACCESS_KEY") + .ok() + .filter(|value| !value.is_empty()); + let mut metadata_headers = Vec::new(); + if (region.is_none() || access.is_none() || secret.is_none()) + && let Some(url) = source + .get("imdsv2_session_token_url") + .and_then(Value::as_str) + { + let token = text( + fetch.clone(), + url, + reqwest::Method::PUT, + &[("x-aws-ec2-metadata-token-ttl-seconds".into(), "300".into())], + ) + .await?; + metadata_headers.push(("x-aws-ec2-metadata-token".into(), token)); + } + let region = match region { + Some(region) => region, + None => { + let mut zone = text( + fetch.clone(), + required(source, "region_url")?, + reqwest::Method::GET, + &metadata_headers, + ) + .await?; + zone.pop(); + zone + } + }; + let (access, secret, token) = match (access, secret) { + (Some(access), Some(secret)) => (access, secret, std::env::var("AWS_SESSION_TOKEN").ok()), + _ => { + let url = required(source, "url")?; + let role = text(fetch.clone(), url, reqwest::Method::GET, &metadata_headers).await?; + let value: Value = serde_json::from_str( + &text( + fetch, + &format!("{url}/{role}"), + reqwest::Method::GET, + &metadata_headers, + ) + .await?, + ) + .map_err(|error| AiMuxError::InvalidResponseData(error.to_string()))?; + ( + required(&value, "AccessKeyId")?.into(), + required(&value, "SecretAccessKey")?.into(), + value + .get("Token") + .and_then(Value::as_str) + .map(str::to_string), + ) + } + }; + let url = required(source, "regional_cred_verification_url")?.replace("{region}", ®ion); + let parsed = + url::Url::parse(&url).map_err(|error| AiMuxError::InvalidArgument(error.to_string()))?; + let host = parsed.host_str().ok_or_else(|| { + AiMuxError::InvalidArgument("AWS credential verification URL has no host".into()) + })?; + let host = parsed + .port() + .map_or_else(|| host.into(), |port| format!("{host}:{port}")); + let now = chrono::Utc::now(); + let date = now.format("%Y%m%d").to_string(); + let time = now.format("%Y%m%dT%H%M%SZ").to_string(); + let mut headers = BTreeMap::from([("host", host.clone()), ("x-amz-date", time.clone())]); + if let Some(token) = token.filter(|value| !value.is_empty()) { + headers.insert("x-amz-security-token", token); + } + let names = headers.keys().copied().collect::>().join(";"); + let canonical_headers = headers + .iter() + .map(|(key, value)| format!("{key}:{value}\n")) + .collect::(); + let canonical = format!( + "POST\n{}\n{}\n{canonical_headers}\n{names}\n{}", + parsed.path(), + parsed.query().unwrap_or(""), + sha256(b"") + ); + let service = host.split('.').next().unwrap_or("sts"); + let scope = format!("{date}/{region}/{service}/aws4_request"); + let string_to_sign = format!( + "AWS4-HMAC-SHA256\n{time}\n{scope}\n{}", + sha256(canonical.as_bytes()) + ); + let key = hmac(format!("AWS4{secret}").as_bytes(), &date); + let key = hmac(&key, ®ion); + let key = hmac(&key, service); + let key = hmac(&key, "aws4_request"); + let signature = hex(&hmac(&key, &string_to_sign)); + headers.insert("authorization", format!("AWS4-HMAC-SHA256 Credential={access}/{scope}, SignedHeaders={names}, Signature={signature}")); + headers.insert( + "x-goog-cloud-target-resource", + required(credentials, "audience")?.into(), + ); + let headers = headers + .into_iter() + .map(|(key, value)| json!({"key": key, "value": value})) + .collect::>(); + let token = json!({"url": url, "method": "POST", "headers": headers}).to_string(); + Ok( + percent_encoding::utf8_percent_encode(&token, percent_encoding::NON_ALPHANUMERIC) + .to_string() + .replace("%2D", "-") + .replace("%5F", "_") + .replace("%2E", ".") + .replace("%7E", "~") + .replace("%21", "!") + .replace("%27", "'") + .replace("%28", "(") + .replace("%29", ")") + .replace("%2A", "*"), + ) +} + +fn hmac(key: &[u8], message: &str) -> Vec { + ring::hmac::sign( + &ring::hmac::Key::new(ring::hmac::HMAC_SHA256, key), + message.as_bytes(), + ) + .as_ref() + .to_vec() +} +fn sha256(value: &[u8]) -> String { + hex(ring::digest::digest(&ring::digest::SHA256, value).as_ref()) +} +fn hex(value: &[u8]) -> String { + value.iter().map(|byte| format!("{byte:02x}")).collect() +} + +pub(super) async fn executable(credentials: &Value) -> Result { + if std::env::var("GOOGLE_EXTERNAL_ACCOUNT_ALLOW_EXECUTABLES").as_deref() != Ok("1") { + return Err(AiMuxError::InvalidArgument("Google external account executables require GOOGLE_EXTERNAL_ACCOUNT_ALLOW_EXECUTABLES=1".into())); + } + let options = &credentials["credential_source"]["executable"]; + let output_file = options + .get("output_file") + .and_then(Value::as_str) + .filter(|path| !path.is_empty()); + if let Some(path) = output_file + && let Ok(output) = std::fs::read_to_string(path) + && !output.is_empty() + { + let value: Value = serde_json::from_str(&output).map_err(|error| { + AiMuxError::InvalidResponseData(format!("Invalid cached executable response: {error}")) + })?; + validate_executable_response(&value)?; + if value.get("success").and_then(Value::as_bool) == Some(true) + && value + .get("expiration_time") + .and_then(Value::as_i64) + .is_none_or(|expiry| expiry >= (chrono::Utc::now().timestamp_millis() + 500) / 1000) + { + return executable_response(&value, true); + } + } + let timeout = options + .get("timeout_millis") + .and_then(Value::as_u64) + .unwrap_or(30_000); + if !(5_000..=120_000).contains(&timeout) { + return Err(AiMuxError::InvalidArgument( + "Google credential executable timeout must be between 5000 and 120000 milliseconds" + .into(), + )); + } + let command = required(options, "command")?; + let pattern = regex::Regex::new(r#"(?:[^\s"]+|"[^"]*")+"#).expect("static command pattern"); + let components = pattern + .find_iter(command) + .map(|part| { + let part = part.as_str(); + part.strip_prefix('"') + .and_then(|part| part.strip_suffix('"')) + .unwrap_or(part) + }) + .collect::>(); + let program = components.first().ok_or_else(|| { + AiMuxError::InvalidArgument("Google credential executable command is empty".into()) + })?; + let mut command = tokio::process::Command::new(program); + command + .args(&components[1..]) + .kill_on_drop(true) + .env( + "GOOGLE_EXTERNAL_ACCOUNT_AUDIENCE", + required(credentials, "audience")?, + ) + .env( + "GOOGLE_EXTERNAL_ACCOUNT_TOKEN_TYPE", + required(credentials, "subject_token_type")?, + ) + .env("GOOGLE_EXTERNAL_ACCOUNT_INTERACTIVE", "0"); + if let Some(path) = output_file { + command.env("GOOGLE_EXTERNAL_ACCOUNT_OUTPUT_FILE", path); + } + if let Some(url) = credentials + .get("service_account_impersonation_url") + .and_then(Value::as_str) + && let Some(email) = url + .rsplit('/') + .next() + .and_then(|target| target.strip_suffix(":generateAccessToken")) + { + command.env("GOOGLE_EXTERNAL_ACCOUNT_IMPERSONATED_EMAIL", email); + } + let output = tokio::time::timeout(std::time::Duration::from_millis(timeout), command.output()) + .await + .map_err(|_| { + AiMuxError::InvalidResponseData("Google credential executable timed out".into()) + })? + .map_err(|error| { + AiMuxError::InvalidResponseData(format!("Google credential executable failed: {error}")) + })?; + if !output.status.success() { + return Err(AiMuxError::InvalidResponseData(format!( + "Google credential executable failed with {}", + output.status + ))); + } + let mut bytes = output.stdout; + bytes.extend(output.stderr); + let response: Value = serde_json::from_slice(&bytes).map_err(|error| { + AiMuxError::InvalidResponseData(format!("Invalid credential executable response: {error}")) + })?; + executable_response(&response, output_file.is_some()) +} + +fn executable_response(response: &Value, output_file: bool) -> Result { + validate_executable_response(response)?; + if response + .get("version") + .and_then(Value::as_u64) + .is_some_and(|version| version > 1) + { + return Err(AiMuxError::InvalidResponseData( + "Google credential executable response requires version 1".into(), + )); + } + if response.get("success").and_then(Value::as_bool) != Some(true) { + return Err(AiMuxError::InvalidResponseData(format!( + "Google credential executable returned {}: {}", + response["code"], response["message"] + ))); + } + let expiry = response.get("expiration_time").and_then(Value::as_i64); + if (output_file && expiry.is_none()) + || expiry + .is_some_and(|expiry| expiry < (chrono::Utc::now().timestamp_millis() + 500) / 1000) + { + return Err(AiMuxError::InvalidResponseData( + "Google credential executable response has missing or expired expiration_time".into(), + )); + } + let field = match required(response, "token_type")? { + "urn:ietf:params:oauth:token-type:saml2" => "saml_response", + "urn:ietf:params:oauth:token-type:id_token" | "urn:ietf:params:oauth:token-type:jwt" => { + "id_token" + } + _ => { + return Err(AiMuxError::InvalidResponseData( + "Unsupported Google credential executable token_type".into(), + )); + } + }; + Ok(required(response, field)?.into()) +} + +fn validate_executable_response(response: &Value) -> Result<(), AiMuxError> { + if response + .get("version") + .and_then(Value::as_u64) + .is_none_or(|version| version == 0) + || response.get("success").and_then(Value::as_bool).is_none() + { + return Err(AiMuxError::InvalidResponseData( + "Google credential executable response requires version and success".into(), + )); + } + if response.get("success").and_then(Value::as_bool) == Some(false) { + required(response, "code")?; + required(response, "message")?; + } else { + let field = match required(response, "token_type")? { + "urn:ietf:params:oauth:token-type:saml2" => "saml_response", + "urn:ietf:params:oauth:token-type:id_token" + | "urn:ietf:params:oauth:token-type:jwt" => "id_token", + _ => { + return Err(AiMuxError::InvalidResponseData( + "Unsupported Google credential executable token_type".into(), + )); + } + }; + required(response, field)?; + } + Ok(()) +} diff --git a/aimux-providers/src/vertex/embedding.rs b/aimux-providers/src/vertex/embedding.rs index 0807adb06..697d92c81 100644 --- a/aimux-providers/src/vertex/embedding.rs +++ b/aimux-providers/src/vertex/embedding.rs @@ -56,7 +56,7 @@ impl EmbeddingModel for VertexEmbeddingModel { if uses_embed_content_endpoint(&self.model_id) { Some(1) } else { - Some(2048) + Some(250) } } @@ -69,7 +69,16 @@ impl EmbeddingModel for VertexEmbeddingModel { options: &EmbeddingCallOptions, ) -> Result { // Parse provider options: `googleVertex`, then `google`. - let vertex_options = parse_vertex_provider_options(options.provider_options.as_ref()); + let vertex_options = parse_vertex_provider_options(options.provider_options.as_ref())?; + let max = self.max_embeddings_per_call().unwrap_or(250); + if options.values.len() > max as usize { + return Err(AiMuxError::InvalidArgument(format!( + "Too many embedding values for {} model {}: maximum {max}, received {}", + self.provider(), + self.model_id, + options.values.len() + ))); + } let exchange = self.config.exchange(options.headers.as_ref()).await?; @@ -88,7 +97,7 @@ impl EmbeddingModel for VertexEmbeddingModel { ); let mut embed_config = Map::new(); - if let Some(dim) = vertex_options.output_dimensionality { + if let Some(dim) = &vertex_options.output_dimensionality { embed_config.insert("outputDimensionality".to_string(), json!(dim)); } if let Some(task_type) = &vertex_options.task_type { @@ -171,7 +180,7 @@ impl EmbeddingModel for VertexEmbeddingModel { .collect(); let mut parameters = Map::new(); - if let Some(dim) = vertex_options.output_dimensionality { + if let Some(dim) = &vertex_options.output_dimensionality { parameters.insert("outputDimensionality".to_string(), json!(dim)); } if let Some(auto_truncate) = vertex_options.auto_truncate { @@ -247,7 +256,7 @@ impl EmbeddingModel for VertexEmbeddingModel { // ── Provider options parsing ───────────────────────────────────────────────── struct VertexEmbeddingProviderOptions { - output_dimensionality: Option, + output_dimensionality: Option, task_type: Option, title: Option, auto_truncate: Option, @@ -255,18 +264,54 @@ struct VertexEmbeddingProviderOptions { /// Parse Vertex embedding provider options. /// -/// Tries the `googleVertex` key first, then `google` (the shared Gemini -/// model's key). +/// Tries `googleVertex`, then legacy `vertex`, then `google`. fn parse_vertex_provider_options( options: Option<&SharedProviderOptions>, -) -> VertexEmbeddingProviderOptions { - let provider_opts = Namespace::Vertex.read_in(options); +) -> Result { + let provider_opts = Namespace::Vertex.read(options); + if let Some(opts) = provider_opts { + for (name, valid) in [ + ( + "outputDimensionality", + opts.get("outputDimensionality") + .is_none_or(Value::is_number), + ), + ("title", opts.get("title").is_none_or(Value::is_string)), + ( + "autoTruncate", + opts.get("autoTruncate").is_none_or(Value::is_boolean), + ), + ( + "taskType", + opts.get("taskType").is_none_or(|v| { + v.as_str().is_some_and(|v| { + matches!( + v, + "SEMANTIC_SIMILARITY" + | "CLASSIFICATION" + | "CLUSTERING" + | "RETRIEVAL_DOCUMENT" + | "RETRIEVAL_QUERY" + | "QUESTION_ANSWERING" + | "FACT_VERIFICATION" + | "CODE_RETRIEVAL_QUERY" + ) + }) + }), + ), + ] { + if !valid { + return Err(AiMuxError::InvalidArgument(format!( + "Invalid Vertex embedding provider option: {name}" + ))); + } + } + } - VertexEmbeddingProviderOptions { + Ok(VertexEmbeddingProviderOptions { output_dimensionality: provider_opts .and_then(|o| o.get("outputDimensionality")) - .and_then(serde_json::Value::as_u64) - .map(|d| d as u32), + .cloned(), task_type: provider_opts .and_then(|o| o.get("taskType")) .and_then(|v| v.as_str()) @@ -278,5 +323,5 @@ fn parse_vertex_provider_options( auto_truncate: provider_opts .and_then(|o| o.get("autoTruncate")) .and_then(serde_json::Value::as_bool), - } + }) } diff --git a/aimux-providers/src/vertex/gemini_transcription.rs b/aimux-providers/src/vertex/gemini_transcription.rs new file mode 100644 index 000000000..00485b88a --- /dev/null +++ b/aimux-providers/src/vertex/gemini_transcription.rs @@ -0,0 +1,459 @@ +//! Gemini transcription through Vertex generateContent and Live API. + +use async_trait::async_trait; +use serde::Deserialize; +use serde_json::{Value, json}; + +use super::ProjectLocationFn; +use crate::google::options::{GOOGLE, Namespace}; +use crate::shared::EndpointConfig; +use aimux_core::error::AiMuxError; +use aimux_core::shared::{SharedProviderOptions, provider_namespace}; +use aimux_core::transcription_model::{ + AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionResponse, + TranscriptionResult, TranscriptionSegment, TranscriptionStreamOptions, + TranscriptionStreamResult, +}; + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct GeminiOptions { + #[serde(default)] + language_codes: Option>, + #[serde(default)] + custom_vocabulary: Option>, + #[serde(default)] + word_timestamp: Option, + #[serde(default)] + diarization: Option, + #[serde(default)] + mode: Option, +} + +#[derive(Deserialize)] +enum Mode { + #[serde(rename = "SMART")] + Smart, + #[serde(rename = "VERBATIM")] + Verbatim, +} + +fn transcription_config(options: Option<&SharedProviderOptions>) -> Result { + let Some(options) = Namespace::Vertex.read(options) else { + return Ok(json!({})); + }; + // z.optional() accepts absence, but not explicit null. + for key in [ + "languageCodes", + "customVocabulary", + "wordTimestamp", + "diarization", + "mode", + ] { + if options.get(key).is_some_and(Value::is_null) { + return Err(AiMuxError::InvalidArgument(format!( + "Google Vertex transcription option {key} cannot be null" + ))); + } + } + let parsed: GeminiOptions = serde_json::from_value(Value::Object(options.clone())) + .map_err(|error| AiMuxError::InvalidArgument(error.to_string()))?; + let mut config = json!({}); + if let Some(value) = parsed.language_codes { + config["languageCodes"] = json!(value); + } + if let Some(value) = parsed.custom_vocabulary { + config["customVocabulary"] = json!(value); + } + if let Some(value) = parsed.word_timestamp { + config["wordTimestamp"] = json!(value); + } + if let Some(value) = parsed.diarization { + config["diarization"] = json!(value); + } + if let Some(value) = parsed.mode { + config["mode"] = json!(match value { + Mode::Smart => "SMART", + Mode::Verbatim => "VERBATIM", + }); + } + Ok(config) +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct Word { + word: Option, + start_offset: Option, + end_offset: Option, +} +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct AudioTranscription { + text: Option, + language_code: Option, + #[serde(rename = "speakerLabel")] + _speaker_label: Option, + words: Option>, +} +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct Part { + text: Option, + audio_transcription: Option, +} +#[derive(Deserialize)] +struct Content { + parts: Option>, +} +#[derive(Deserialize)] +struct Candidate { + content: Option, +} +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct Response { + candidates: Option>, + usage_metadata: Option>, +} + +fn offset(value: &Option) -> Option { + // Number.parseFloat accepts a numeric prefix, including duration suffixes. + let value = value.as_ref()?.trim_start(); + (1..=value.len()) + .rev() + .filter(|end| value.is_char_boundary(*end)) + .find_map(|end| value[..end].parse::().ok()) + .filter(|value| value.is_finite()) +} + +/// Gemini transcription model, distinct from Cloud Speech-to-Text. +pub struct VertexGeminiTranscriptionModel { + model_id: String, + config: EndpointConfig, + project_location: ProjectLocationFn, + #[cfg(feature = "realtime")] + web_socket: Option>, +} + +impl VertexGeminiTranscriptionModel { + pub(crate) fn from_config( + model_id: String, + config: EndpointConfig, + project_location: ProjectLocationFn, + #[cfg(feature = "realtime")] web_socket: Option< + std::sync::Arc, + >, + ) -> Self { + Self { + model_id, + config, + project_location, + #[cfg(feature = "realtime")] + web_socket, + } + } +} + +#[async_trait] +impl TranscriptionModel for VertexGeminiTranscriptionModel { + fn provider(&self) -> &str { + &self.config.provider + } + fn model_id(&self) -> &str { + &self.model_id + } + + async fn do_generate( + &self, + options: &TranscriptionCallOptions, + ) -> Result { + if self.model_id.contains("-live") { + return Err(AiMuxError::InvalidArgument(format!( + "Model '{}' only supports streaming transcription. Use stream_transcribe or a unary model.", + self.model_id + ))); + } + (self.project_location)().await?; + let timestamp = chrono::Utc::now().to_rfc3339(); + let config = transcription_config(options.provider_options.as_ref())?; + let audio = match &options.audio { + AudioInput::Base64(value) => value.clone(), + AudioInput::Binary(value) => { + base64::Engine::encode(&base64::engine::general_purpose::STANDARD, value) + } + }; + let mut body = json!({"contents": [{"role": "user", "parts": [{"inlineData": {"mimeType": options.media_type, "data": audio}}]}]}); + if config.as_object().is_some_and(|config| !config.is_empty()) { + body["generationConfig"] = json!({"audioTranscriptionConfig": config}); + } + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let response = aimux_provider_utils::post_json_to_api( + exchange.request( + exchange.url(&format!("/models/{}:generateContent", self.model_id)), + options, + ), + exchange.transform_body(body), + aimux_provider_utils::create_json_response_handler::(), + crate::google::google_failed_response_handler(), + ) + .await?; + let parts = response + .value + .candidates + .unwrap_or_default() + .into_iter() + .next() + .and_then(|candidate| candidate.content) + .and_then(|content| content.parts) + .unwrap_or_default(); + let mut text: String = parts + .iter() + .filter_map(|part| part.text.as_deref()) + .collect(); + if text.is_empty() { + text = parts + .iter() + .filter_map(|part| part.audio_transcription.as_ref()?.text.as_deref()) + .collect(); + } + let mut language = None; + let mut segments = Vec::new(); + for part in parts { + let Some(transcription) = part.audio_transcription else { + continue; + }; + if language.is_none() { + language = transcription.language_code; + } + for word in transcription.words.unwrap_or_default() { + if let (Some(text), Some(start_second), Some(end_second)) = ( + word.word.clone(), + offset(&word.start_offset), + offset(&word.end_offset), + ) { + segments.push(TranscriptionSegment { + text, + start_second, + end_second, + }); + } + } + } + Ok(TranscriptionResult { + text, + segments, + language, + duration_in_seconds: None, + warnings: Vec::new(), + request: None, + response: TranscriptionResponse { + timestamp: Some(timestamp), + model_id: Some(self.model_id.clone()), + headers: Some(response.response_headers), + body: response.raw_value, + }, + provider_metadata: response.value.usage_metadata.map(|usage| { + provider_namespace(GOOGLE, json!({"usageMetadata": usage})) + .expect("provider metadata must be an object") + }), + }) + } + + async fn do_stream( + &self, + options: TranscriptionStreamOptions, + ) -> Result { + if !self.model_id.contains("-live") { + return Err(AiMuxError::InvalidArgument(format!( + "Model '{}' does not support streaming transcription. Use a live model.", + self.model_id + ))); + } + #[cfg(feature = "realtime")] + { + self.stream_live(options).await + } + #[cfg(not(feature = "realtime"))] + { + let _ = options; + Err(AiMuxError::UnsupportedFunctionality( + "Enable realtime for Vertex Live transcription".into(), + )) + } + } +} + +#[cfg(feature = "realtime")] +impl VertexGeminiTranscriptionModel { + async fn stream_live( + &self, + options: TranscriptionStreamOptions, + ) -> Result { + use aimux_core::transcription_model::{ + AudioChunk, TranscriptionRequest, TranscriptionStreamPart, + }; + use aimux_provider_utils::ws::{WebSocketRequest, WsMessage, ws_connect}; + use futures::StreamExt; + use std::time::Duration; + use tokio::time::Instant; + + if options.input_audio_format.format_type != "audio/pcm" + || options + .input_audio_format + .rate + .is_some_and(|rate| rate != 16000) + { + return Err(AiMuxError::InvalidArgument( + "The Gemini Live transcription API only supports 16kHz 16-bit PCM input audio." + .into(), + )); + } + let target = (self.project_location)().await?; + let config = transcription_config(options.provider_options.as_ref())?; + let setup = json!({"setup": { + "model": format!("projects/{}/locations/{}/publishers/google/models/{}", target.project, target.location, self.model_id), + "inputAudioTranscription": config, + }}); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let url = format!( + "wss://{}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent", + super::location_host(&target.location) + ); + let mut socket = ws_connect(&WebSocketRequest::for_transcription( + url, + exchange.headers(), + &options, + self.web_socket.clone(), + )) + .await?; + socket.send_text(&setup.to_string()).await?; + let request = Some(TranscriptionRequest { + body: Some(setup["setup"].to_string()), + }); + let response = Some(TranscriptionResponse { + timestamp: Some(chrono::Utc::now().to_rfc3339()), + model_id: Some(self.model_id.clone()), + ..Default::default() + }); + let include_raw = options.include_raw_chunks; + let mut audio = options.audio; + let stream = async_stream::stream! { + yield Ok(TranscriptionStreamPart::StreamStart { warnings: Vec::new() }); + let mut ready = false; + let mut audio_ended = false; + let mut deadline = None; + let mut segment_counter = 0; + let mut segment_buffer = String::new(); + let mut latest_interim = String::new(); + let mut full_text = String::new(); + let mut language = None; + let mut usage_metadata = None; + let mut complete = false; + loop { + tokio::select! { + chunk = audio.next(), if ready && !audio_ended => { + let message = match chunk { + Some(chunk) => { + let data = match chunk { + AudioChunk::Base64(data) => data, + AudioChunk::Binary(data) => base64::Engine::encode(&base64::engine::general_purpose::STANDARD, data), + }; + json!({"realtimeInput": {"audio": {"data": data, "mimeType": "audio/pcm;rate=16000"}}}) + } + None => { + audio_ended = true; + deadline = Some(Instant::now() + Duration::from_secs(3)); + json!({"realtimeInput": {"audioStreamEnd": true}}) + } + }; + if let Err(error) = socket.send_text(&message.to_string()).await { + yield Err(error); break; + } + } + incoming = socket.next() => { + let incoming = match incoming { + Some(Ok(incoming)) => incoming, + Some(Err(error)) => { yield Err(error); break; } + None if audio_ended => { complete = true; break; } + None => { + yield Err(AiMuxError::InvalidResponseData("Vertex Live transcription WebSocket closed unexpectedly before finishing".into())); + break; + } + }; + let text = match incoming { + WsMessage::Text(text) => text, + WsMessage::Binary(data) => String::from_utf8_lossy(&data).into_owned(), + }; + let Ok(message) = serde_json::from_str::(&text) else { continue; }; + if include_raw { yield Ok(TranscriptionStreamPart::Raw { raw_value: message.clone() }); } + if message.get("setupComplete").is_some_and(|value| !value.is_null()) { ready = true; } + if let Some(value) = message.get("usageMetadata").filter(|value| !value.is_null()) { usage_metadata = Some(value.clone()); } + if let Some(error) = message.get("error").filter(|value| !value.is_null()) { + yield Err(AiMuxError::InvalidResponseData(error.get("message").and_then(Value::as_str).unwrap_or("Vertex Live API error").into())); + break; + } + let content = message.get("serverContent"); + if let Some(text) = content.and_then(|value| value.get("interimInputTranscription")).and_then(|value| value.get("text")).and_then(Value::as_str).filter(|text| !text.is_empty()) { + if audio_ended { deadline = Some(Instant::now() + Duration::from_secs(3)); } + latest_interim = text.to_owned(); + yield Ok(TranscriptionStreamPart::TranscriptPartial { + id: Some(format!("google-segment-{segment_counter}")), text: text.into(), start_second: None, duration_in_seconds: None, channel_index: None, provider_metadata: None, + }); + } + let transcription = content.and_then(|value| value.get("inputTranscription")).filter(|value| !value.is_null()).or_else(|| message.get("inputTranscription")); + if let Some(transcription) = transcription { + if let Some(value) = transcription.get("languageCode").and_then(Value::as_str) { language = Some(value.to_owned()); } + if let Some(text) = transcription.get("text").and_then(Value::as_str).filter(|text| !text.is_empty()) { + if audio_ended { deadline = Some(Instant::now() + Duration::from_secs(3)); } + latest_interim.clear(); + segment_buffer.push_str(text); + yield Ok(TranscriptionStreamPart::TranscriptDelta { id: Some(format!("google-segment-{segment_counter}")), delta: text.into(), provider_metadata: None }); + } + } + let turn_complete = content.and_then(|value| value.get("turnComplete")).and_then(Value::as_bool) == Some(true); + if turn_complete || transcription.and_then(|value| value.get("finished")).and_then(Value::as_bool) == Some(true) { + if segment_buffer.is_empty() { segment_buffer = std::mem::take(&mut latest_interim); } + latest_interim.clear(); + if !segment_buffer.is_empty() { + if !full_text.is_empty() { full_text.push(' '); } + full_text.push_str(&segment_buffer); + yield Ok(TranscriptionStreamPart::TranscriptFinal { + id: Some(format!("google-segment-{segment_counter}")), text: std::mem::take(&mut segment_buffer), start_second: None, end_second: None, channel_index: None, provider_metadata: None, + }); + segment_counter += 1; + } + } + let status = content.and_then(|value| value.get("interactionStatus")).and_then(Value::as_str); + if audio_ended && (matches!(status, Some("IDLE" | "REQUIRES_ACTION")) || (turn_complete && status.is_none())) { + complete = true; break; + } + } + _ = async { + match deadline { Some(deadline) => tokio::time::sleep_until(deadline).await, None => std::future::pending().await } + } => { complete = true; break; } + } + } + if complete { + if segment_buffer.is_empty() { segment_buffer = latest_interim; } + if !segment_buffer.is_empty() { + if !full_text.is_empty() { full_text.push(' '); } + full_text.push_str(&segment_buffer); + yield Ok(TranscriptionStreamPart::TranscriptFinal { + id: Some(format!("google-segment-{segment_counter}")), text: segment_buffer, start_second: None, end_second: None, channel_index: None, provider_metadata: None, + }); + } + yield Ok(TranscriptionStreamPart::Finish { + text: full_text, segments: Vec::new(), language, duration_in_seconds: None, + provider_metadata: usage_metadata.map(|usage| provider_namespace(GOOGLE, json!({"usageMetadata": usage})).expect("provider metadata must be an object")), + }); + } + socket.close().await; + }; + Ok(TranscriptionStreamResult { + stream: Box::pin(stream), + request, + response, + }) + } +} diff --git a/aimux-providers/src/vertex/image.rs b/aimux-providers/src/vertex/image.rs index eb23c746e..87b5f3bf3 100644 --- a/aimux-providers/src/vertex/image.rs +++ b/aimux-providers/src/vertex/image.rs @@ -3,23 +3,25 @@ //! Aligned with Vercel AI SDK `GoogleVertexImageModel` //! (`reference/ai/packages/google-vertex/src/google-vertex-image-model.ts`). //! -//! Two code paths: -//! - **Imagen** models (non-`gemini-*`): `POST {base_url}/models/{id}:predict` -//! - **Gemini** image models (`gemini-*`): `POST {base_url}/models/{id}:generateContent` +//! Uses the Gemini language model endpoint for image generation. use async_trait::async_trait; -use serde_json::{Map, Value, json}; +use serde_json::{Value, json}; +use crate::google::convert::build_vertex_request_body; use crate::google::google_failed_response_handler; -use crate::google::options::{vertex_metadata_map, vertex_options}; +use crate::google::options::{GOOGLE_VERTEX, Namespace}; use crate::shared::EndpointConfig; +use aimux_core::language_model_message::{FilePart, LanguageModelMessage, TextPart, UserPart}; +use aimux_core::options::CallOptions; +use aimux_core::shared::provider_namespace; use aimux_core::error::AiMuxError; use aimux_core::image_model::{ ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, ImageResponse, ImageResult, ImageUsage, }; -use aimux_core::shared::Warning; +use aimux_core::shared::{FileBytes, FileData, Warning}; fn is_gemini_model(model_id: &str) -> bool { model_id.starts_with("gemini-") @@ -39,10 +41,6 @@ impl VertexImageModel { Self { model_id, config } } - fn predict_path(&self) -> String { - format!("/models/{}:predict", self.model_id) - } - fn generate_content_path(&self) -> String { let mp = if self.model_id.contains('/') { self.model_id.clone() @@ -52,165 +50,6 @@ impl VertexImageModel { format!("/{mp}:generateContent") } - fn get_base64_data(file: &ImageFile) -> Result { - match file { - ImageFile::Url { .. } => Err(AiMuxError::InvalidArgument( - "URL-based images are not supported for Google Vertex image editing.".into(), - )), - ImageFile::File { data, .. } => match data { - ImageFileData::Base64(s) => Ok(s.clone()), - ImageFileData::Binary(b) => Ok(base64::Engine::encode( - &base64::engine::general_purpose::STANDARD, - b, - )), - }, - } - } - - async fn do_generate_imagen( - &self, - options: &ImageCallOptions, - ) -> Result { - let mut warnings = Vec::new(); - if options.size.is_some() { - warnings.push(Warning::Unsupported { - feature: "size".into(), - details: Some( - "This model does not support the `size` option. Use `aspectRatio` instead." - .into(), - ), - }); - } - - // Parse provider options (`googleVertex`) - let gv_opts = vertex_options(Some(&options.provider_options)); - let edit_opts = gv_opts.and_then(|o| o.get("edit")); - let edit_mode = edit_opts - .and_then(|e| e.get("mode")) - .and_then(|v| v.as_str()); - let base_steps = edit_opts.and_then(|e| e.get("baseSteps")); - let mask_mode = edit_opts - .and_then(|e| e.get("maskMode")) - .and_then(|v| v.as_str()); - let mask_dilation = edit_opts.and_then(|e| e.get("maskDilation")); - - let is_edit_mode = options.files.as_ref().is_some_and(|f| !f.is_empty()); - - let mut parameters = Map::new(); - parameters.insert("sampleCount".into(), json!(options.n)); - if let Some(ar) = options.aspect_ratio { - parameters.insert("aspectRatio".into(), json!(ar.to_string())); - } - if let Some(seed) = options.seed { - parameters.insert("seed".into(), json!(seed)); - } - - let mut reference_images: Vec = Vec::new(); - - if is_edit_mode { - if let Some(ref files) = options.files { - for (i, file) in files.iter().enumerate() { - reference_images.push(json!({ - "referenceType": "REFERENCE_TYPE_RAW", - "referenceId": i + 1, - "referenceImage": { "bytesBase64Encoded": Self::get_base64_data(file)? } - })); - } - } - if let Some(ref mask) = options.mask { - let mut mask_config = Map::new(); - mask_config.insert( - "maskMode".into(), - json!(mask_mode.unwrap_or("MASK_MODE_USER_PROVIDED")), - ); - if let Some(d) = mask_dilation { - mask_config.insert("dilation".into(), d.clone()); - } - let file_count = options.files.as_ref().map_or(0, std::vec::Vec::len); - reference_images.push(json!({ - "referenceType": "REFERENCE_TYPE_MASK", - "referenceId": file_count + 1, - "referenceImage": { "bytesBase64Encoded": Self::get_base64_data(mask)? }, - "maskImageConfig": Value::Object(mask_config), - })); - } - parameters.insert( - "editMode".into(), - json!(edit_mode.unwrap_or("EDIT_MODE_INPAINT_INSERTION")), - ); - if let Some(bs) = base_steps { - parameters.insert("editConfig".into(), json!({ "baseSteps": bs })); - } - } - - let body = if is_edit_mode { - json!({ - "instances": [{ "prompt": options.prompt, "referenceImages": reference_images }], - "parameters": parameters, - }) - } else { - json!({ "instances": [{ "prompt": options.prompt }], "parameters": parameters }) - }; - - let exchange = self.config.exchange(options.headers.as_ref()).await?; - - let resp = aimux_provider_utils::post_json_to_api( - exchange.request(exchange.url(&self.predict_path()), options), - exchange.transform_body(body), - aimux_provider_utils::create_json_response_handler(), - crate::google::google_failed_response_handler(), - ) - .await?; - let rh = resp.response_headers; - let rb: Value = resp.value; - - let images: Vec = rb - .get("predictions") - .and_then(|p| p.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|p| { - p.get("bytesBase64Encoded") - .and_then(|v| v.as_str()) - .map(String::from) - }) - .collect() - }) - .unwrap_or_default(); - - // Build provider metadata with revisedPrompt - let image_metas: Vec = rb - .get("predictions") - .and_then(|p| p.as_array()) - .map(|arr| { - arr.iter() - .map(|p| { - let mut e = Map::new(); - if let Some(rp) = p.get("prompt").and_then(|v| v.as_str()) { - e.insert("revisedPrompt".into(), json!(rp)); - } - Value::Object(e) - }) - .collect() - }) - .unwrap_or_default(); - - let payload = json!({ "images": image_metas }); - let metadata = vertex_metadata_map(&payload); - - Ok(ImageResult { - images: ImageOutputs::Base64(images), - warnings, - provider_metadata: Some(metadata), - response: ImageResponse { - timestamp: Some(chrono::Utc::now().to_rfc3339()), - model_id: Some(self.model_id.clone()), - headers: Some(rh), - }, - usage: None, - }) - } - async fn do_generate_gemini( &self, options: &ImageCallOptions, @@ -221,9 +60,6 @@ impl VertexImageModel { "Gemini image models do not support mask-based image editing.".into(), )); } - if options.n > 1 { - return Err(AiMuxError::UnsupportedFunctionality("Gemini image models do not support generating a set number of images per call. Use n=1 or omit the n parameter.".into())); - } if options.size.is_some() { warnings.push(Warning::Unsupported { feature: "size".into(), @@ -234,58 +70,78 @@ impl VertexImageModel { }); } - let mut parts: Vec = Vec::new(); + let mut parts = Vec::new(); if let Some(ref p) = options.prompt { - parts.push(json!({ "text": p })); + parts.push(UserPart::Text(TextPart { + text: p.clone(), + provider_options: None, + })); } if let Some(ref files) = options.files { for file in files { match file { ImageFile::Url { url } => { - return Err(AiMuxError::UnsupportedFunctionality(format!( - "URL-based input images with media type \"image/*\" are not passed as inline bytes. URL: {url}" - ))); + let file = FilePart { + data: FileData::Url { + url: url::Url::parse(url) + .map_err(|error| { + AiMuxError::InvalidArgument(error.to_string()) + })? + .to_string(), + original_url: None, + }, + media_type: "image/*".into(), + filename: None, + provider_options: None, + }; + aimux_provider_utils::resolve_full_media_type(&file)?; + parts.push(UserPart::File(file)); } ImageFile::File { media_type, data } => { - let ds = match data { - ImageFileData::Base64(s) => s.clone(), - ImageFileData::Binary(b) => base64::Engine::encode( - &base64::engine::general_purpose::STANDARD, - b, - ), - }; - parts.push(json!({ "inlineData": { "mimeType": media_type, "data": ds } })); + parts.push(UserPart::File(FilePart { + data: FileData::Data { + data: match data { + ImageFileData::Base64(data) => FileBytes::Base64(data.clone()), + ImageFileData::Binary(data) => FileBytes::Binary(data.clone()), + }, + }, + media_type: media_type.clone(), + filename: None, + provider_options: None, + })); } } } } - let mut gc = Map::new(); - gc.insert("responseModalities".into(), json!(["IMAGE"])); + let mut inner_options = Namespace::Vertex + .write_keys() + .iter() + .find_map(|key| options.provider_options.get(*key)) + .cloned() + .unwrap_or_default(); + inner_options.insert("responseModalities".into(), json!(["IMAGE"])); + let mut image_config = inner_options + .get("imageConfig") + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); if let Some(ar) = options.aspect_ratio { - gc.insert( - "imageConfig".into(), - json!({ "aspectRatio": ar.to_string() }), - ); - } - if let Some(seed) = options.seed { - gc.insert("seed".into(), json!(seed)); - } - - // Only declared generation-config options survive the upstream language schema. - let gv_opts = vertex_options(Some(&options.provider_options)); - if let Some(obj) = gv_opts { - for key in ["audioTimestamp", "thinkingConfig", "mediaResolution"] { - if let Some(value) = obj.get(key) { - gc.insert( - key.to_string(), - crate::google::convert::image_generation_option(key, value), - ); - } - } - } - - let body = json!({ "contents": [{ "role": "user", "parts": parts }], "generationConfig": Value::Object(gc) }); + image_config.insert("aspectRatio".into(), json!(ar.to_string())); + } + if options.aspect_ratio.is_some() || inner_options.contains_key("imageConfig") { + inner_options.insert("imageConfig".into(), Value::Object(image_config)); + } + let mut call_options = CallOptions::new(vec![LanguageModelMessage::User { + content: parts, + provider_options: None, + }]); + call_options.seed = options.seed; + call_options.provider_options = Some( + provider_namespace(GOOGLE_VERTEX, Value::Object(inner_options)) + .expect("provider metadata must be an object"), + ); + let body = build_vertex_request_body(&self.model_id, &call_options)?; let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( @@ -299,48 +155,47 @@ impl VertexImageModel { let rb: Value = resp.value; let mut images: Vec = Vec::new(); - if let Some(candidates) = rb.get("candidates").and_then(|c| c.as_array()) { - for c in candidates { - if let Some(parts) = c - .get("content") - .and_then(|c| c.get("parts")) - .and_then(|p| p.as_array()) + if let Some(c) = rb + .get("candidates") + .and_then(|c| c.as_array()) + .and_then(|c| c.first()) + && let Some(parts) = c + .get("content") + .and_then(|c| c.get("parts")) + .and_then(|p| p.as_array()) + { + for p in parts { + if p.get("thought").and_then(Value::as_bool) != Some(true) + && let Some(id) = p.get("inlineData") + && let Some(mt) = id.get("mimeType").and_then(|m| m.as_str()) + && mt.starts_with("image/") + && let Some(d) = id.get("data").and_then(|d| d.as_str()) { - for p in parts { - if let Some(id) = p.get("inlineData") - && let Some(mt) = id.get("mimeType").and_then(|m| m.as_str()) - && mt.starts_with("image/") - && let Some(d) = id.get("data").and_then(|d| d.as_str()) - { - images.push(d.to_string()); - } - } + images.push(d.to_string()); } } } - let usage = rb.get("usageMetadata").map(|u| { - let inp = u - .get("promptTokenCount") - .and_then(serde_json::Value::as_u64) - .map(|x| x as u32); - let out = u - .get("candidatesTokenCount") - .and_then(serde_json::Value::as_u64) - .map(|x| x as u32); - let tot = u - .get("totalTokenCount") - .and_then(serde_json::Value::as_u64) - .map(|x| x as u32); - ImageUsage { - input_tokens: inp, - output_tokens: out, - total_tokens: tot.or_else(|| Some(inp.unwrap_or(0) + out.unwrap_or(0))), - } - }); + let usage = Some( + if let Some(u) = rb.get("usageMetadata").filter(|u| !u.is_null()) { + let count = |name| u.get(name).and_then(Value::as_u64).unwrap_or(0) as u32; + let input_tokens = count("promptTokenCount") + count("toolUsePromptTokenCount"); + let output_tokens = count("candidatesTokenCount") + count("thoughtsTokenCount"); + ImageUsage { + input_tokens: Some(input_tokens), + output_tokens: Some(output_tokens), + total_tokens: Some(input_tokens + output_tokens), + } + } else { + ImageUsage { + total_tokens: Some(0), + ..Default::default() + } + }, + ); let payload = json!({ "images": images.iter().map(|_| json!({})).collect::>() }); - let metadata = vertex_metadata_map(&payload); + let metadata = Namespace::Vertex.metadata(payload); Ok(ImageResult { images: ImageOutputs::Base64(images), @@ -365,18 +220,16 @@ impl ImageModel for VertexImageModel { &self.model_id } fn max_images_per_call(&self) -> Option { - if is_gemini_model(&self.model_id) { - Some(10) - } else { - Some(4) - } + Some(1) } async fn do_generate(&self, options: &ImageCallOptions) -> Result { if is_gemini_model(&self.model_id) { self.do_generate_gemini(options).await } else { - self.do_generate_imagen(options).await + Err(AiMuxError::UnsupportedFunctionality( + "Google image models other than Gemini are no longer supported. Use a model ID that starts with `gemini-`.".into(), + )) } } } diff --git a/aimux-providers/src/vertex/mod.rs b/aimux-providers/src/vertex/mod.rs index d4f20223d..c943292be 100644 --- a/aimux-providers/src/vertex/mod.rs +++ b/aimux-providers/src/vertex/mod.rs @@ -1,33 +1,7 @@ //! Google Vertex AI provider. //! -//! [`create_google_vertex`] is the Rust form of the AI SDK's -//! `createGoogleVertex`: it takes [`VertexProviderSettings`], validates what -//! was given explicitly and returns a [`VertexProvider`]. Nothing is read from -//! the environment and no project, location, key or token is needed until a -//! request is made: each call resolves its mode, base URL and headers then. -//! [`google_vertex()`] is the default instance. -//! -//! Two modes, chosen per request: -//! -//! - **Express mode** when an API key is available (the `api_key` setting, else -//! `GOOGLE_VERTEX_API_KEY`): requests carry `x-goog-api-key` and go to -//! `https://aiplatform.googleapis.com/v1/publishers/google`. -//! - **Standard mode** otherwise: requests carry `Authorization: Bearer ` -//! and go to the project/location-scoped host. The token is the -//! `access_token` setting (any [`Resolvable`], so a host can refresh it) or -//! `GOOGLE_VERTEX_ACCESS_TOKEN`; aimux does not implement Application -//! Default Credentials. A host that authenticates some other way supplies an -//! `Authorization` header through `headers` instead. -//! -//! Vertex serves the Gemini request format, so these models share the -//! conversion in [`crate::google::convert`]; they differ in endpoint, -//! authentication and the providerOptions/metadata namespace (options are -//! read from `googleVertex`, then `google`; metadata is written under -//! `googleVertex` only). Claude on -//! Vertex is the shared Anthropic Messages model (see -//! [`VertexProvider::anthropic_model`]). -//! -//! Reference: +//! API keys select Express mode when the provider is created. Standard mode +//! uses project/location endpoints and application default credentials. use std::sync::Arc; @@ -42,22 +16,32 @@ use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_core::transcription_model::TranscriptionModel; use aimux_core::video_model::VideoModel; use aimux_provider_utils::{ - FetchFunction, HeaderMapOpt, Resolvable, combine_headers, load_api_key, load_setting, - validate_base_url, + FetchFunction, HeaderMapOpt, Resolvable, combine_headers, load_setting, }; use std::sync::OnceLock; -use crate::shared::{Endpoint, EndpointConfig, TransformRequestBody, is_valid_hostname_part}; +use crate::shared::{Endpoint, EndpointConfig, TransformRequestBody}; mod anthropic_model; +mod anthropic_provider; +mod auth; +mod auth_certificate; +mod auth_external; mod embedding; +mod gemini_transcription; pub mod image; mod model; mod transcription; mod video; pub use anthropic_model::VertexAnthropicModel; +pub use anthropic_provider::{ + VertexAnthropicProvider, VertexAnthropicProviderSettings, create_google_vertex_anthropic, + google_vertex_anthropic, +}; +pub use auth::{GoogleAuthOptions, GoogleAuthScopes}; pub use embedding::VertexEmbeddingModel; +pub use gemini_transcription::VertexGeminiTranscriptionModel; pub use image::VertexImageModel; pub use model::VertexModel; pub use transcription::VertexTranscriptionModel; @@ -70,27 +54,21 @@ const ACCESS_TOKEN_ENV_VAR: &str = "GOOGLE_VERTEX_ACCESS_TOKEN"; const PROJECT_ENV_VAR: &str = "GOOGLE_VERTEX_PROJECT"; const LOCATION_ENV_VAR: &str = "GOOGLE_VERTEX_LOCATION"; -/// The provider string of the language, embedding and image models. -const PROVIDER: &str = "google.vertex"; - /// Tuned models are addressed by their `endpoints/{id}` resource. const ENDPOINT_MODEL_PREFIX: &str = "endpoints/"; /// Settings of [`create_google_vertex`] (the AI SDK's -/// `GoogleVertexProviderSettings`, plus the token source aimux needs in place -/// of google-auth-library). +/// `GoogleVertexProviderSettings`). /// -/// Every field is optional. Nothing here is evaluated when the provider is -/// created except an explicit `base_url` and `location`; everything else is -/// evaluated on every request. +/// Every field is optional. The API-key environment setting is captured when +/// the provider is created; header and token producers run on each request. #[derive(Clone, Default)] pub struct VertexProviderSettings { /// Express-mode API key, sent as `x-goog-api-key`. `None` loads - /// `GOOGLE_VERTEX_API_KEY` when a request is made; with no key at all the - /// request uses standard mode. An empty or whitespace-only key also means - /// standard mode. + /// `GOOGLE_VERTEX_API_KEY` when the provider is created. An empty key means + /// standard mode; whitespace is retained as in the upstream provider. pub api_key: Option>, - /// The location (region) of standard mode: a single DNS label such as + /// The location (region) of standard mode, such as /// `us-central1`, `global`, `us` or `eu`. `None` loads /// `GOOGLE_VERTEX_LOCATION` when a request is made and fails that request /// with `AiMuxError::LoadSetting` if it is unset. Use `base_url` for a @@ -108,16 +86,17 @@ pub struct VertexProviderSettings { /// are layered over `Authorization`, so they can replace it; in Express /// mode `x-goog-api-key` wins. pub headers: Option>, - /// The OAuth2 access token of standard mode, sent as - /// `Authorization: Bearer`. `None` loads `GOOGLE_VERTEX_ACCESS_TOKEN` when - /// a request is made and fails that request with `AiMuxError::LoadApiKey` - /// if it is unset (unless `headers` carries an `Authorization` header). A - /// [`Resolvable::AsyncFn`] is called on every request, so a host can hand - /// out refreshed tokens. + /// Overrides application default credentials with a token producer. + /// `None` uses `GOOGLE_VERTEX_ACCESS_TOKEN` if set, then Google credentials. pub access_token: Option>, + /// Options for Google application default credentials. + pub google_auth_options: Option, /// The transport: a mock, a signing decorator, a proxy-aware client. /// `None` uses the process default, resolved per request. pub fetch: Option, + /// Custom WebSocket transport for Gemini live transcription. + #[cfg(feature = "realtime")] + pub web_socket: Option>, /// Rewrites every JSON request body once, after it is serialized and /// before it is sent. pub transform_request_body: Option, @@ -133,6 +112,7 @@ impl std::fmt::Debug for VertexProviderSettings { .field("base_url", &self.base_url) .field("headers", &self.headers.is_some()) .field("access_token", &self.access_token) + .field("google_auth_options", &self.google_auth_options) .field("fetch", &self.fetch.is_some()) .field( "transform_request_body", @@ -142,18 +122,6 @@ impl std::fmt::Debug for VertexProviderSettings { } } -fn validate_location(location: &str) -> Result<(), AiMuxError> { - if is_valid_hostname_part(location) { - Ok(()) - } else { - Err(AiMuxError::InvalidArgument( - "Invalid Google Vertex location. Expected a single DNS label (letters, digits, and \ - hyphens). Use `base_url` for custom endpoints." - .to_string(), - )) - } -} - /// The host of a location: `global` is the bare host, `us` and `eu` the /// multi-region `rep` hosts, anything else a regional host. fn location_host(location: &str) -> String { @@ -186,27 +154,30 @@ pub(crate) type ProjectLocationFn = /// go and what authenticates them. struct Resolver { api_key: Option>, + express_mode: bool, location: Option, project: Option, base_url: Option, headers: Option>, access_token: Option>, + auth: auth::GoogleAuth, + use_access_token_env: bool, } impl Resolver { - /// The Express-mode key, when there is a non-empty one. + /// The API key for the mode selected when the provider was created. async fn express_key(&self) -> Result, AiMuxError> { - let key = match &self.api_key { - Some(key) => Some(key.resolve().await?), - None => std::env::var(API_KEY_ENV_VAR).ok(), - }; - Ok(key.filter(|key| !key.trim().is_empty())) + if !self.express_mode { + return Ok(None); + } + match &self.api_key { + Some(key) => Ok(Some(key.resolve().await?)), + None => Ok(None), + } } fn location(&self) -> Result { - let location = load_setting(self.location.as_deref(), LOCATION_ENV_VAR, "location")?; - validate_location(&location)?; - Ok(location) + load_setting(self.location.as_deref(), LOCATION_ENV_VAR, "location") } fn project(&self) -> Result { @@ -231,26 +202,15 @@ impl Resolver { )) } Publisher::Anthropic => { - // Claude is served from the same project/location under the - // `anthropic` publisher; a configured base URL has its Google - // publisher suffix swapped. - let root = if let Some(url) = &self.base_url { - url.strip_suffix("/publishers/google") - .unwrap_or(url) - .to_string() - } else if express { - EXPRESS_MODE_BASE_URL - .strip_suffix("/publishers/google") - .unwrap_or(EXPRESS_MODE_BASE_URL) - .to_string() - } else { - let location = self.location()?; - let project = self.project()?; - format!( - "https://{}/v1/projects/{project}/locations/{location}", - location_host(&location) - ) - }; + if let Some(url) = &self.base_url { + return Ok(url.clone()); + } + let location = self.location()?; + let project = self.project()?; + let root = format!( + "https://{}/v1/projects/{project}/locations/{location}", + location_host(&location) + ); Ok(format!("{root}/publishers/anthropic/models")) } } @@ -258,24 +218,26 @@ impl Resolver { /// The provider headers of a request. async fn headers(&self, express_key: Option<&str>) -> Result { - let user = match &self.headers { - Some(headers) => headers.resolve().await?, - None => HeaderMapOpt::new(), - }; if let Some(key) = express_key { + let user = match &self.headers { + Some(headers) => headers.resolve().await?, + None => HeaderMapOpt::new(), + }; let mut fixed = HeaderMapOpt::new(); fixed.insert("x-goog-api-key".to_string(), Some(key.to_string())); return Ok(combine_headers(&[&user, &fixed])); } - let user_authorizes = user - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") && value.is_some()); - if user_authorizes { - return Ok(user); - } let token = match &self.access_token { Some(token) => token.resolve().await?, - None => load_api_key(None, ACCESS_TOKEN_ENV_VAR, "Google Vertex access token")?, + None if self.use_access_token_env => match std::env::var(ACCESS_TOKEN_ENV_VAR) { + Ok(token) => token, + Err(_) => self.auth.access_token().await?, + }, + None => self.auth.access_token().await?, + }; + let user = match &self.headers { + Some(headers) => headers.resolve().await?, + None => HeaderMapOpt::new(), }; let mut auth = HeaderMapOpt::new(); auth.insert("Authorization".to_string(), Some(format!("Bearer {token}"))); @@ -284,7 +246,10 @@ impl Resolver { /// The endpoint of `publisher`'s models for one request. async fn endpoint(&self, publisher: Publisher) -> Result { - let express_key = self.express_key().await?; + let express_key = match publisher { + Publisher::Google => self.express_key().await?, + Publisher::Anthropic => None, + }; let base_url = self.base_url(express_key.is_some(), publisher)?; let headers = self.headers(express_key.as_deref()).await?; let headers = aimux_provider_utils::headers::with_user_agent_suffix_fn( @@ -302,38 +267,47 @@ impl Resolver { /// /// # Errors /// -/// Returns `AiMuxError::InvalidArgument` when an explicit `base_url` is not an -/// `http(s)` URL with a host, or an explicit `location` is not a single DNS -/// label. Those are the only ways this fails: keys, tokens, and a project or -/// location taken from the environment are loaded per request, not here. +/// Credentials, project and location errors are reported by model requests. pub fn create_google_vertex( settings: VertexProviderSettings, ) -> Result { let base_url = settings .base_url - .as_deref() - .map(validate_base_url) - .transpose()?; - if let Some(location) = settings.location.as_deref() { - validate_location(location)?; + .map(|url| url.trim_end_matches('/').to_string()); + let api_key = settings + .api_key + .or_else(|| std::env::var(API_KEY_ENV_VAR).ok().map(Resolvable::Value)); + let express_mode = match &api_key { + Some(Resolvable::Value(key)) => !key.is_empty(), + Some(_) => true, + None => false, + }; + let mut google_auth_options = settings.google_auth_options.unwrap_or_default(); + if google_auth_options.project_id.is_none() { + google_auth_options.project_id = settings.project.clone(); } Ok(VertexProvider { resolver: Arc::new(Resolver { - api_key: settings.api_key, + api_key, + express_mode, location: settings.location, project: settings.project, base_url, headers: settings.headers, access_token: settings.access_token, + use_access_token_env: true, + auth: auth::GoogleAuth::new(google_auth_options, settings.fetch.clone()), }), fetch: settings.fetch, + #[cfg(feature = "realtime")] + web_socket: settings.web_socket, transform_request_body: settings.transform_request_body, }) } /// The default provider: `create_google_vertex` with default settings, created -/// on first use. Creating it reads nothing from the environment and cannot -/// fail; a missing key, token, project or location surfaces from the first +/// on first use. It captures the API-key environment setting; a missing +/// token, project or location surfaces from the first /// request instead. pub fn google_vertex() -> &'static VertexProvider { static DEFAULT: OnceLock = OnceLock::new(); @@ -348,6 +322,8 @@ pub fn google_vertex() -> &'static VertexProvider { pub struct VertexProvider { resolver: Arc, fetch: Option, + #[cfg(feature = "realtime")] + web_socket: Option>, transform_request_body: Option, } @@ -383,7 +359,15 @@ impl VertexProvider { .to_string(), )); } - resolver.endpoint(Publisher::Google).await + let mut endpoint = resolver.endpoint(Publisher::Google).await?; + if tuned && resolver.base_url.is_none() { + endpoint.base_url = endpoint + .base_url + .strip_suffix("/publishers/google") + .unwrap_or(&endpoint.base_url) + .to_string(); + } + Ok(endpoint) }) }), fetch: self.fetch.clone(), @@ -413,7 +397,7 @@ impl VertexProvider { } /// A Gemini model (e.g. `"gemini-2.0-flash"`, or a tuned - /// `"endpoints/{id}"`); `provider()` is `"google.vertex"`. + /// `"endpoints/{id}"`); `provider()` is `"google.vertex.chat"`. /// /// A tuned model fails its requests with `AiMuxError::InvalidArgument` in /// Express mode. @@ -421,22 +405,31 @@ impl VertexProvider { pub fn chat(&self, model_id: &str) -> VertexModel { VertexModel::from_config( model_id.to_string(), - self.model_config(PROVIDER, model_id.starts_with(ENDPOINT_MODEL_PREFIX)), + self.model_config( + "google.vertex.chat", + model_id.starts_with(ENDPOINT_MODEL_PREFIX), + ), ) } /// An embedding model (e.g. `"textembedding-gecko@001"`); `provider()` is - /// `"google.vertex"`. + /// `"google.vertex.embedding"`. #[must_use] pub fn embedding(&self, model_id: &str) -> VertexEmbeddingModel { - VertexEmbeddingModel::from_config(model_id.to_string(), self.model_config(PROVIDER, false)) + VertexEmbeddingModel::from_config( + model_id.to_string(), + self.model_config("google.vertex.embedding", false), + ) } /// An image generation model (e.g. `"imagen-4.0-generate-001"` or - /// `"gemini-2.5-flash-image"`); `provider()` is `"google.vertex"`. + /// `"gemini-2.5-flash-image"`); `provider()` is `"google.vertex.image"`. #[must_use] pub fn image(&self, model_id: &str) -> VertexImageModel { - VertexImageModel::from_config(model_id.to_string(), self.model_config(PROVIDER, false)) + VertexImageModel::from_config( + model_id.to_string(), + self.model_config("google.vertex.image", false), + ) } /// A video generation model (e.g. `"veo-3.0-generate-001"`); @@ -449,35 +442,25 @@ impl VertexProvider { ) } - /// A Speech-to-Text transcription model (e.g. `"chirp_3"`); `provider()` + /// A Gemini or Speech-to-Text transcription model; `provider()` /// is `"google.vertex.transcription"`. Its requests need a project and a /// location and fail in Express mode. #[must_use] - pub fn transcription(&self, model_id: &str) -> VertexTranscriptionModel { - VertexTranscriptionModel::from_config( + pub fn transcription(&self, model_id: &str) -> Arc { + if model_id.starts_with("gemini") { + return Arc::new(VertexGeminiTranscriptionModel::from_config( + model_id.to_string(), + self.model_config("google.vertex.transcription", false), + self.project_location("transcription"), + #[cfg(feature = "realtime")] + self.web_socket.clone(), + )); + } + Arc::new(VertexTranscriptionModel::from_config( model_id.to_string(), self.model_config("google.vertex.transcription", false), self.project_location("transcription"), - ) - } - - /// A Claude model served by Vertex AI through `rawPredict` (e.g. - /// `"claude-sonnet-4-20250514"`); `provider()` is - /// `"googleVertex.anthropic.messages"`. It uses this provider's - /// credentials, transport and location, and the shared Anthropic Messages - /// model. - #[must_use] - pub fn anthropic_model(&self, model_id: &str) -> VertexAnthropicModel { - let resolver = self.resolver.clone(); - anthropic_model::model( - model_id, - Arc::new(move || { - let resolver = resolver.clone(); - Box::pin(async move { resolver.endpoint(Publisher::Anthropic).await }) - }), - self.fetch.clone(), - self.transform_request_body.clone(), - ) + )) } /// The provider as a function: the default language model for an id. The @@ -509,7 +492,7 @@ impl Provider for VertexProvider { &self, model_id: &str, ) -> Option, AiMuxError>> { - Some(Ok(Arc::new(self.transcription(model_id)))) + Some(Ok(self.transcription(model_id))) } fn video_model(&self, model_id: &str) -> Option, AiMuxError>> { diff --git a/aimux-providers/src/vertex/model.rs b/aimux-providers/src/vertex/model.rs index 4d868bd89..08c7ac2ac 100644 --- a/aimux-providers/src/vertex/model.rs +++ b/aimux-providers/src/vertex/model.rs @@ -16,21 +16,22 @@ use aimux_core::options::CallOptions; use aimux_core::result::{ GenerateContent, GenerateResult, GeneratedFile, ReasoningOutput, Source, StreamResult, }; -use aimux_core::shared::{FileBytes, GeneratedFileData}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{ - FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, + FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, Warning, }; use crate::google::convert::{ - build_vertex_request_body, code_execution_tool_name, convert_usage, extract_sources, - parse_finish_reason, + build_vertex_request_body, code_execution_tool_name, extract_sources, parse_finish_reason, + prepare_all_tools, validate_call_options_for_namespace, }; use crate::google::types::{Candidate, GenerateContentResponse, GoogleStreamEvent}; +use crate::google::utils::{GoogleJsonAccumulator, PartialArg}; -use crate::google::options::Namespace; +use crate::google::options::{GOOGLE, Namespace}; use crate::shared::EndpointConfig; use aimux_core::language_model::SupportedUrls; +use aimux_core::shared::{FileBytes, GeneratedFileData}; /// A Google Vertex AI language model. /// @@ -46,19 +47,6 @@ impl VertexModel { Self { model_id, config } } - /// The effective base URL for this model. Tuned models addressed via - /// `endpoints/{id}` are served from `…/locations/{region}/endpoints/{id}` - /// (no `/publishers/google` suffix), so the suffix is stripped from the - /// resolved base URL. Mirrors the TS `loadBaseURL({ endpoint: true })`. - fn effective_base_url<'a>(&self, base_url: &'a str) -> &'a str { - if self.model_id.starts_with("endpoints/") - && let Some(stripped) = base_url.strip_suffix("/publishers/google") - { - return stripped; - } - base_url - } - fn model_path(&self) -> String { if self.model_id.contains('/') { self.model_id.clone() @@ -69,18 +57,14 @@ impl VertexModel { /// `…/models/{model}:generateContent` fn generate_endpoint(&self, base_url: &str) -> String { - format!( - "{}/{}:generateContent", - self.effective_base_url(base_url), - self.model_path() - ) + format!("{}/{}:generateContent", base_url, self.model_path()) } /// `…/models/{model}:streamGenerateContent?alt=sse` fn stream_endpoint(&self, base_url: &str) -> String { format!( "{}/{}:streamGenerateContent?alt=sse", - self.effective_base_url(base_url), + base_url, self.model_path() ) } @@ -101,74 +85,10 @@ impl LanguageModel for VertexModel { } async fn do_generate(&self, options: &CallOptions) -> Result { - if options - .prompt - .iter() - .filter_map(|message| match message { - aimux_core::language_model_message::LanguageModelMessage::Assistant { - content, - .. - } => Some(content), - _ => None, - }) - .flatten() - .any(|part| { - matches!( - part, - aimux_core::language_model_message::AssistantPart::ReasoningFile( - aimux_core::language_model_message::ReasoningFilePart { - data: aimux_core::shared::GeneratedFileData::Url { .. }, - .. - } - ) | aimux_core::language_model_message::AssistantPart::File( - aimux_core::language_model_message::FilePart { - data: aimux_core::shared::FileData::Url { .. }, - .. - } - ) - ) - }) - { - return Err(AiMuxError::UnsupportedFunctionality( - "File data URLs in assistant messages are not supported".to_string(), - )); - } - if options.prompt.iter().any(|message| match message { - aimux_core::language_model_message::LanguageModelMessage::Assistant { - content, .. - } => content.iter().any(|part| { - matches!( - part, - aimux_core::language_model_message::AssistantPart::File( - aimux_core::language_model_message::FilePart { - data: aimux_core::shared::FileData::Reference { .. }, - .. - } - ) - ) - }), - aimux_core::language_model_message::LanguageModelMessage::User { content, .. } => { - content.iter().any(|part| { - matches!( - part, - aimux_core::language_model_message::UserPart::File( - aimux_core::language_model_message::FilePart { - data: aimux_core::shared::FileData::Reference { .. }, - .. - } - ) - ) - }) - } - _ => false, - }) { - return Err(AiMuxError::UnsupportedFunctionality( - "file parts with provider references".to_string(), - )); - } - crate::google::convert::validate_tool_result_files(&options.prompt)?; + let options = vertex_call_options(options)?; + let options = &options; let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); - let body = build_vertex_request_body(&self.model_id, options)?; + let (body, warnings) = vertex_request_body(&self.model_id, options, false)?; let exchange = self.config.exchange(options.headers.as_ref()).await?; let body = exchange.transform_body(body); let resp = aimux_provider_utils::post_json_to_api( @@ -182,36 +102,50 @@ impl LanguageModel for VertexModel { let response_body = resp.raw_value; let response_headers = resp.response_headers; - let data: GenerateContentResponse = resp.value; - - let candidate = data.candidates.into_iter().next().ok_or_else(|| { - AiMuxError::InvalidResponseData("no candidates in response".to_string()) - })?; + let raw: Value = resp.value; + let mut parsed_response = raw.clone(); + if parsed_response + .get("candidates") + .is_some_and(Value::is_null) + { + parsed_response["candidates"] = json!([]); + } + let data: GenerateContentResponse = serde_json::from_value(parsed_response) + .map_err(|error| AiMuxError::InvalidResponseData(error.to_string()))?; + let candidate = data.candidates.into_iter().next().unwrap_or(Candidate { + content: None, + finish_reason: None, + finish_message: None, + safety_ratings: None, + grounding_metadata: None, + url_context_metadata: None, + index: None, + }); let (content, has_tool_calls) = extract_content_from_candidate(&candidate, &code_execution_tool_name); - + let block_reason = confirmed_block_reason(data.prompt_feedback.as_ref()); let finish_reason = candidate .finish_reason .as_deref() .map(|r| parse_finish_reason(r, has_tool_calls)) .unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, + unified: if block_reason.is_some() { + FinishReasonUnified::ContentFilter + } else { + FinishReasonUnified::Other + }, + raw: block_reason, }); - - let usage = data - .usage_metadata - .as_ref() - .map(convert_usage) - .unwrap_or_default(); + let usage = vertex_usage(raw.get("usageMetadata")); let provider_metadata = Some(vertex_provider_metadata(json!({ "promptFeedback": data.prompt_feedback, "groundingMetadata": candidate.grounding_metadata, "urlContextMetadata": candidate.url_context_metadata, "safetyRatings": candidate.safety_ratings, - "usageMetadata": data.usage_metadata, + "usageMetadata": raw.get("usageMetadata"), + "serviceTier": raw.pointer("/usageMetadata/serviceTier"), "finishMessage": candidate.finish_message, }))); @@ -219,14 +153,12 @@ impl LanguageModel for VertexModel { content, finish_reason, usage, - warnings: Vec::new(), + warnings, provider_metadata, response: Some(aimux_core::shared::ResponseInfo { id: data.response_id, - // Gemini responses carry no timestamp field; use the response - // `Date` header (RFC1123) like Bedrock does. - timestamp: response_headers.get("date").cloned(), - model_id: data.model_version, + timestamp: None, + model_id: None, headers: Some(response_headers), body: response_body, }), @@ -235,116 +167,40 @@ impl LanguageModel for VertexModel { } async fn do_stream(&self, options: &CallOptions) -> Result { - if options - .prompt - .iter() - .filter_map(|message| match message { - aimux_core::language_model_message::LanguageModelMessage::Assistant { - content, - .. - } => Some(content), - _ => None, - }) - .flatten() - .any(|part| { - matches!( - part, - aimux_core::language_model_message::AssistantPart::ReasoningFile( - aimux_core::language_model_message::ReasoningFilePart { - data: aimux_core::shared::GeneratedFileData::Url { .. }, - .. - } - ) | aimux_core::language_model_message::AssistantPart::File( - aimux_core::language_model_message::FilePart { - data: aimux_core::shared::FileData::Url { .. }, - .. - } - ) - ) - }) - { - return Err(AiMuxError::UnsupportedFunctionality( - "File data URLs in assistant messages are not supported".to_string(), - )); - } - if options.prompt.iter().any(|message| match message { - aimux_core::language_model_message::LanguageModelMessage::Assistant { - content, .. - } => content.iter().any(|part| { - matches!( - part, - aimux_core::language_model_message::AssistantPart::File( - aimux_core::language_model_message::FilePart { - data: aimux_core::shared::FileData::Reference { .. }, - .. - } - ) - ) - }), - aimux_core::language_model_message::LanguageModelMessage::User { content, .. } => { - content.iter().any(|part| { - matches!( - part, - aimux_core::language_model_message::UserPart::File( - aimux_core::language_model_message::FilePart { - data: aimux_core::shared::FileData::Reference { .. }, - .. - } - ) - ) - }) - } - _ => false, - }) { - return Err(AiMuxError::UnsupportedFunctionality( - "file parts with provider references".to_string(), - )); - } - crate::google::convert::validate_tool_result_files(&options.prompt)?; + let options = vertex_call_options(options)?; + let options = &options; let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); - let body = build_vertex_request_body(&self.model_id, options)?; + let (body, warnings) = vertex_request_body(&self.model_id, options, true)?; + let include_raw_chunks = options.include_raw_chunks == Some(true); let exchange = self.config.exchange(options.headers.as_ref()).await?; let body = exchange.transform_body(body); let endpoint = self.stream_endpoint(exchange.base_url()); let resp = aimux_provider_utils::post_json_to_api( exchange.request(endpoint.clone(), options), body.clone(), - aimux_provider_utils::create_event_source_response_handler::(), + aimux_provider_utils::create_event_source_response_handler::(), crate::google::google_failed_response_handler(), ) .await?; let response_headers = resp.response_headers; - // Same source as the non-stream path: the response `Date` header. - let response_timestamp = response_headers.get("date").cloned(); - - let mut sse_stream = resp.value; - let first_event = match sse_stream.next().await { - Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), - first_event => first_event, - }; - if let Some(Ok(GoogleStreamEvent::Error(error))) = first_event.as_ref() { - return Err(crate::google::google_stream_error( - &error.error, - &endpoint, - body.clone(), - response_headers, - )); - } + let sse_stream = resp.value; let stream_error_url = endpoint; let stream_request_body = body.clone(); let stream_error_headers = response_headers.clone(); let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings: vec![] }); + yield Ok(StreamPart::StreamStart { warnings }); - let mut sse_stream = futures::stream::iter(first_event.into_iter()).chain(sse_stream); + let mut sse_stream = sse_stream; let mut text_id: Option = None; let mut reasoning_id: Option = None; + let mut prompt_blocked = false; let mut block_counter = 0usize; let mut final_usage: Usage = Usage::default(); let mut final_finish_reason: Option = None; let mut has_tool_calls = false; + let mut active_tool_calls: Vec<(RawToolCall, GoogleJsonAccumulator)> = Vec::new(); let mut response_metadata_emitted = false; let mut stream_errored = false; @@ -372,6 +228,18 @@ impl LanguageModel for VertexModel { if stream_errored { break; } + let event = match event { + Ok(raw) => { + if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: raw.clone() }); } + if let Some(usage) = raw.get("usageMetadata") { + final_usage = vertex_usage(Some(usage)); + last_usage_metadata_value = Some(usage.clone()); + } + serde_json::from_value::(raw) + .map_err(|error| AiMuxError::InvalidResponseData(error.to_string())) + } + Err(error) => Err(error), + }; match event { Ok(GoogleStreamEvent::Chunk(chunk)) => { @@ -380,17 +248,17 @@ impl LanguageModel for VertexModel { response_metadata_emitted = true; yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { id: Some(id.clone()), - timestamp: response_timestamp.clone(), - model_id: chunk.model_version.clone(), + timestamp: None, + model_id: None, })); } - if let Some(usage) = &chunk.usage_metadata { - final_usage = convert_usage(usage); - last_usage_metadata_value = serde_json::to_value(usage).ok(); - } - if let Some(pf) = &chunk.prompt_feedback { + if !prompt_blocked && let Some(pf) = &chunk.prompt_feedback { last_prompt_feedback = Some(pf.clone()); + if let Some(reason) = confirmed_block_reason(Some(pf)) { + prompt_blocked = true; + final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::ContentFilter, raw: Some(reason) }); + } } let Some(candidates) = chunk.candidates else { @@ -407,6 +275,15 @@ impl LanguageModel for VertexModel { last_url_context_metadata = Some(ucm.clone()); } + if let Some(sr) = &candidate.safety_ratings { + last_safety_ratings = serde_json::to_value(sr).ok(); + } + if let Some(fm) = &candidate.finish_message { + last_finish_message = Some(json!(fm)); + } + + if prompt_blocked { continue; } + // Extract url sources from this chunk's grounding metadata // (deduplicated across chunks; document sources are not // emitted in the stream, matching TS). @@ -432,88 +309,103 @@ impl LanguageModel for VertexModel { if let Some(parts) = candidate.content.as_ref().and_then(|c| c.parts.as_ref()) { - for part in parts { - let thought_signature_metadata = part - .get("thoughtSignature") - .and_then(|v| v.as_str()) - .filter(|signature| !signature.is_empty()) - .map(|signature| { - vertex_provider_metadata(json!({ - "thoughtSignature": signature, - })) - }); - if let Some(text) = part.get("text").and_then(|v| v.as_str()) { + for part in parts.iter().filter(|part| part.get("functionCall").is_none()) + .chain(parts.iter().filter(|part| part.get("functionCall").is_some())) { + if let Some(text) = part.get("text").and_then(Value::as_str) { + let metadata = thought_metadata(part); if text.is_empty() { - if let (Some(metadata), Some(id)) = (&thought_signature_metadata, &text_id) { - yield Ok(StreamPart::TextDelta { - id: id.clone(), - delta: String::new(), - provider_metadata: Some(metadata.clone()), - }); - } - } else if part.get("thought").and_then(Value::as_bool) == Some(true) { - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); + if metadata.is_some() && let Some(id) = &text_id { + yield Ok(StreamPart::TextDelta { id: id.clone(), delta: String::new(), provider_metadata: metadata }); } + } else if part.get("thought") == Some(&Value::Bool(true)) { + if let Some(id) = text_id.take() { yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); } if reasoning_id.is_none() { - let id = format!("{block_counter}"); - block_counter += 1; + let id = format!("{block_counter}"); block_counter += 1; reasoning_id = Some(id.clone()); - yield Ok(StreamPart::ReasoningStart { - id, - provider_metadata: thought_signature_metadata.clone(), - }); - } - if let Some(id) = &reasoning_id { - yield Ok(StreamPart::ReasoningDelta { - id: id.clone(), - delta: text.to_string(), - provider_metadata: thought_signature_metadata, - }); + yield Ok(StreamPart::ReasoningStart { id, provider_metadata: metadata.clone() }); } + yield Ok(StreamPart::ReasoningDelta { id: reasoning_id.clone().unwrap(), delta: text.to_string(), provider_metadata: metadata }); } else { - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); - } + if let Some(id) = reasoning_id.take() { yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); } if text_id.is_none() { - let id = format!("{block_counter}"); - block_counter += 1; + let id = format!("{block_counter}"); block_counter += 1; text_id = Some(id.clone()); - yield Ok(StreamPart::TextStart { - id, - provider_metadata: thought_signature_metadata.clone(), - }); - } - if let Some(id) = &text_id { - yield Ok(StreamPart::TextDelta { - id: id.clone(), - delta: text.to_string(), - provider_metadata: thought_signature_metadata, - }); + yield Ok(StreamPart::TextStart { id, provider_metadata: metadata.clone() }); } + yield Ok(StreamPart::TextDelta { id: text_id.clone().unwrap(), delta: text.to_string(), provider_metadata: metadata }); } + } else if let Some(file) = inline_file(part) { + if let Some(id) = text_id.take() { yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); } + if let Some(id) = reasoning_id.take() { yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); } + yield Ok(if part.get("thought") == Some(&Value::Bool(true)) { StreamPart::ReasoningFile(file) } else { StreamPart::File(file) }); } else if let Some(fc) = part.get("functionCall") { - let name = fc - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or(""); - let id = fc - .get("id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - .unwrap_or_else(|| format!("call-{block_counter}")); - block_counter += 1; - let args = fc.get("args").cloned().unwrap_or(json!({})); - let thought_signature = part - .get("thoughtSignature") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); + let name = fc.get("name").and_then(Value::as_str); + let args = fc.get("args").filter(|value| !value.is_null()); + let partial_args = fc.get("partialArgs").filter(|value| !value.is_null()); + let will_continue = fc.get("willContinue").and_then(Value::as_bool) == Some(true); + let thought_signature = part.get("thoughtSignature") + .and_then(Value::as_str).filter(|signature| !signature.is_empty()) + .map(str::to_owned); let tool_metadata = thought_signature.as_deref().map(|signature| { - vertex_provider_metadata(json!({ - "thoughtSignature": signature, - })) + vertex_provider_metadata(json!({ "thoughtSignature": signature })) }); - + if partial_args.is_some() || (name.is_some() && will_continue) { + if let Some(name) = name { + let id = fc.get("id").and_then(Value::as_str) + .filter(|id| !id.is_empty()).map(str::to_owned) + .unwrap_or_else(aimux_provider_utils::generate_id); + yield Ok(StreamPart::ToolInputStart { + id: id.clone(), tool_name: name.to_owned(), + provider_executed: None, dynamic: None, title: None, + provider_metadata: tool_metadata.clone(), + }); + active_tool_calls.push((RawToolCall { + tool_call_id: id, tool_name: name.to_owned(), input: String::new(), + provider_executed: None, dynamic: None, + provider_metadata: tool_metadata.clone(), + }, GoogleJsonAccumulator::new())); + } + if let Some(partial_args) = partial_args.and_then(Value::as_array) + && let Some((call, accumulator)) = active_tool_calls.last_mut() { + let partial_args: Vec = partial_args.iter().map(|arg| PartialArg { + json_path: arg.get("jsonPath").and_then(Value::as_str).unwrap_or_default().to_owned(), + string_value: arg.get("stringValue").and_then(Value::as_str).map(str::to_owned), + number_value: arg.get("numberValue").and_then(Value::as_f64), + bool_value: arg.get("boolValue").and_then(Value::as_bool), + null_value: arg.get("nullValue").map(|_| ()), + will_continue: arg.get("willContinue").and_then(Value::as_bool), + }).collect(); + match accumulator.process_partial_args(&partial_args) { + Ok(result) => { + if !result.text_delta.is_empty() { + yield Ok(StreamPart::ToolInputDelta { + id: call.tool_call_id.clone(), delta: result.text_delta, + provider_metadata: tool_metadata, + }); + } + } + Err(error) => { yield Ok(StreamPart::Error { error }); return; } + } + if !will_continue && partial_args.iter().all(|arg| arg.will_continue != Some(true)) { + for event in finish_streaming_tool_call(&mut active_tool_calls) { yield Ok(event); } + has_tool_calls = true; + } + } + continue; + } + if name.is_none() && args.is_none() && partial_args.is_none() + && fc.get("willContinue").is_none_or(Value::is_null) { + if !active_tool_calls.is_empty() { + for event in finish_streaming_tool_call(&mut active_tool_calls) { yield Ok(event); } + has_tool_calls = true; + } + continue; + } + let Some(name) = name else { continue; }; + let id = fc.get("id").and_then(Value::as_str) + .filter(|id| !id.is_empty()).map(str::to_owned) + .unwrap_or_else(aimux_provider_utils::generate_id); + let args = args.cloned().unwrap_or(json!({})); yield Ok(StreamPart::ToolInputStart { id: id.clone(), tool_name: name.to_string(), @@ -522,12 +414,13 @@ impl LanguageModel for VertexModel { title: None, provider_metadata: tool_metadata.clone(), }); - let args_str = args.to_string(); - yield Ok(StreamPart::ToolInputDelta { + let args_str = args.as_str().map(str::to_string).unwrap_or_else(|| args.to_string()); + if fc.get("args").is_some_and(|value| !value.is_null()) { yield Ok(StreamPart::ToolInputDelta { id: id.clone(), - delta: args_str, + delta: args_str.clone(), provider_metadata: tool_metadata.clone(), }); + } yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: tool_metadata.clone(), @@ -535,7 +428,7 @@ impl LanguageModel for VertexModel { yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id, tool_name: name.to_string(), - input: args.to_string(), + input: args_str, provider_executed: None, dynamic: None, provider_metadata: tool_metadata, @@ -555,7 +448,7 @@ impl LanguageModel for VertexModel { let file = GeneratedFile { data: GeneratedFileData::Data { data: FileBytes::Base64(data.to_string()) }, media_type: media_type.to_string(), - provider_metadata: thought_signature_metadata, + provider_metadata: thought_metadata(part), }; yield Ok(if part.get("thought").and_then(Value::as_bool) == Some(true) { StreamPart::ReasoningFile(file) @@ -626,8 +519,9 @@ impl LanguageModel for VertexModel { let id = tc .get("id") .and_then(|v| v.as_str()) + .filter(|id| !id.is_empty()) .map(std::string::ToString::to_string) - .unwrap_or_else(|| format!("call-{block_counter}")); + .unwrap_or_else(aimux_provider_utils::generate_id); block_counter += 1; last_server_tool_call_id = Some(id.clone()); let args = tc.get("args").cloned().unwrap_or(json!({})); @@ -685,20 +579,6 @@ impl LanguageModel for VertexModel { } if let Some(reason) = candidate.finish_reason.as_deref() { - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); - } - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); - } - // Snapshot the finishReason-chunk metadata. - if let Some(sr) = &candidate.safety_ratings { - last_safety_ratings = - Some(serde_json::to_value(sr).unwrap_or(Value::Null)); - } - if let Some(fm) = &candidate.finish_message { - last_finish_message = Some(json!(fm)); - } final_finish_reason = Some(parse_finish_reason(reason, has_tool_calls)); } @@ -717,7 +597,7 @@ impl LanguageModel for VertexModel { } Err(error) => { let recoverable = error.is_recoverable_stream_error(); - yield Err(error); + yield Ok(StreamPart::Error { error }); if !recoverable { return; } @@ -732,12 +612,12 @@ impl LanguageModel for VertexModel { if let Some(id) = reasoning_id.take() { yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); } - let provider_metadata = Some(vertex_provider_metadata(json!({ "promptFeedback": last_prompt_feedback, "groundingMetadata": last_grounding_metadata, "urlContextMetadata": last_url_context_metadata, "safetyRatings": last_safety_ratings, + "serviceTier": last_usage_metadata_value.as_ref().and_then(|usage| usage.get("serviceTier")), "usageMetadata": last_usage_metadata_value, "finishMessage": last_finish_message, }))); @@ -750,7 +630,7 @@ impl LanguageModel for VertexModel { } } else { final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Stop, + unified: FinishReasonUnified::Other, raw: None, }) }, @@ -771,7 +651,36 @@ impl LanguageModel for VertexModel { // ── Helpers ────────────────────────────────────────────────────────────────── -fn vertex_provider_metadata(payload: Value) -> ProviderMetadata { +fn finish_streaming_tool_call( + active_tool_calls: &mut Vec<(RawToolCall, GoogleJsonAccumulator)>, +) -> Vec { + let Some((mut call, accumulator)) = active_tool_calls.pop() else { + return Vec::new(); + }; + let result = accumulator.finalize(); + let mut events = Vec::new(); + if !result.closing_delta.is_empty() { + events.push(StreamPart::ToolInputDelta { + id: call.tool_call_id.clone(), + delta: result.closing_delta, + provider_metadata: call.provider_metadata.clone(), + }); + } + events.push(StreamPart::ToolInputEnd { + id: call.tool_call_id.clone(), + provider_metadata: call.provider_metadata.clone(), + }); + call.input = result.final_json; + events.push(StreamPart::ToolCall(call)); + events +} + +fn vertex_provider_metadata(mut payload: Value) -> ProviderMetadata { + if let Some(object) = payload.as_object_mut() + && object.get("serviceTier").is_some_and(Value::is_null) + { + object.remove("serviceTier"); + } Namespace::Vertex.metadata(payload) } @@ -795,7 +704,7 @@ fn extract_content_from_candidate( candidate: &Candidate, code_execution_tool_name: &str, ) -> (Vec, bool) { - let mut content = Vec::new(); + let mut content: Vec = Vec::new(); let mut has_tool_calls = false; let mut source_id = 0usize; let mut last_code_execution_tool_call_id: Option = None; @@ -849,52 +758,67 @@ fn extract_content_from_candidate( })); } } else if let Some(text) = part.get("text").and_then(|v| v.as_str()) { - let provider_metadata = part - .get("thoughtSignature") - .and_then(Value::as_str) - .filter(|signature| !signature.is_empty()) - .map(|signature| { - vertex_provider_metadata(json!({ "thoughtSignature": signature })) + let provider_metadata = thought_metadata(part); + if !text.is_empty() { + content.push(if part.get("thought") == Some(&Value::Bool(true)) { + GenerateContent::Reasoning(ReasoningOutput { + text: text.to_string(), + provider_metadata, + }) + } else { + GenerateContent::Text { + text: text.to_string(), + provider_metadata, + } }); - if text.is_empty() { - if let ( - Some(metadata), - Some( - GenerateContent::Text { - provider_metadata, .. + } else if provider_metadata.is_some() + && let Some(last) = content.last_mut() + { + match last { + GenerateContent::Text { + provider_metadata: metadata, + .. + } => *metadata = provider_metadata, + GenerateContent::Reasoning(reasoning) => { + reasoning.provider_metadata = provider_metadata + } + GenerateContent::File(file) | GenerateContent::ReasoningFile(file) => { + file.provider_metadata = provider_metadata + } + GenerateContent::Custom { + provider_metadata: metadata, + .. + } => *metadata = provider_metadata, + GenerateContent::ToolApprovalRequest(request) => { + request.provider_metadata = provider_metadata + } + GenerateContent::ToolCall(call) => { + call.provider_metadata = provider_metadata + } + GenerateContent::ToolResult(result) => { + result.provider_metadata = provider_metadata + } + GenerateContent::Source( + Source::Url { + provider_metadata: metadata, + .. } - | GenerateContent::Reasoning(ReasoningOutput { - provider_metadata, .. - }) - | GenerateContent::File(GeneratedFile { - provider_metadata, .. - }) - | GenerateContent::ReasoningFile(GeneratedFile { - provider_metadata, .. - }) - | GenerateContent::ToolCall(RawToolCall { - provider_metadata, .. - }) - | GenerateContent::ToolResult(ToolResult { - provider_metadata, .. - }), - ), - ) = (provider_metadata, content.last_mut()) - { - *provider_metadata = Some(metadata); + | Source::Document { + provider_metadata: metadata, + .. + }, + ) => *metadata = provider_metadata, } - } else if part.get("thought").and_then(Value::as_bool) == Some(true) { - content.push(GenerateContent::Reasoning(ReasoningOutput { - text: text.to_string(), - provider_metadata, - })); - } else { - content.push(GenerateContent::Text { - text: text.to_string(), - provider_metadata, - }); } - } else if let Some(fc) = part.get("functionCall") { + } else if let Some(file) = inline_file(part) { + content.push(if part.get("thought") == Some(&Value::Bool(true)) { + GenerateContent::ReasoningFile(file) + } else { + GenerateContent::File(file) + }); + } else if let Some(fc) = part.get("functionCall") + && fc.get("name").is_some_and(|name| !name.is_null()) + { let name = fc .get("name") .and_then(|v| v.as_str()) @@ -903,9 +827,14 @@ fn extract_content_from_candidate( let id = fc .get("id") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let input = fc.get("args").cloned().unwrap_or(json!({})); + .filter(|id| !id.is_empty()) + .map(str::to_string) + .unwrap_or_else(aimux_provider_utils::generate_id); + let input = fc + .get("args") + .filter(|value| !value.is_null()) + .cloned() + .unwrap_or(json!({})); let thought_signature = part .get("thoughtSignature") .and_then(|v| v.as_str()) @@ -922,32 +851,6 @@ fn extract_content_from_candidate( provider_metadata, })); has_tool_calls = true; - } else if let Some(inline) = part.get("inlineData") { - if let (Some(data), Some(media_type)) = ( - inline.get("data").and_then(Value::as_str), - inline.get("mimeType").and_then(Value::as_str), - ) { - let file = GeneratedFile { - data: GeneratedFileData::Data { - data: FileBytes::Base64(data.to_string()), - }, - media_type: media_type.to_string(), - provider_metadata: part - .get("thoughtSignature") - .and_then(Value::as_str) - .filter(|signature| !signature.is_empty()) - .map(|signature| { - vertex_provider_metadata(json!({ "thoughtSignature": signature })) - }), - }; - content.push( - if part.get("thought").and_then(Value::as_bool) == Some(true) { - GenerateContent::ReasoningFile(file) - } else { - GenerateContent::File(file) - }, - ); - } } else if let Some(tc) = part.get("toolCall") { let tool_type = tc.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); let id = tc @@ -1006,3 +909,140 @@ fn extract_content_from_candidate( (content, has_tool_calls) } + +fn thought_metadata(part: &Value) -> Option { + part.get("thoughtSignature") + .and_then(Value::as_str) + .filter(|signature| !signature.is_empty()) + .map(|signature| vertex_provider_metadata(json!({ "thoughtSignature": signature }))) +} + +fn inline_file(part: &Value) -> Option { + let inline = part.get("inlineData")?; + Some(GeneratedFile { + media_type: inline.get("mimeType")?.as_str()?.to_string(), + data: GeneratedFileData::Data { + data: FileBytes::Base64(inline.get("data")?.as_str()?.to_string()), + }, + provider_metadata: thought_metadata(part), + }) +} + +fn confirmed_block_reason(feedback: Option<&Value>) -> Option { + feedback? + .get("blockReason")? + .as_str() + .filter(|reason| { + !matches!( + *reason, + "" | "BLOCK_REASON_UNSPECIFIED" | "BLOCKED_REASON_UNSPECIFIED" + ) + }) + .map(str::to_string) +} + +fn vertex_usage(raw: Option<&Value>) -> Usage { + use aimux_core::types::{InputTokenUsage, OutputTokenUsage}; + let Some(raw) = raw.filter(|value| !value.is_null()) else { + return Usage::default(); + }; + let count = |key: &str| raw.get(key).and_then(Value::as_u64).unwrap_or(0) as u32; + let input = count("promptTokenCount") + count("toolUsePromptTokenCount"); + let cached = count("cachedContentTokenCount"); + let text = count("candidatesTokenCount"); + let reasoning = count("thoughtsTokenCount"); + Usage { + input_tokens: InputTokenUsage { + total: Some(input), + no_cache: Some(input.saturating_sub(cached)), + cache_read: Some(cached), + ..Default::default() + }, + output_tokens: OutputTokenUsage { + total: Some(text + reasoning), + text: Some(text), + reasoning: Some(reasoning), + }, + raw: raw.as_object().cloned(), + } +} + +fn vertex_call_options(options: &CallOptions) -> Result { + let mut options = options.clone(); + // The shared validator reads the Google key; validate the effective Vertex options. + if let Some(effective) = Namespace::Vertex + .read(options.provider_options.as_ref()) + .cloned() + { + options + .provider_options + .as_mut() + .unwrap() + .insert(GOOGLE.into(), effective); + } + validate_call_options_for_namespace(&options, Namespace::Vertex)?; + if let Some(provider_options) = Namespace::Vertex.read(options.provider_options.as_ref()) { + let mut headers = options.headers.clone().unwrap_or_default(); + for (option, header) in [ + ("sharedRequestType", "X-Vertex-AI-LLM-Shared-Request-Type"), + ("requestType", "X-Vertex-AI-LLM-Request-Type"), + ] { + if let Some(value) = provider_options.get(option).and_then(Value::as_str) { + headers.retain(|name, _| !name.eq_ignore_ascii_case(header)); + headers.insert(header.to_string(), value.to_string()); + } + } + options.headers = Some(headers); + } + Ok(options) +} + +fn vertex_request_body( + model_id: &str, + options: &CallOptions, + streaming: bool, +) -> Result<(Value, Vec), AiMuxError> { + let mut body = build_vertex_request_body(model_id, options)?; + let mut warnings = + prepare_all_tools(&options.tools, options.tool_choice.as_ref(), model_id).warnings; + body.as_object_mut().unwrap().remove("serviceTier"); + if Namespace::Vertex + .read(options.provider_options.as_ref()) + .and_then(|options| options.get("serviceTier")) + .is_some_and(|value| !value.is_null()) + { + warnings.push(Warning::Other { message: "'serviceTier' is a Gemini API option and is not supported on Vertex AI. Use 'sharedRequestType' (and optionally 'requestType') instead. See https://docs.cloud.google.com/vertex-ai/generative-ai/docs/priority-paygo".to_string() }); + } + if let Some(contents) = body.get_mut("contents").and_then(Value::as_array_mut) { + for content in contents { + if let Some(parts) = content.get_mut("parts").and_then(Value::as_array_mut) { + for part in parts { + for key in ["functionCall", "functionResponse"] { + if let Some(call) = part.get_mut(key).and_then(Value::as_object_mut) { + call.remove("id"); + } + } + } + } + } + } + if streaming + && Namespace::Vertex + .read(options.provider_options.as_ref()) + .and_then(|options| options.get("streamFunctionCallArguments")) + == Some(&Value::Bool(true)) + { + let config = body + .as_object_mut() + .unwrap() + .entry("toolConfig") + .or_insert_with(|| json!({})); + let function_config = config + .as_object_mut() + .unwrap() + .entry("functionCallingConfig") + .or_insert_with(|| json!({})); + function_config["streamFunctionCallArguments"] = json!(true); + } + Ok((body, warnings)) +} diff --git a/aimux-providers/src/vertex/transcription.rs b/aimux-providers/src/vertex/transcription.rs index b63ab1e45..ab3e3d1ef 100644 --- a/aimux-providers/src/vertex/transcription.rs +++ b/aimux-providers/src/vertex/transcription.rs @@ -74,6 +74,7 @@ fn parse_duration_seconds(value: &Option) -> Option { value .as_ref() .and_then(|s| s.trim_end_matches('s').parse::().ok()) + .filter(|value| value.is_finite()) } /// Convert a BCP 47 language tag (e.g. `"en-US"`) to an ISO 639-1 code @@ -82,8 +83,15 @@ fn convert_bcp47_to_iso6391(value: &Option) -> Option { value .as_ref() .and_then(|s| s.split('-').next()) - .filter(|s| s.len() == 2) - .map(std::string::ToString::to_string) + .map(str::to_ascii_lowercase) + .map(|s| match s.as_str() { + "cmn" => "zh".to_string(), + "iw" => "he".to_string(), + "in" => "id".to_string(), + "ji" => "yi".to_string(), + _ => s, + }) + .filter(|s| s.len() == 2 && s.bytes().all(|b| b.is_ascii_alphabetic())) } fn audio_input_to_base64(audio: &AudioInput) -> Result { @@ -205,8 +213,13 @@ impl TranscriptionModel for VertexTranscriptionModel { let exchange = self.config.exchange(options.headers.as_ref()).await?; let url = self.endpoint(exchange.base_url(), &target.project, ®ion); + // The recognizer lives on the Speech-to-Text host, not on the Vertex + // host the endpoint resolves to; it is the host these credentials + // are meant for. + let mut request = exchange.request(url.clone(), options); + request.credentialed_origin = Some(url); let resp = aimux_provider_utils::post_json_to_api( - exchange.request(url, options), + request, exchange.transform_body(request_body.clone()), aimux_provider_utils::create_json_response_handler::(), crate::google::google_failed_response_handler(), @@ -223,11 +236,12 @@ impl TranscriptionModel for VertexTranscriptionModel { // Concatenate transcript from all results. let text: String = results .iter() - .filter_map(|r| { + .map(|r| { r.alternatives .as_ref() .and_then(|a| a.first()) .and_then(|alt| alt.transcript.clone()) + .unwrap_or_default() }) .collect::>() .join(" ") @@ -240,8 +254,8 @@ impl TranscriptionModel for VertexTranscriptionModel { .flat_map(|r| { r.alternatives .as_ref() + .and_then(|alternatives| alternatives.first()) .into_iter() - .flatten() .flat_map(|alt| alt.words.as_deref().unwrap_or(&[])) .filter_map(|w| { let word = w.word.as_ref()?; diff --git a/aimux-providers/src/vertex/video.rs b/aimux-providers/src/vertex/video.rs index 63a182032..fdea68047 100644 --- a/aimux-providers/src/vertex/video.rs +++ b/aimux-providers/src/vertex/video.rs @@ -5,7 +5,7 @@ //! //! Uses the Long Running Operations API: //! 1. POST `{base_url}/models/{model}:predictLongRunning` → returns operation name -//! 2. POST `:fetchPredictOperation` — polled by Core until `done: true` +//! 2. POST fetchPredictOperation — polled by Core until `done: true` //! 3. Return video URL(s) use async_trait::async_trait; @@ -17,7 +17,10 @@ use aimux_core::video_model::{ VideoResponse, VideoResult, }; +use crate::google::options::Namespace; use crate::shared::EndpointConfig; +use aimux_core::shared::{Warning, provider_namespace}; +use aimux_core::video_model::{VideoFile, VideoFileData, VideoFrameType}; /// A Google Vertex AI video generation model. /// @@ -34,6 +37,32 @@ impl VertexVideoModel { } } +fn convert_image(file: &VideoFile, warnings: &mut Vec) -> Option { + match file { + VideoFile::Url { url, .. } if url.starts_with("gs://") => { + Some(json!({"gcsUri": url, "mimeType": "image/png"})) + } + VideoFile::Url { .. } => { + warnings.push(Warning::Unsupported { + feature: "URL-based image input".into(), + details: Some("Vertex AI video models require base64-encoded images or GCS URIs. URL will be ignored.".into()), + }); + None + } + VideoFile::File { media_type, data } => { + let data = match data { + VideoFileData::Base64(data) => data.clone(), + VideoFileData::Binary(data) => { + base64::Engine::encode(&base64::engine::general_purpose::STANDARD, data) + } + }; + Some( + json!({"bytesBase64Encoded": data, "mimeType": if media_type.is_empty() { "image/png" } else { media_type }}), + ) + } + } +} + #[async_trait] impl VideoModel for VertexVideoModel { fn provider(&self) -> &str { @@ -50,28 +79,101 @@ impl VideoModel for VertexVideoModel { &self, options: &VideoCallOptions, ) -> Result { - let mut instances = vec![json!({"prompt": options.prompt})]; - if let Some(ref image) = options.image - && let aimux_core::video_model::VideoFile::Url { url, .. } = image + let mut warnings = Vec::new(); + let provider_options = Namespace::Vertex + .write_keys() + .iter() + .find_map(|key| options.provider_options.get(*key)); + let mut instance = Map::new(); + if let Some(prompt) = &options.prompt { + instance.insert("prompt".into(), json!(prompt)); + } + let frames = options.frame_images.as_deref().unwrap_or(&[]); + let start_image = frames + .iter() + .find(|f| f.frame_type == VideoFrameType::FirstFrame) + .map(|f| &f.image) + .or(options.image.as_ref()); + if let Some(image) = start_image.and_then(|file| convert_image(file, &mut warnings)) { + instance.insert("image".into(), image); + } + if let Some(image) = frames + .iter() + .find(|f| f.frame_type == VideoFrameType::LastFrame) + .and_then(|f| convert_image(&f.image, &mut warnings)) { - instances[0]["image"] = json!({"gcsUri": url, "mimeType": "image/png"}); + instance.insert("lastFrame".into(), image); } - + if frames.is_empty() + && let Some(references) = &options.input_references + && !references.is_empty() + { + instance.insert( + "referenceImages".into(), + Value::Array( + references + .iter() + .filter_map(|file| convert_image(file, &mut warnings)) + .map(|image| json!({"image": image, "referenceType": "asset"})) + .collect(), + ), + ); + } else if let Some(references) = provider_options + .and_then(|p| p.get("referenceImages")) + .filter(|value| !value.is_null()) + { + instance.insert("referenceImages".into(), references.clone()); + } + let instances = vec![Value::Object(instance)]; let mut parameters = Map::new(); - parameters.insert("sampleCount".to_string(), json!(options.n)); + parameters.insert("sampleCount".into(), json!(options.n)); + if let Some(resolution) = options.resolution { + let resolution = resolution.to_string(); + let mapped = match resolution.as_str() { + "1280x720" => "720p", + "1920x1080" => "1080p", + "3840x2160" => "4k", + _ => &resolution, + }; + parameters.insert("resolution".into(), json!(mapped)); + } if let Some(ar) = options.aspect_ratio { parameters.insert("aspectRatio".to_string(), json!(ar.to_string())); } - if let Some(seed) = options.seed { + if let Some(seed) = options.seed.filter(|seed| *seed != 0) { parameters.insert("seed".to_string(), json!(seed)); } - if let Some(duration) = options.duration { + if let Some(duration) = options.duration.filter(|duration| *duration != 0) { parameters.insert("durationSeconds".to_string(), json!(duration)); } if let Some(ga) = options.generate_audio { parameters.insert("generateAudio".to_string(), json!(ga)); } + if let Some(provider_options) = provider_options { + for (key, value) in provider_options { + if ![ + "pollIntervalMs", + "pollTimeoutMs", + "referenceImages", + "generateAudio", + ] + .contains(&key.as_str()) + && (!value.is_null() + || !["personGeneration", "negativePrompt", "gcsOutputDirectory"] + .contains(&key.as_str())) + { + parameters.insert(key.clone(), value.clone()); + } + } + if options.generate_audio.is_none() + && let Some(value) = provider_options.get("generateAudio") + && !value.is_null() + { + parameters.insert("generateAudio".into(), value.clone()); + } + } + let body = json!({ "instances": instances, "parameters": Value::Object(parameters), @@ -93,14 +195,15 @@ impl VideoModel for VertexVideoModel { let operation_name = predict_response .get("name") .and_then(|v| v.as_str()) + .filter(|name| !name.is_empty()) .ok_or_else(|| { AiMuxError::InvalidResponseData("No operation name returned from API".to_string()) })? .to_string(); Ok(VideoOperationStart { - operation: json!({ "operation_name": operation_name }), - warnings: Vec::new(), + operation: json!({ "operationName": operation_name }), + warnings, provider_metadata: None, response: VideoResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), @@ -116,20 +219,19 @@ impl VideoModel for VertexVideoModel { options: &VideoCallOptions, ) -> Result { let operation_name = operation - .get("operation_name") + .get("operationName") .and_then(Value::as_str) .ok_or_else(|| { AiMuxError::InvalidArgument( - "google.vertex operation reference is missing operation_name".to_string(), + "google.vertex operation reference is missing operationName".to_string(), ) })?; let exchange = self.config.exchange(options.headers.as_ref()).await?; let poll_url = exchange.url(&format!("/models/{}:fetchPredictOperation", self.model_id)); - let body = json!({"operationName": operation_name}); let resp = aimux_provider_utils::post_json_to_api( exchange.request(poll_url.clone(), options), - exchange.transform_body(body), + exchange.transform_body(json!({"operationName": operation_name})), aimux_provider_utils::create_json_response_handler::(), crate::google::google_failed_response_handler(), ) @@ -141,7 +243,7 @@ impl VideoModel for VertexVideoModel { if raw_body.get("done").and_then(serde_json::Value::as_bool) != Some(true) { return Ok(VideoOperationStatus::Pending); } - if let Some(err) = raw_body.get("error") { + if let Some(err) = raw_body.get("error").filter(|value| !value.is_null()) { let msg = err .get("message") .and_then(|v| v.as_str()) @@ -156,61 +258,68 @@ impl VideoModel for VertexVideoModel { ..ApiCallError::new(msg, poll_url, serde_json::json!({})) }))); } - - let mut videos = Vec::new(); - let mut video_metadata = Vec::new(); let outputs = raw_body .get("response") .and_then(|r| r.get("videos")) .and_then(Value::as_array) + .filter(|outputs| !outputs.is_empty()) .ok_or_else(|| { AiMuxError::InvalidResponseData(format!( "No videos in response. Response: {raw_body}" )) })?; - if outputs.is_empty() { - return Err(AiMuxError::InvalidResponseData(format!( - "No videos in response. Response: {raw_body}" - ))); - } - for video in outputs { - let mime_type = video.get("mimeType").cloned().unwrap_or(Value::Null); - let media_type = mime_type - .as_str() - .filter(|s| !s.is_empty()) - .unwrap_or("video/mp4") - .to_string(); - if let Some(data) = video - .get("bytesBase64Encoded") - .and_then(Value::as_str) - .filter(|s| !s.is_empty()) - { - videos.push(VideoData::Base64 { - data: data.to_string(), - media_type, - }); - let mut metadata = Map::new(); - if video.get("mimeType").is_some() { - metadata.insert("mimeType".to_string(), mime_type); - } - video_metadata.push(Value::Object(metadata)); - } else if let Some(url) = video - .get("gcsUri") - .and_then(Value::as_str) - .filter(|s| !s.is_empty()) - { - videos.push(VideoData::Url { - url: url.to_string(), - media_type, - }); - let mut metadata = Map::new(); - metadata.insert("gcsUri".to_string(), json!(url)); - if video.get("mimeType").is_some() { - metadata.insert("mimeType".to_string(), mime_type); + let mut video_metadata = Vec::new(); + let videos: Vec = outputs + .iter() + .filter_map(|video| { + let media_type = video + .get("mimeType") + .and_then(Value::as_str) + .filter(|mime| !mime.is_empty()) + .unwrap_or("video/mp4") + .to_string(); + if let Some(data) = video + .get("bytesBase64Encoded") + .and_then(Value::as_str) + .filter(|data| !data.is_empty()) + { + let mut metadata = Map::new(); + if let Some(mime) = video.get("mimeType") { + metadata.insert("mimeType".into(), mime.clone()); + } + video_metadata.push(Value::Object(metadata)); + Some(VideoData::Base64 { + data: data.to_string(), + media_type, + }) + } else if let Some(url) = video + .get("gcsUri") + .and_then(Value::as_str) + .filter(|url| !url.is_empty()) + { + let mut metadata = Map::new(); + metadata.insert("gcsUri".into(), json!(url)); + if let Some(mime) = video.get("mimeType") { + metadata.insert("mimeType".into(), mime.clone()); + } + video_metadata.push(Value::Object(metadata)); + Some(VideoData::Url { + url: url.to_string(), + media_type, + }) + } else { + None } - video_metadata.push(Value::Object(metadata)); - } - } + }) + .collect(); + let payload = json!({"videos": video_metadata}); + let metadata = crate::google::options::VERTEX_VIDEO_METADATA_KEYS + .into_iter() + .flat_map(|key| { + provider_namespace(key, payload.clone()) + .expect("provider metadata must be an object") + }) + .collect(); if videos.is_empty() { return Err(AiMuxError::InvalidResponseData( @@ -221,19 +330,7 @@ impl VideoModel for VertexVideoModel { Ok(VideoOperationStatus::Completed(VideoResult { videos, warnings: Vec::new(), - provider_metadata: Some( - crate::google::options::VERTEX_VIDEO_METADATA_KEYS - .into_iter() - .map(|key| { - ( - key.to_string(), - [("videos".to_string(), json!(video_metadata))] - .into_iter() - .collect(), - ) - }) - .collect(), - ), + provider_metadata: Some(metadata), response: VideoResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), model_id: Some(self.model_id.clone()), diff --git a/aimux-providers/src/xai/convert.rs b/aimux-providers/src/xai/convert.rs index c28f643d6..d9699238d 100644 --- a/aimux-providers/src/xai/convert.rs +++ b/aimux-providers/src/xai/convert.rs @@ -1,11 +1,18 @@ //! Helpers shared by the xAI Responses converter. //! -//! Mirrors the TS xai package's `supports-reasoning-effort.ts` and -//! `remove-additional-properties.ts`, plus the media-type and provider -//! reference helpers its input conversion uses. +//! Mirrors `supports-reasoning-effort.ts` and the provider-reference helper +//! used by upstream input conversion. -use aimux_core::shared::SharedProviderReference; -use serde_json::Value; +use std::collections::HashMap; + +use aimux_core::{ + error::AiMuxError, + language_model_message::FilePart, + shared::{FileBytes, FileData}, +}; +use aimux_provider_utils::{ + MediaTypeData, detect_media_type, get_top_level_media_type, is_full_media_type, +}; // ── Reasoning effort ───────────────────────────────────────────────────────── @@ -33,27 +40,6 @@ pub fn supports_reasoning_effort(model_id: &str) -> bool { !is_model_without_reasoning_effort(model_id) } -// ── JSON schema ────────────────────────────────────────────────────────────── - -pub fn remove_additional_properties_false(value: &Value) -> Value { - match value { - Value::Array(arr) => { - Value::Array(arr.iter().map(remove_additional_properties_false).collect()) - } - Value::Object(obj) => { - let mut result = serde_json::Map::new(); - for (key, val) in obj { - if key == "additionalProperties" && val == &Value::Bool(false) { - continue; - } - result.insert(key.clone(), remove_additional_properties_false(val)); - } - Value::Object(result) - } - other => other.clone(), - } -} - /// Resolve the provider-specific reference string from a file part's /// `provider` map. /// @@ -62,7 +48,7 @@ pub fn remove_additional_properties_false(value: &Value) -> Value { /// Returns a `String` listing the available providers when the requested key /// is absent. pub fn resolve_provider_reference( - reference: &SharedProviderReference, + reference: &HashMap, provider: &str, ) -> Result { if let Some(value) = reference.get(provider) { @@ -76,3 +62,31 @@ pub fn resolve_provider_reference( available.join(", ") )) } + +pub(super) fn resolve_full_media_type(part: &FilePart) -> Result { + let media_type = &part.media_type; + if is_full_media_type(media_type) { + return Ok(media_type.clone()); + } + let data = match &part.data { + FileData::Data { + data: FileBytes::Binary(bytes), + } => Some(MediaTypeData::Bytes(bytes)), + FileData::Data { + data: FileBytes::Base64(data), + } => Some(MediaTypeData::Base64(data)), + _ => None, + }; + let reason = if let Some(data) = data { + if let Some(detected) = detect_media_type(data, Some(get_top_level_media_type(media_type)))? + { + return Ok(detected.into()); + } + "it could not be auto-detected" + } else { + "it is not passed as inline bytes" + }; + Err(AiMuxError::UnsupportedFunctionality(format!( + "file of media type \"{media_type}\" must specify subtype since {reason}" + ))) +} diff --git a/aimux-providers/src/xai/mod.rs b/aimux-providers/src/xai/mod.rs index d3859e8b0..aff464878 100644 --- a/aimux-providers/src/xai/mod.rs +++ b/aimux-providers/src/xai/mod.rs @@ -1,7 +1,7 @@ //! xAI (Grok) provider. //! //! [`create_xai`] is the Rust form of the AI SDK's `createXai`: it takes -//! [`XAIProviderSettings`], validates the base URL, fixes the provider name and +//! [`XAIProviderSettings`], removes a trailing slash, fixes the provider name and //! returns an [`XAIProvider`]. The API key is not read there; it is loaded in //! the request headers of every call, from the setting or from //! `XAI_API_KEY`. [`xai()`] is the default instance. @@ -30,13 +30,12 @@ use aimux_core::image_model::ImageModel; use aimux_core::language_model::LanguageModel; use aimux_core::model_catalogue::RuntimeModel; use aimux_core::provider::{Provider, ProviderDiscovery}; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable}; use crate::shared::{Credential, EndpointConfig, provider_headers}; const DEFAULT_BASE_URL: &str = "https://api.x.ai/v1"; const API_KEY_ENV_VAR: &str = "XAI_API_KEY"; -const DEFAULT_NAME: &str = "xai"; pub(crate) fn xai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -52,8 +51,8 @@ pub(crate) fn xai_failed_response_handler() -> aimux_provider_utils::ResponseHan let error = data.get("error").unwrap_or(data); aimux_provider_utils::ProviderErrorParts { message: error - .get("message") - .and_then(Value::as_str) + .as_str() + .or_else(|| error.get("message").and_then(Value::as_str)) .unwrap_or_default() .to_owned(), provider_code: error @@ -74,7 +73,14 @@ pub(crate) fn xai_stream_error( request_body_values: Value, response_headers: std::collections::HashMap, ) -> AiMuxError { - let error = event.get("error").unwrap_or(event); + let error = event + .get("error") + .or_else(|| { + event + .get("response") + .and_then(|response| response.get("error")) + }) + .unwrap_or(event); let message = error .as_str() .or_else(|| error.get("message").and_then(Value::as_str)) @@ -85,8 +91,15 @@ pub(crate) fn xai_stream_error( .or_else(|| event.get("code")) .or_else(|| error.get("status")) .or_else(|| error.get("code")) - .and_then(Value::as_u64) - .and_then(|status| u16::try_from(status).ok()) + .and_then(|value| match value { + Value::String(value) + if value.len() == 3 && value.bytes().all(|byte| byte.is_ascii_digit()) => + { + value.parse::().ok() + } + Value::Number(value) => value.as_u64().and_then(|status| u16::try_from(status).ok()), + _ => None, + }) .filter(|status| (400..=599).contains(status)); let provider_code = event .get("code") @@ -97,7 +110,19 @@ pub(crate) fn xai_stream_error( Value::Number(value) => Some(value.to_string()), _ => None, }); - aimux_provider_utils::stream_error_api_call( + let status_code = status_code.or(match provider_code.as_deref() { + Some("rate_limit_exceeded" | "rate_limit_error" | "insufficient_quota") => Some(429), + Some("api_error" | "internal_server_error" | "server_error") => Some(500), + Some("overloaded_error" | "service_unavailable") => Some(503), + Some("timeout" | "timeout_error") => Some(504), + Some("authentication_error" | "invalid_api_key") => Some(401), + Some("permission_error") => Some(403), + Some("not_found_error" | "model_not_found") => Some(404), + Some("bad_request" | "context_length_exceeded" | "invalid_request_error") => Some(400), + _ => None, + }); + let insufficient_quota = provider_code.as_deref() == Some("insufficient_quota"); + let mut result = aimux_provider_utils::stream_error_api_call( message, provider_code, status_code, @@ -105,7 +130,11 @@ pub(crate) fn xai_stream_error( url, request_body_values, response_headers, - ) + ); + if insufficient_quota && let AiMuxError::ApiCall(error) = &mut result { + error.is_retryable = false; + } + result } pub(crate) fn xai_successful_response_handler() -> aimux_provider_utils::ResponseHandler @@ -265,16 +294,14 @@ impl std::fmt::Debug for XAIProviderSettings { /// /// # Errors /// -/// Returns `AiMuxError::InvalidArgument` when `base_url` is not an `http(s)` -/// URL with a host. That is the only way this fails: the key is loaded per -/// request, not here. +/// Creation does not validate the URL or load the API key; request errors +/// surface when the model is called. pub fn create_xai(settings: XAIProviderSettings) -> Result { - let base_url = match settings.base_url.as_deref() { - Some(url) => validate_base_url(url)?, - None => DEFAULT_BASE_URL.to_string(), - }; + let base_url = settings + .base_url + .unwrap_or_else(|| DEFAULT_BASE_URL.to_string()); + let base_url = base_url.strip_suffix('/').unwrap_or(&base_url).to_string(); Ok(XAIProvider { - name: DEFAULT_NAME.to_string(), base_url, headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( provider_headers( @@ -306,7 +333,6 @@ pub fn xai() -> &'static XAIProvider { /// An xAI provider (the AI SDK's `XaiProvider`). Cheap to clone the models out /// of; it holds no HTTP client. pub struct XAIProvider { - name: String, base_url: String, headers: HeadersFn, fetch: Option, @@ -315,7 +341,7 @@ pub struct XAIProvider { impl XAIProvider { fn model_config(&self, method: &str) -> EndpointConfig { EndpointConfig::fixed( - format!("{}.{method}", self.name), + format!("xai.{method}"), self.base_url.clone(), self.headers.clone(), self.fetch.clone(), @@ -324,7 +350,7 @@ impl XAIProvider { } /// A Responses model (e.g. `"grok-4"`); `provider()` is - /// `"{name}.responses"`. + /// `"xai.responses"`. /// /// Uses the xAI `/responses` endpoint with the Responses API wire format /// (input items, reasoning objects, provider-executed tools, etc.). diff --git a/aimux-providers/src/xai/options.rs b/aimux-providers/src/xai/options.rs index 1078e3228..34824b72e 100644 --- a/aimux-providers/src/xai/options.rs +++ b/aimux-providers/src/xai/options.rs @@ -20,5 +20,5 @@ pub(crate) fn xai_options(provider_options: Option<&SharedProviderOptions>) -> O /// `{ "xai": payload }`. pub(crate) fn xai_metadata(payload: Value) -> ProviderMetadata { - provider_namespace(NAMESPACE, payload).expect("provider metadata must be an object") + provider_namespace(NAMESPACE, payload).expect("metadata payload is an object") } diff --git a/aimux-providers/src/xai/responses/convert.rs b/aimux-providers/src/xai/responses/convert.rs index e47826042..6bfd80c73 100644 --- a/aimux-providers/src/xai/responses/convert.rs +++ b/aimux-providers/src/xai/responses/convert.rs @@ -16,12 +16,12 @@ use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; use aimux_core::tool::Tool; use aimux_core::types::{FinishReasonUnified, ReasoningEffort, Warning}; -use aimux_provider_utils::{get_top_level_media_type, resolve_full_media_type}; +use aimux_provider_utils::get_top_level_media_type; use serde_json::{Value, json}; use super::types::XaiResponsesUsage; use crate::xai::convert::{ - remove_additional_properties_false, resolve_provider_reference, supports_reasoning_effort, + resolve_full_media_type, resolve_provider_reference, supports_reasoning_effort, }; // ── Finish reason ──────────────────────────────────────────────────────────── @@ -33,7 +33,6 @@ pub fn map_xai_responses_finish_reason(s: &str) -> FinishReasonUnified { "length" | "max_output_tokens" => FinishReasonUnified::Length, "tool_calls" | "function_call" => FinishReasonUnified::ToolCalls, "content_filter" => FinishReasonUnified::ContentFilter, - "error" => FinishReasonUnified::Error, _ => FinishReasonUnified::Other, } } @@ -95,12 +94,18 @@ pub struct PreparedResponsesTools { pub provider_tool_names: std::collections::HashMap, } -#[must_use] +/// # Errors +/// Rejects provider-tool arguments that do not match the upstream schemas. pub fn prepare_responses_tools( tools: &Option>, tool_choice: Option<&ToolChoice>, -) -> PreparedResponsesTools { +) -> Result { let non_empty = tools.as_ref().filter(|t| !t.is_empty()); + for tool in non_empty.into_iter().flatten() { + if let Tool::Provider(tool) = tool { + validate_provider_tool(tool)?; + } + } let mut tool_warnings: Vec = Vec::new(); let mut provider_tool_names: std::collections::HashMap = std::collections::HashMap::new(); @@ -115,7 +120,7 @@ pub fn prepare_responses_tools( let mut func = json!({ "type": "function", "name": ft.name, - "parameters": remove_additional_properties_false(&ft.input_schema), + "parameters": ft.input_schema, }); if let Some(ref desc) = ft.description { func["description"] = json!(desc); @@ -131,11 +136,7 @@ pub fn prepare_responses_tools( } }) .collect(); - if xai_tools.is_empty() { - None - } else { - Some(xai_tools) - } + Some(xai_tools) } }; @@ -150,7 +151,7 @@ pub fn prepare_responses_tools( ToolChoice::None => Some(json!("none")), ToolChoice::Required => Some(json!("required")), ToolChoice::Tool { tool_name } => { - let selected = non_empty_tools.iter().find(|t| match t { + let selected = non_empty_tools.iter().rev().find(|t| match t { Tool::Function(ft) => ft.name == *tool_name, Tool::Provider(pt) => pt.name == *tool_name, }); @@ -174,12 +175,12 @@ pub fn prepare_responses_tools( } }; - PreparedResponsesTools { + Ok(PreparedResponsesTools { tools: tools_opt, tool_choice: tool_choice_opt, tool_warnings, provider_tool_names, - } + }) } fn prepare_provider_tool( @@ -229,6 +230,13 @@ fn prepare_provider_tool( "xai.code_execution" => Some(json!({ "type": "code_interpreter" })), "xai.view_image" => Some(json!({ "type": "view_image" })), "xai.view_x_video" => Some(json!({ "type": "view_x_video" })), + "xai.image_generation" => { + let mut tool = json!({ "type": "image_generation" }); + if let Some(action) = args.get("action") { + tool["action"] = action.clone(); + } + Some(tool) + } "xai.file_search" => { let mut tool = json!({ "type": "file_search" }); if let Some(v) = args.get("vectorStoreIds") { @@ -281,6 +289,131 @@ fn xai_option(options: &Option, key: &str) -> Option bool, +) -> Result<(), AiMuxError> { + if valid(value) { + Ok(()) + } else { + Err(AiMuxError::InvalidArgument(format!( + "Invalid xai option {key}: {value}" + ))) + } +} + +fn string_array(value: &Value, maximum: usize) -> bool { + value + .as_array() + .is_some_and(|items| items.len() <= maximum && items.iter().all(Value::is_string)) +} + +fn validate_response_options(options: &Option) -> Result<(), AiMuxError> { + for (key, value) in options + .as_ref() + .and_then(|options| crate::xai::options::xai_options(Some(options))) + .into_iter() + .flatten() + { + validate_value(key, value, |value| match key.as_str() { + "reasoningEffort" => value + .as_str() + .is_some_and(|v| ["none", "low", "medium", "high", "xhigh"].contains(&v)), + "reasoningSummary" => value + .as_str() + .is_some_and(|v| ["auto", "concise", "detailed"].contains(&v)), + "serviceTier" => value + .as_str() + .is_some_and(|v| ["default", "priority"].contains(&v)), + "topLogprobs" => value + .as_f64() + .is_some_and(|v| v.fract() == 0.0 && (0.0..=8.0).contains(&v)), + "minP" => value.as_f64().is_some_and(|v| (0.0..=1.0).contains(&v)), + "maxTurns" => value + .as_f64() + .is_some_and(|v| v.fract() == 0.0 && v.abs() <= 9_007_199_254_740_991.0), + "logprobs" | "store" | "parallelToolCalls" => value.is_boolean(), + "promptCacheKey" | "safetyIdentifier" | "user" | "previousResponseId" => { + value.is_string() + } + "include" => { + value.is_null() + || value.as_array().is_some_and(|items| { + items.iter().all(|item| { + item.as_str().is_some_and(|v| { + [ + "file_search_call.results", + "web_search_call.action.sources", + "code_interpreter_call.outputs", + "reasoning.encrypted_content", + "no_inline_citations", + ] + .contains(&v) + }) + }) + }) + } + _ => true, + })?; + } + Ok(()) +} + +fn validate_provider_tool(tool: &aimux_core::tool::ProviderTool) -> Result<(), AiMuxError> { + let id = tool.id.as_str(); + if ![ + "xai.web_search", + "xai.x_search", + "xai.image_generation", + "xai.file_search", + "xai.mcp", + ] + .contains(&id) + { + return Ok(()); + } + let args = &tool.args; + for (key, value) in args { + validate_value(key, value, |value| match (id, key.as_str()) { + ("xai.web_search", "allowedDomains" | "excludedDomains") => string_array(value, 5), + ("xai.x_search", "allowedXHandles" | "excludedXHandles") => string_array(value, 10), + ("xai.web_search", "enableImageSearch" | "enableImageUnderstanding") + | ("xai.x_search", "enableImageUnderstanding" | "enableVideoUnderstanding") => { + value.is_boolean() + } + ("xai.x_search", "fromDate" | "toDate") => value.is_string(), + ("xai.image_generation", "action") => value + .as_str() + .is_some_and(|v| ["auto", "generate", "edit"].contains(&v)), + ("xai.file_search", "vectorStoreIds") | ("xai.mcp", "allowedTools") => { + string_array(value, usize::MAX) + } + ("xai.file_search", "maxNumResults") => value.is_number(), + ("xai.mcp", "serverUrl" | "serverLabel" | "serverDescription" | "authorization") => { + value.is_string() + } + ("xai.mcp", "headers") => value + .as_object() + .is_some_and(|headers| headers.values().all(Value::is_string)), + _ => true, + })?; + } + for key in match id { + "xai.file_search" => &["vectorStoreIds"][..], + "xai.mcp" => &["serverUrl"][..], + _ => &[], + } { + if !args.contains_key(*key) { + return Err(AiMuxError::InvalidArgument(format!( + "Missing {key} for {}", + tool.name + ))); + } + } + Ok(()) +} + // ── Input conversion ───────────────────────────────────────────────────────── /// Convert a [`LanguageModelPrompt`] into the xAI Responses `input` array, @@ -383,16 +516,7 @@ pub fn convert_to_xai_responses_input( .. }) => { // Skip provider-executed tool calls. - // The standardized field is authoritative. Fall back - // to the legacy provider option only when it is absent. - let is_provider_executed = provider_executed.unwrap_or_else(|| { - provider_options - .as_ref() - .and_then(|po| crate::xai::options::xai_options(Some(po))) - .and_then(|x| x.get("providerExecuted")) - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) - }); + let is_provider_executed = provider_executed.unwrap_or(false); if is_provider_executed { continue; } @@ -454,13 +578,12 @@ pub fn convert_to_xai_responses_input( } } AssistantPart::File(_) - | AssistantPart::Custom(_) - | AssistantPart::ReasoningFile(_) => { + | AssistantPart::ReasoningFile(_) + | AssistantPart::Custom(_) => { let kind = match part { AssistantPart::File(_) => "file", - AssistantPart::Custom(_) => "custom", AssistantPart::ReasoningFile(_) => "reasoning-file", - _ => unreachable!(), + _ => "custom", }; warnings.push(Warning::Other { message: format!("xAI Responses API does not support {kind} in assistant messages"), @@ -479,7 +602,7 @@ pub fn convert_to_xai_responses_input( else { continue; }; - let output_value = convert_tool_result_output(output)?; + let output_value = convert_tool_output(output)?; input.push(json!({ "type": "function_call_output", "call_id": tool_call_id, @@ -493,10 +616,10 @@ pub fn convert_to_xai_responses_input( Ok((input, warnings)) } -fn convert_tool_result_output(output: &ToolResultOutput) -> Result { - Ok(match output { +fn convert_tool_output(output: &ToolResultOutput) -> Result { + match output { ToolResultOutput::ExecutionDenied { reason, .. } => { - json!(reason.as_deref().unwrap_or("tool execution denied")) + Ok(json!(reason.as_deref().unwrap_or("tool execution denied"))) } ToolResultOutput::Content { value } => { let mut parts = Vec::new(); @@ -506,31 +629,31 @@ fn convert_tool_result_output(output: &ToolResultOutput) -> Result + if get_top_level_media_type(&part.media_type) == "image" => { - let url = match &part.data { + let image_url = match &part.data { FileData::Url { url, .. } => url.clone(), FileData::Data { data } => { use base64::Engine; - let b64 = match data { + let encoded = match data { + FileBytes::Base64(data) => data.clone(), FileBytes::Binary(bytes) => { base64::engine::general_purpose::STANDARD.encode(bytes) } - FileBytes::Base64(data) => data.clone(), }; - format!("data:{};base64,{}", resolve_full_media_type(part)?, b64) + format!("data:{};base64,{encoded}", resolve_full_media_type(part)?) } _ => continue, }; - parts.push(json!({"type":"input_image", "image_url":url})); + parts.push(json!({"type":"input_image", "image_url":image_url})); } _ => {} } } - json!(parts) + Ok(json!(parts)) } - _ => crate::openai::convert::tool_result_to_content(output), - }) + _ => Ok(crate::openai::convert::tool_result_to_content(output)), + } } fn convert_image_part( @@ -565,6 +688,11 @@ fn convert_image_part( .and_then(|po| crate::xai::options::xai_options(Some(po))) .and_then(|x| x.get("imageDetail")) { + validate_value("imageDetail", detail, |value| { + value + .as_str() + .is_some_and(|value| ["auto", "low", "high"].contains(&value)) + })?; part["detail"] = detail.clone(); } @@ -592,64 +720,8 @@ pub fn build_responses_request_body( stream: bool, ) -> Result { let mut warnings: Vec = Vec::new(); - let parsed_options = options - .provider_options - .as_ref() - .map(|namespaces| { - let mut namespaces = namespaces.clone(); - if let Some(raw) = namespaces.get(crate::xai::options::NAMESPACE) { - let parsed = crate::openai::convert::parse_option_fields( - raw, - crate::xai::options::NAMESPACE, - |key, value| { - let valid = match key { - "reasoningEffort" => value.as_str().is_some_and(|s| { - matches!(s, "none" | "low" | "medium" | "high" | "xhigh") - }), - "reasoningSummary" => value - .as_str() - .is_some_and(|s| matches!(s, "auto" | "concise" | "detailed")), - "logprobs" | "parallelToolCalls" | "store" => value.is_boolean(), - "topLogprobs" => value - .as_f64() - .is_some_and(|n| n.fract() == 0.0 && (0.0..=8.0).contains(&n)), - "minP" => value.as_f64().is_some_and(|n| (0.0..=1.0).contains(&n)), - "maxTurns" => value.as_f64().is_some_and(|n| { - n.fract() == 0.0 && n.abs() <= 9_007_199_254_740_991.0 - }), - "promptCacheKey" | "safetyIdentifier" | "previousResponseId" - | "user" => value.is_string(), - "serviceTier" => value - .as_str() - .is_some_and(|s| matches!(s, "default" | "priority")), - "include" => { - value.is_null() - || value.as_array().is_some_and(|items| { - items.iter().all(|item| { - item.as_str().is_some_and(|s| { - matches!( - s, - "file_search_call.results" - | "web_search_call.action.sources" - | "code_interpreter_call.outputs" - | "reasoning.encrypted_content" - | "no_inline_citations" - ) - }) - }) - }) - } - _ => return None, - }; - Some(if valid { Ok(value.clone()) } else { Err(()) }) - }, - )?; - namespaces.insert(crate::xai::options::NAMESPACE.to_string(), parsed); - } - Ok::<_, AiMuxError>(namespaces) - }) - .transpose()?; - let xai_opts = &parsed_options; + let xai_opts = &options.provider_options; + validate_response_options(xai_opts)?; if options.frequency_penalty.is_some() { warnings.push(Warning::Unsupported { @@ -675,7 +747,7 @@ pub fn build_responses_request_body( let (input, input_warnings) = convert_to_xai_responses_input(&options.prompt)?; warnings.extend(input_warnings); - let prepared = prepare_responses_tools(&options.tools, options.tool_choice.as_ref()); + let prepared = prepare_responses_tools(&options.tools, options.tool_choice.as_ref())?; for tw in &prepared.tool_warnings { warnings.push(tw.clone()); } @@ -704,14 +776,22 @@ pub fn build_responses_request_body( ReasoningEffort::Low => "low".to_string(), ReasoningEffort::Medium => "medium".to_string(), ReasoningEffort::High => "high".to_string(), + ReasoningEffort::Xhigh if model_id == "grok-4.6" => "xhigh".to_string(), ReasoningEffort::Xhigh => "high".to_string(), _ => reasoning.to_string(), }); + if reasoning_effort.as_deref() != Some(reasoning.to_string().as_str()) { + warnings.push(Warning::Compatibility { + feature: "reasoning".to_string(), + details: Some(format!( + "reasoning \"{reasoning}\" is not directly supported by this model. mapped to effort \"{}\".", + reasoning_effort.as_deref().unwrap_or_default() + )), + }); + } } } - let reasoning_summary = xai_option(xai_opts, "reasoningSummary"); - // Store option. let store = xai_option(xai_opts, "store").and_then(|v| v.as_bool()); @@ -771,16 +851,21 @@ pub fn build_responses_request_body( body["top_logprobs"] = tlp; } - // Reasoning object. - if reasoning_effort.is_some() || reasoning_summary.is_some() { - let mut reasoning_obj = json!({}); - if let Some(effort) = &reasoning_effort { - reasoning_obj["effort"] = json!(effort); - } - if let Some(summary) = &reasoning_summary { - reasoning_obj["summary"] = json!(summary); + if let Some(effort) = &reasoning_effort { + body["reasoning"] = json!({ "effort": effort }); + } + for (option, parameter) in [ + ("minP", "min_p"), + ("maxTurns", "max_turns"), + ("parallelToolCalls", "parallel_tool_calls"), + ("promptCacheKey", "prompt_cache_key"), + ("safetyIdentifier", "safety_identifier"), + ("serviceTier", "service_tier"), + ("user", "user"), + ] { + if let Some(value) = xai_option(xai_opts, option) { + body[parameter] = value; } - body["reasoning"] = reasoning_obj; } // Store. @@ -871,18 +956,10 @@ pub fn resolve_tool_name( .get("xai.code_execution") .cloned() .unwrap_or_else(|| "code_execution".to_string()); - let default_view_image = provider_tool_names - .get("xai.view_image") - .cloned() - .unwrap_or_else(|| "view_image".to_string()); - let default_view_x_video = provider_tool_names - .get("xai.view_x_video") - .cloned() - .unwrap_or_else(|| "view_x_video".to_string()); let default_mcp = provider_tool_names .get("xai.mcp") .cloned() - .unwrap_or_else(|| "mcp".to_string()); + .unwrap_or_else(|| part_name.unwrap_or("mcp").to_string()); let default_file_search = provider_tool_names .get("xai.file_search") .cloned() @@ -897,14 +974,15 @@ pub fn resolve_tool_name( || part_type == "code_execution_call" { default_code - } else if name == "view_image" || part_type == "view_image_call" { - default_view_image - } else if name == "view_x_video" || part_type == "view_x_video_call" { - default_view_x_video } else if part_type == "mcp_call" { default_mcp } else if part_type == "file_search_call" { default_file_search + } else if part_type == "image_generation_call" { + provider_tool_names + .get("xai.image_generation") + .cloned() + .unwrap_or_else(|| "image_generation".to_string()) } else if !name.is_empty() { name.to_string() } else { diff --git a/aimux-providers/src/xai/responses/mod.rs b/aimux-providers/src/xai/responses/mod.rs index 290108c21..0d062001f 100644 --- a/aimux-providers/src/xai/responses/mod.rs +++ b/aimux-providers/src/xai/responses/mod.rs @@ -44,6 +44,134 @@ fn generate_source_id() -> String { format!("id-{}", SOURCE_ID_COUNTER.fetch_add(1, Ordering::SeqCst)) } +fn zero_usage() -> Usage { + Usage { + input_tokens: aimux_core::types::InputTokenUsage { + total: Some(0), + no_cache: Some(0), + cache_read: Some(0), + cache_write: Some(0), + }, + output_tokens: aimux_core::types::OutputTokenUsage { + total: Some(0), + text: Some(0), + reasoning: Some(0), + }, + raw: None, + } +} + +fn response_provider_metadata( + cost: Option, + service_tier: Option<&str>, + prompt_cache_key: Option<&str>, + safety_identifier: Option<&str>, +) -> Option { + let mut metadata = serde_json::Map::new(); + if let Some(cost) = cost { + metadata.insert("costInUsdTicks".into(), json!(cost)); + } + for (key, value) in [ + ("serviceTier", service_tier), + ("promptCacheKey", prompt_cache_key), + ("safetyIdentifier", safety_identifier), + ] { + if let Some(value) = value { + metadata.insert(key.into(), json!(value)); + } + } + (!metadata.is_empty()).then(|| crate::xai::options::xai_metadata(Value::Object(metadata))) +} + +fn map_web_search_action(action: &Value) -> Value { + let Some(kind) = action.get("type").and_then(Value::as_str) else { + return json!({}); + }; + let fields: &[&str] = match kind { + "search" => &["query"], + "open_page" => &["url"], + "find_in_page" => &["url", "pattern"], + _ => return json!({}), + }; + if fields.iter().any(|key| { + action + .get(key) + .is_some_and(|v| !v.is_null() && !v.is_string()) + }) || action + .get("sources") + .is_some_and(|v| !v.is_null() && !v.is_array()) + || (kind == "search" + && action.get("queries").is_some_and(|v| { + !v.is_null() && !v.as_array().is_some_and(|a| a.iter().all(Value::is_string)) + })) + { + return json!({}); + } + let mapped_type = match kind { + "open_page" => "openPage", + "find_in_page" => "findInPage", + _ => kind, + }; + let mut mapped_action = json!({ "type": mapped_type }); + for key in fields { + if let Some(value) = action.get(key) + && (kind != "search" || !value.is_null()) + { + mapped_action[key] = value.clone(); + } + } + if kind == "search" + && let Some(queries) = action.get("queries").filter(|v| !v.is_null()) + { + mapped_action["queries"] = queries.clone(); + } + let mut result = json!({ "action": mapped_action }); + if let Some(sources) = action.get("sources").and_then(Value::as_array) { + let sources: Vec = sources + .iter() + .filter(|source| source.get("type").and_then(Value::as_str) == Some("url")) + .filter_map(|source| source.get("url").and_then(Value::as_str)) + .map(|url| json!({ "type": "url", "url": url })) + .collect(); + if !sources.is_empty() { + result["sources"] = json!(sources); + } + } + result +} + +fn image_generation_result(part: &Value, tool_name: String) -> ToolResult { + let (result, is_error) = match part.get("result").filter(|value| !value.is_null()) { + Some(result) => { + let mut output = json!({ "result": result }); + if let Some(prompt) = part.get("prompt").filter(|value| !value.is_null()) { + output["prompt"] = prompt.clone(); + } + (output, None) + } + None => ( + json!(format!( + "Image generation failed (status: {}).", + part.get("status").and_then(Value::as_str).unwrap_or("") + )), + Some(true), + ), + }; + ToolResult { + tool_call_id: part + .get("id") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(), + tool_name, + result, + is_error, + preliminary: None, + dynamic: None, + provider_metadata: None, + } +} + /// An xAI Responses language model (Grok). /// /// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the @@ -85,7 +213,7 @@ impl LanguageModel for XaiResponsesModel { let response_headers = resp.response_headers; - let response_value = resp.raw_value.unwrap_or(Value::Null); + let raw_value = resp.raw_value.unwrap_or(Value::Null); let data = resp.value; let mut content: Vec = Vec::new(); @@ -94,6 +222,26 @@ impl LanguageModel for XaiResponsesModel { for part in &data.output { let part_type = part.get("type").and_then(|v| v.as_str()).unwrap_or(""); + if part_type == "image_generation_call" { + let tool_name = resolve_tool_name(part_type, None, &provider_tool_names); + content.push(GenerateContent::ToolCall(RawToolCall { + tool_call_id: part + .get("id") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(), + tool_name: tool_name.clone(), + input: "{}".to_string(), + provider_executed: Some(true), + dynamic: None, + provider_metadata: None, + })); + content.push(GenerateContent::ToolResult(image_generation_result( + part, tool_name, + ))); + continue; + } + // ── file_search_call ── if part_type == "file_search_call" { let tool_name = resolve_tool_name(part_type, None, &provider_tool_names); @@ -157,12 +305,23 @@ impl LanguageModel for XaiResponsesModel { content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: part_id.to_string(), - tool_name, + tool_name: tool_name.clone(), input: tool_input, provider_executed: Some(true), dynamic: None, provider_metadata: None, })); + if part_type == "web_search_call" { + content.push(GenerateContent::ToolResult(ToolResult { + tool_call_id: part_id.to_string(), + tool_name, + result: map_web_search_action(part.get("action").unwrap_or(&Value::Null)), + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: None, + })); + } continue; } @@ -253,8 +412,14 @@ impl LanguageModel for XaiResponsesModel { }; let reasoning_text = texts.join(""); - let encrypted_content = part.get("encrypted_content").and_then(|v| v.as_str()); - let item_id = part.get("id").and_then(|v| v.as_str()); + let encrypted_content = part + .get("encrypted_content") + .and_then(|v| v.as_str()) + .filter(|v| !v.is_empty()); + let item_id = part + .get("id") + .and_then(|v| v.as_str()) + .filter(|v| !v.is_empty()); if !reasoning_text.is_empty() || encrypted_content.is_some() { let mut meta = json!({}); @@ -266,7 +431,8 @@ impl LanguageModel for XaiResponsesModel { } content.push(GenerateContent::Reasoning(ReasoningOutput { text: reasoning_text, - provider_metadata: Some(crate::xai::options::xai_metadata(json!(meta))), + provider_metadata: (encrypted_content.is_some() || item_id.is_some()) + .then(|| crate::xai::options::xai_metadata(meta)), })); } } @@ -286,18 +452,19 @@ impl LanguageModel for XaiResponsesModel { raw: data.status.clone(), }; - let (usage, provider_metadata) = if let Some(u) = &data.usage { - let meta = if u.cost_in_usd_ticks.is_some() { - Some(crate::xai::options::xai_metadata( - json!({ "costInUsdTicks": u.cost_in_usd_ticks }), - )) - } else { - None - }; - (convert_xai_responses_usage(u), meta) - } else { - (Usage::default(), None) - }; + let usage = data.usage.as_ref().map_or_else(zero_usage, |u| { + let mut usage = convert_xai_responses_usage(u); + usage.raw = raw_value + .get("usage") + .and_then(|value| value.as_object().cloned()); + usage + }); + let provider_metadata = response_provider_metadata( + data.usage.as_ref().and_then(|u| u.cost_in_usd_ticks), + data.service_tier.as_deref(), + data.prompt_cache_key.as_deref(), + data.safety_identifier.as_deref(), + ); let timestamp = data .created_at @@ -314,7 +481,7 @@ impl LanguageModel for XaiResponsesModel { timestamp, model_id: data.model, headers: Some(response_headers), - body: Some(response_value), + body: Some(raw_value), }), request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), }) @@ -341,9 +508,9 @@ impl LanguageModel for XaiResponsesModel { let mut sse_stream = resp.value; let first_event = match sse_stream.next().await { Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), - first_event => first_event, + first => first, }; - if let Some(Ok(ref event)) = first_event + if let Some(Ok(event)) = &first_event && types::event_type(event) == "error" { return Err(super::xai_stream_error( @@ -356,6 +523,7 @@ impl LanguageModel for XaiResponsesModel { let stream_error_url = endpoint; let stream_request_body = body.clone(); let stream_response_headers = response_headers.clone(); + let include_raw_chunks = options.include_raw_chunks == Some(true); let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings }); @@ -366,11 +534,14 @@ impl LanguageModel for XaiResponsesModel { raw: None, }; let mut cost_in_usd_ticks: Option = None; + let mut service_tier: Option = None; + let mut prompt_cache_key: Option = None; + let mut safety_identifier: Option = None; let mut has_function_call = false; let mut is_first_chunk = true; // Track content blocks. - let mut content_blocks: HashMap = HashMap::new(); // block_id -> ended (text blocks) + let mut content_blocks: Vec = Vec::new(); let mut seen_tool_calls: std::collections::HashSet = std::collections::HashSet::new(); let mut active_reasoning: HashMap = HashMap::new(); @@ -378,10 +549,12 @@ impl LanguageModel for XaiResponsesModel { let mut ongoing_tool_calls: HashMap = HashMap::new(); // output_index -> (tool_call_id, tool_name) let mut event_iter = futures::stream::iter(first_event.into_iter()).chain(sse_stream); - while let Some(event) = event_iter.next().await { match event { Ok(parsed) => { + if include_raw_chunks { + yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); + } let event_type = types::event_type(&parsed); // ── response.created / response.in_progress ── @@ -462,8 +635,8 @@ impl LanguageModel for XaiResponsesModel { let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); let block_id = format!("text-{item_id}"); - if !content_blocks.contains_key(&block_id) { - content_blocks.insert(block_id.clone(), false); + if !content_blocks.contains(&block_id) { + content_blocks.push(block_id.clone()); yield Ok(StreamPart::TextStart { id: block_id.clone(), provider_metadata: None}); } yield Ok(StreamPart::TextDelta { @@ -517,10 +690,16 @@ impl LanguageModel for XaiResponsesModel { let response = parsed.get("response").cloned().unwrap_or(Value::Null); if let Some(usage) = response.get("usage") && let Ok(u) = serde_json::from_value::(usage.clone()) { - final_usage = Some(convert_xai_responses_usage(&u)); + let mut converted = convert_xai_responses_usage(&u); + converted.raw = usage.as_object().cloned(); + final_usage = Some(converted); cost_in_usd_ticks = u.cost_in_usd_ticks; } + service_tier = response.get("service_tier").and_then(Value::as_str).map(str::to_string); + prompt_cache_key = response.get("prompt_cache_key").and_then(Value::as_str).map(str::to_string); + safety_identifier = response.get("safety_identifier").and_then(Value::as_str).map(str::to_string); + if event_type == "response.incomplete" { let reason = response .get("incomplete_details") @@ -560,8 +739,20 @@ impl LanguageModel for XaiResponsesModel { }; if let Some(usage) = response.get("usage") && let Ok(u) = serde_json::from_value::(usage.clone()) { - final_usage = Some(convert_xai_responses_usage(&u)); + let mut converted = convert_xai_responses_usage(&u); + converted.raw = usage.as_object().cloned(); + final_usage = Some(converted); } + if response.get("error").is_some_and(|error| !error.is_null()) { + yield Ok(StreamPart::Error { + error: super::xai_stream_error( + &parsed, + &stream_error_url, + stream_request_body.clone(), + stream_response_headers.clone(), + ), + }); + } continue; } @@ -579,9 +770,6 @@ impl LanguageModel for XaiResponsesModel { unified: FinishReasonUnified::Error, raw: Some("error".to_string()), }; - // A terminal error ends this stream; waiting for - // more events can hang on a source that keeps the - // connection open. break; } @@ -614,8 +802,16 @@ impl LanguageModel for XaiResponsesModel { // ── output_item.added / output_item.done ── if event_type == "response.output_item.added" || event_type == "response.output_item.done" + || matches!(event_type, + "response.image_generation_call.in_progress" + | "response.image_generation_call.generating" + | "response.image_generation_call.completed") { - let item = parsed.get("item").cloned().unwrap_or(Value::Null); + let item = if event_type.starts_with("response.image_generation_call.") { + json!({ "type": "image_generation_call", "id": parsed.get("item_id") }) + } else { + parsed.get("item").cloned().unwrap_or(Value::Null) + }; let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); let output_index = parsed.get("output_index").and_then(serde_json::Value::as_u64).unwrap_or(0); @@ -624,17 +820,20 @@ impl LanguageModel for XaiResponsesModel { if event_type == "response.output_item.done" { let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); let block_id = format!("reasoning-{part_id}"); - let encrypted = item.get("encrypted_content").and_then(|v| v.as_str()); + let encrypted = item.get("encrypted_content").and_then(|v| v.as_str()).filter(|v| !v.is_empty()); if !active_reasoning.contains_key(part_id) { active_reasoning.insert(part_id.to_string(), ()); yield Ok(StreamPart::ReasoningStart { id: block_id.clone(), - provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": part_id }))), + provider_metadata: Some(crate::xai::options::xai_metadata(if part_id.is_empty() { json!({}) } else { json!({ "itemId": part_id }) })), }); } - let mut meta = json!({ "itemId": part_id }); + let mut meta = json!({}); + if !part_id.is_empty() { + meta["itemId"] = json!(part_id); + } if let Some(ec) = encrypted { meta["reasoningEncryptedContent"] = json!(ec); } @@ -648,23 +847,24 @@ impl LanguageModel for XaiResponsesModel { } // ── file_search_call item ── - if item_type == "file_search_call" { + if item_type == "file_search_call" || item_type == "image_generation_call" { let tool_name = resolve_tool_name(item_type, None, &provider_tool_names); let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); + let input = if item_type == "image_generation_call" { "{}" } else { "" }; if !seen_tool_calls.contains(part_id) { seen_tool_calls.insert(part_id.to_string()); yield Ok(StreamPart::ToolInputStart { id: part_id.to_string(), tool_name: tool_name.clone(), - provider_executed: Some(true), + provider_executed: None, dynamic: None, title: None, provider_metadata: None, }); yield Ok(StreamPart::ToolInputDelta { id: part_id.to_string(), - delta: String::new(), + delta: input.to_string(), provider_metadata: None, }); yield Ok(StreamPart::ToolInputEnd { @@ -674,7 +874,7 @@ impl LanguageModel for XaiResponsesModel { yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: part_id.to_string(), tool_name: tool_name.clone(), - input: String::new(), + input: input.to_string(), provider_executed: Some(true), dynamic: None, provider_metadata: None, @@ -682,6 +882,10 @@ impl LanguageModel for XaiResponsesModel { } if event_type == "response.output_item.done" { + if item_type == "image_generation_call" { + yield Ok(StreamPart::ToolResult(image_generation_result(&item, tool_name))); + continue; + } let queries = item.get("queries").and_then(|v| v.as_array()).cloned().unwrap_or_default(); let results = item.get("results").and_then(|v| v.as_array()).map(|arr| { json!(arr.iter().map(|r| json!({ @@ -727,7 +931,7 @@ impl LanguageModel for XaiResponsesModel { yield Ok(StreamPart::ToolInputStart { id: part_id.to_string(), tool_name: tool_name.clone(), - provider_executed: Some(true), + provider_executed: None, dynamic: None, title: None, provider_metadata: None, @@ -755,7 +959,11 @@ impl LanguageModel for XaiResponsesModel { yield Ok(StreamPart::ToolResult(ToolResult { tool_call_id: part_id.to_string(), tool_name, - result: json!({}), + result: if item_type == "web_search_call" { + map_web_search_action(item.get("action").unwrap_or(&Value::Null)) + } else { + json!({}) + }, is_error: None, preliminary: None, dynamic: None, @@ -777,8 +985,8 @@ impl LanguageModel for XaiResponsesModel { && !text.is_empty() { let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); let block_id = format!("text-{part_id}"); - if !content_blocks.contains_key(&block_id) { - content_blocks.insert(block_id.clone(), false); + if !content_blocks.contains(&block_id) { + content_blocks.push(block_id.clone()); yield Ok(StreamPart::TextStart { id: block_id.clone(), provider_metadata: None}); yield Ok(StreamPart::TextDelta { id: block_id, @@ -846,9 +1054,13 @@ impl LanguageModel for XaiResponsesModel { // All other event types (web_search_call.in_progress, etc.) are ignored. } Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { + if error.is_recoverable_stream_error() { + if include_raw_chunks { + yield Ok(StreamPart::Raw { raw_value: Value::Null }); + } + yield Ok(StreamPart::Error { error }); + } else { + yield Err(error); return; } } @@ -856,31 +1068,21 @@ impl LanguageModel for XaiResponsesModel { } // Close any remaining open text blocks. - for (block_id, ended) in &content_blocks { - if !ended { - yield Ok(StreamPart::TextEnd { id: block_id.clone(), provider_metadata: None}); - } + for block_id in &content_blocks { + yield Ok(StreamPart::TextEnd { id: block_id.clone(), provider_metadata: None}); } // Final part: Finish. - let provider_meta = cost_in_usd_ticks.map(|cost| crate::xai::options::xai_metadata(json!({ "costInUsdTicks": cost }))); + let provider_meta = response_provider_metadata( + cost_in_usd_ticks, + service_tier.as_deref(), + prompt_cache_key.as_deref(), + safety_identifier.as_deref(), + ); yield Ok(StreamPart::Finish { finish_reason: final_finish_reason, - usage: final_usage.unwrap_or(Usage { - input_tokens: aimux_core::types::InputTokenUsage { - total: Some(0), - no_cache: Some(0), - cache_read: Some(0), - cache_write: Some(0), - }, - output_tokens: aimux_core::types::OutputTokenUsage { - total: Some(0), - text: Some(0), - reasoning: Some(0), - }, - raw: None, - }), + usage: final_usage.unwrap_or_else(zero_usage), provider_metadata: provider_meta, }); }; diff --git a/aimux-providers/src/xai/responses/types.rs b/aimux-providers/src/xai/responses/types.rs index ff8b93d01..b86bfd0c8 100644 --- a/aimux-providers/src/xai/responses/types.rs +++ b/aimux-providers/src/xai/responses/types.rs @@ -25,6 +25,12 @@ pub struct XaiResponsesResponse { pub usage: Option, #[serde(default)] pub status: Option, + #[serde(default)] + pub service_tier: Option, + #[serde(default)] + pub prompt_cache_key: Option, + #[serde(default)] + pub safety_identifier: Option, } /// Usage object for the Responses API (`xaiResponsesUsageSchema`). @@ -32,30 +38,36 @@ pub struct XaiResponsesResponse { pub struct XaiResponsesUsage { pub input_tokens: u64, pub output_tokens: u64, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub total_tokens: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub input_tokens_details: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub output_tokens_details: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub num_sources_used: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub num_server_side_tools_used: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub cost_in_usd_ticks: Option, + #[serde(flatten)] + pub extra: serde_json::Map, } #[derive(Debug, Deserialize, Clone, Default, Serialize)] pub struct XaiResponsesInputTokensDetails { - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub cached_tokens: Option, + #[serde(flatten)] + pub extra: serde_json::Map, } #[derive(Debug, Deserialize, Clone, Default, Serialize)] pub struct XaiResponsesOutputTokensDetails { - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub reasoning_tokens: Option, + #[serde(flatten)] + pub extra: serde_json::Map, } // ── Streaming ── diff --git a/aimux-providers/tests/alibaba_test.rs b/aimux-providers/tests/alibaba_test.rs index a1d2dcf0f..04e9e4b5f 100644 --- a/aimux-providers/tests/alibaba_test.rs +++ b/aimux-providers/tests/alibaba_test.rs @@ -34,7 +34,10 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; -use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; +use aimux_core::language_model_message::{ + AssistantPart, LanguageModelMessage, LanguageModelPrompt, ToolCallPart, ToolPart, + ToolResultOutput, ToolResultPart, +}; use aimux_core::options::CallOptions; use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; @@ -298,12 +301,63 @@ async fn sends_correct_request_body() { assert_eq!(body["messages"][0]["content"], "Hello"); } +/// TS: a tool-call assistant message round-trips as `tool_calls` and a tool +/// result as a `tool` role message (shared OpenAI conversion). +#[tokio::test] +async fn converts_tool_call_and_tool_result_messages() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(text_completion_body())) + .mount(&server) + .await; + + let model = make_provider(&server); + + let prompt: LanguageModelPrompt = vec![ + LanguageModelMessage::Assistant { + content: vec![AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "call-1".to_string(), + tool_name: "get_weather".to_string(), + input: json!({"location": "SF"}), + provider_executed: None, + provider_options: None, + })], + provider_options: None, + }, + LanguageModelMessage::Tool { + content: vec![ToolPart::ToolResult(ToolResultPart { + tool_call_id: "call-1".to_string(), + output: ToolResultOutput::Json { + value: json!({"temp": 72}), + provider_options: None, + }, + tool_name: "get_weather".into(), + provider_options: None, + })], + provider_options: None, + }, + ]; + let _ = model.do_generate(&default_options(prompt)).await.unwrap(); + + let requests = server.received_requests().await.expect("requests recorded"); + let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + // Assistant tool-call message. + assert_eq!(body["messages"][0]["role"], "assistant"); + assert_eq!(body["messages"][0]["tool_calls"][0]["id"], "call-1"); + assert_eq!( + body["messages"][0]["tool_calls"][0]["function"]["name"], + "get_weather" + ); + // Tool result message. + assert_eq!(body["messages"][1]["role"], "tool"); + assert_eq!(body["messages"][1]["tool_call_id"], "call-1"); +} + // ════════════════════════════════════════════════════════════════════════════ // Usage conversion (convert-alibaba-usage.test.ts + chat model usage tests) // ════════════════════════════════════════════════════════════════════════════ -/// TS (convert-alibaba-usage): cache tokens distribute correctly �?/// `cached_tokens` �?cacheRead, `cache_creation_input_tokens` �?cacheWrite, -/// and `noCache = prompt - cacheRead - cacheWrite`. #[tokio::test] async fn extracts_usage_with_cache_tokens() { let server = MockServer::start().await; diff --git a/aimux-providers/tests/amazon_bedrock_image_test.rs b/aimux-providers/tests/amazon_bedrock_image_test.rs index 8a8d30ade..a7f355ad6 100644 --- a/aimux-providers/tests/amazon_bedrock_image_test.rs +++ b/aimux-providers/tests/amazon_bedrock_image_test.rs @@ -17,7 +17,7 @@ fn bedrock_response() -> Value { async fn mock_bedrock(server: &MockServer, body: Value) { Mock::given(method("POST")) - .and(path("/model/amazon.nova-canvas-v1:0/invoke")) + .and(path("/model/amazon.nova-canvas-v1%3A0/invoke")) .respond_with(ResponseTemplate::new(200).set_body_json(body)) .mount(server) .await; @@ -156,7 +156,7 @@ async fn should_support_inpainting() { async fn should_handle_moderated_response() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/model/amazon.nova-canvas-v1:0/invoke")) + .and(path("/model/amazon.nova-canvas-v1%3A0/invoke")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "fe7256d1-50d9-4663-8592-85eaf002e80c", "status": "Request Moderated", diff --git a/aimux-providers/tests/anthropic_remaining_test.rs b/aimux-providers/tests/anthropic_remaining_test.rs new file mode 100644 index 000000000..5f265803f --- /dev/null +++ b/aimux-providers/tests/anthropic_remaining_test.rs @@ -0,0 +1,524 @@ +//! Remaining Anthropic provider tests — ported from the TS SDK suite. +//! +//! Sources: +//! +//! - `anthropic-provider.test.ts` — provider configuration (baseURL, auth, +//! custom provider name, `supportedUrls`). +//! - `anthropic-unknown-model-max-output-tokens.test.ts` — unknown-model +//! `max_tokens` defaulting + compatibility warning. +//! - `sanitize-json-schema.test.ts` — JSON Schema sanitization. +//! - `anthropic-language-model.test.ts` → `mid-conversation tool changes` +//! describe block. +//! +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; +use aimux_core::options::CallOptions; +use aimux_core::shared::provider_namespace; +use aimux_core::tool::{FunctionTool, Tool}; +use aimux_core::types::Warning; +use aimux_providers::anthropic::sanitize_json_schema::sanitize_json_schema; +use aimux_providers::anthropic::{ + AnthropicMessagesModel, AnthropicProviderSettings, create_anthropic, +}; +use mock_fetch::{Canned, EnvVar, MockFetch}; +use serde_json::{Value, json}; +use serial_test::serial; +use std::sync::Arc; + +/// The TS `TEST_PROMPT`: a single user text message "Hello". +fn test_prompt() -> LanguageModelPrompt { + vec![LanguageModelMessage::user_text("Hello")] +} + +fn default_options(prompt: LanguageModelPrompt) -> CallOptions { + CallOptions::new(prompt) +} + +fn make_model_with_config(settings: AnthropicProviderSettings) -> AnthropicMessagesModel { + create_anthropic(settings) + .unwrap() + .messages("claude-3-haiku-20240307") +} + +fn text_response(text: &str) -> Value { + json!({ + "id": "msg_123", "type": "message", "role": "assistant", + "content": [{ "type": "text", "text": text }], + "model": "claude-3-haiku-20240307", "stop_reason": "end_turn", + "stop_sequence": null, "usage": { "input_tokens": 1, "output_tokens": 1 } + }) +} + +fn mock() -> Arc { + MockFetch::new(vec![Canned::json(&text_response("Hi"))]) +} + +fn settings(fetch: &Arc) -> AnthropicProviderSettings { + AnthropicProviderSettings { + api_key: Some("test-api-key".into()), + fetch: Some(fetch.transport()), + ..Default::default() + } +} + +async fn request_url(settings: AnthropicProviderSettings, fetch: &Arc) -> String { + make_model_with_config(settings) + .do_generate(&default_options(test_prompt())) + .await + .unwrap(); + let requests = fetch.seen(); + assert_eq!(requests.len(), 1); + requests[0].url.clone() +} + +/// TS: "uses the default Anthropic base URL when not provided". +#[tokio::test] +#[serial] +async fn default_base_url_when_not_provided() { + let _env = EnvVar::set("ANTHROPIC_BASE_URL", None); + let fetch = mock(); + assert_eq!( + request_url(settings(&fetch), &fetch).await, + "https://api.anthropic.com/v1/messages" + ); +} + +/// TS: "uses ANTHROPIC_BASE_URL when set". +#[tokio::test] +#[serial] +async fn uses_anthropic_base_url_env_when_set() { + let _env = EnvVar::set( + "ANTHROPIC_BASE_URL", + Some("https://proxy.anthropic.example/v1/"), + ); + let fetch = mock(); + assert_eq!( + request_url(settings(&fetch), &fetch).await, + "https://proxy.anthropic.example/v1/messages" + ); +} + +/// TS: "normalizes a bare Anthropic API URL from ANTHROPIC_BASE_URL". +#[tokio::test] +#[serial] +async fn normalizes_bare_anthropic_url_from_env() { + let _env = EnvVar::set("ANTHROPIC_BASE_URL", Some("https://api.anthropic.com/")); + let fetch = mock(); + assert_eq!( + request_url(settings(&fetch), &fetch).await, + "https://api.anthropic.com/v1/messages" + ); +} + +/// TS: "normalizes a bare Anthropic API URL from the baseURL option". +#[tokio::test] +async fn normalizes_bare_anthropic_url_from_base_url_option() { + let fetch = mock(); + let settings = AnthropicProviderSettings { + base_url: Some("https://api.anthropic.com/".into()), + ..settings(&fetch) + }; + assert_eq!( + request_url(settings, &fetch).await, + "https://api.anthropic.com/v1/messages" + ); +} + +/// TS: "prefers the baseURL option over ANTHROPIC_BASE_URL". +#[tokio::test] +#[serial] +async fn prefers_base_url_option_over_env() { + let _env = EnvVar::set( + "ANTHROPIC_BASE_URL", + Some("https://env.anthropic.example/v1"), + ); + let fetch = mock(); + let settings = AnthropicProviderSettings { + base_url: Some("https://option.anthropic.example/v1/".into()), + ..settings(&fetch) + }; + assert_eq!( + request_url(settings, &fetch).await, + "https://option.anthropic.example/v1/messages" + ); +} + +/// TS: "rejects an empty baseURL option during provider creation". +#[test] +fn rejects_empty_base_url_option() { + let result = create_anthropic(AnthropicProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(String::new()), + ..Default::default() + }); + assert!( + matches!(result, Err(AiMuxError::InvalidArgument(msg)) if msg == "baseURL must be a non-empty string.") + ); +} + +/// TS: "sends Authorization Bearer header when authToken is provided". +#[tokio::test] +async fn sends_authorization_bearer_when_auth_token_provided() { + let fetch = mock(); + let settings = AnthropicProviderSettings { + auth_token: Some("test-auth-token".into()), + api_key: None, + ..settings(&fetch) + }; + request_url(settings, &fetch).await; + let requests = fetch.seen(); + assert_eq!( + requests[0].headers.get("authorization").map(String::as_str), + Some("Bearer test-auth-token") + ); + assert!(!requests[0].headers.contains_key("x-api-key")); +} + +/// TS: "throws error when both apiKey and authToken options are provided". +#[test] +fn throws_when_both_api_key_and_auth_token_provided() { + let result = create_anthropic(AnthropicProviderSettings { + api_key: Some("test-api-key".into()), + auth_token: Some("test-auth-token".into()), + ..Default::default() + }); + assert!( + matches!(result, Err(AiMuxError::InvalidArgument(msg)) if msg == "Both apiKey and authToken were provided. Please use only one authentication method.") + ); +} + +/// TS: "should use custom provider name when specified". +#[test] +fn uses_custom_provider_name_when_specified() { + let model = make_model_with_config(AnthropicProviderSettings { + name: Some("my-proxy".into()), + api_key: Some("test-api-key".into()), + ..Default::default() + }); + assert_eq!(model.provider(), "my-proxy"); +} + +/// TS: "should default to anthropic.messages when name not specified". +#[test] +fn defaults_to_anthropic_messages_when_not_specified() { + let model = make_model_with_config(AnthropicProviderSettings { + api_key: Some("test-api-key".into()), + ..Default::default() + }); + assert_eq!(model.provider(), "anthropic.messages"); +} + +/// TS: "should support image/* URLs". +#[test] +fn supports_image_urls() { + let model = make_model_with_config(AnthropicProviderSettings::default()); + assert!(model.supported_urls().0["image/*"][0].is_match("https://example.com/image.png")); +} + +/// TS: "should support application/pdf URLs". +#[test] +fn supports_pdf_urls() { + let model = make_model_with_config(AnthropicProviderSettings::default()); + assert!( + model.supported_urls().0["application/pdf"][0].is_match("https://arxiv.org/pdf/2401.00001") + ); +} + +/// TS: "should warn when using the default max output token limit". +#[tokio::test] +async fn warns_when_using_default_max_output_token_limit() { + let server = mock(); + + let model = create_anthropic(settings(&server)) + .unwrap() + .messages("future-model"); + let result = model + .do_generate(&default_options(test_prompt())) + .await + .expect("do_generate should succeed"); + + let requests = server.seen(); + let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!(body["model"], "future-model"); + assert_eq!(body["max_tokens"], 4096); + + assert_eq!(result.warnings.len(), 1); + match &result.warnings[0] { + Warning::Compatibility { feature, details } => { + assert_eq!(feature, "maxOutputTokens"); + assert_eq!( + details.as_deref(), + Some( + "The model \"future-model\" is unknown. The max output tokens have been \ + limited to 4096. Set maxOutputTokens explicitly to override this limit." + ) + ); + } + other => panic!("expected Compatibility warning, got {other:?}"), + } +} + +/// TS: "should not warn when max output tokens are provided". +#[tokio::test] +async fn does_not_warn_when_max_output_tokens_provided() { + let server = mock(); + + let model = create_anthropic(settings(&server)) + .unwrap() + .messages("future-model"); + let mut opts = default_options(test_prompt()); + opts.max_output_tokens = Some(123456); + let result = model + .do_generate(&opts) + .await + .expect("do_generate should succeed"); + + let requests = server.seen(); + let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!(body["model"], "future-model"); + assert_eq!(body["max_tokens"], 123456); + assert!( + result.warnings.is_empty(), + "warnings: {:?}", + result.warnings + ); +} + +/// TS: unknown model current-generation default and compatibility warning. +#[tokio::test] +async fn uses_current_gen_default_and_warns_for_unknown_model() { + let server = mock(); + + let model = create_anthropic(settings(&server)) + .unwrap() + .messages("claude-future-9"); + let result = model + .do_generate(&default_options(test_prompt())) + .await + .expect("do_generate should succeed"); + + let requests = server.seen(); + let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!(body["model"], "claude-future-9"); + assert_eq!(body["max_tokens"], 128000); + + assert_eq!(result.warnings.len(), 1); + match &result.warnings[0] { + Warning::Compatibility { feature, details } => { + assert_eq!(feature, "maxOutputTokens"); + assert_eq!( + details.as_deref(), + Some( + "The model \"claude-future-9\" is unknown. The max output tokens have been \ + limited to 128000. Set maxOutputTokens explicitly to override this limit." + ) + ); + } + other => panic!("expected Compatibility warning, got {other:?}"), + } +} + +// ═════════════════════════════════════════════════════════════════════════════ +// sanitize-json-schema (TS: sanitize-json-schema.test.ts) +// ═════════════════════════════════════════════════════════════════════════════ + +/// TS: "strips unsupported number constraints and adds readable descriptions". +#[test] +fn sanitize_strips_unsupported_number_constraints() { + let schema = json!({ + "type": "object", + "properties": { + "recurringIntervalMinutes": { + "type": "number", + "exclusiveMinimum": 0, + "minimum": 1, + "maximum": 60, + "exclusiveMaximum": 120 + } + }, + "required": ["recurringIntervalMinutes"], + "additionalProperties": false + }); + + let expected = json!({ + "additionalProperties": false, + "properties": { + "recurringIntervalMinutes": { + "description": + "minimum: 1; maximum: 60; exclusive minimum: 0; exclusive maximum: 120.", + "type": "number" + } + }, + "required": ["recurringIntervalMinutes"], + "type": "object" + }); + + assert_eq!(sanitize_json_schema(&schema), expected); +} + +/// TS: "strips unsupported string constraints and unsupported formats". +#[test] +fn sanitize_strips_unsupported_string_constraints_and_formats() { + let schema = json!({ + "type": "object", + "properties": { + "slug": { + "type": "string", + "description": "A URL slug", + "minLength": 1, + "maxLength": 20, + "pattern": "^[a-z0-9-]+$", + "format": "regex" + } + } + }); + + let expected = json!({ + "additionalProperties": false, + "properties": { + "slug": { + "description": + "A URL slug\nmin length: 1; max length: 20; pattern: ^[a-z0-9-]+$; format: regex.", + "type": "string" + } + }, + "type": "object" + }); + + assert_eq!(sanitize_json_schema(&schema), expected); +} + +/// TS: "recursively sanitizes arrays, definitions, and composition schemas". +#[test] +fn sanitize_recursively_handles_arrays_defs_and_composition() { + let schema = json!({ + "type": "object", + "$defs": { + "PositiveInteger": { "type": "integer", "minimum": 1 } + }, + "properties": { + "count": { "$ref": "#/$defs/PositiveInteger" }, + "tags": { + "type": "array", + "minItems": 2, + "maxItems": 4, + "uniqueItems": true, + "items": { + "anyOf": [ + { "type": "string", "minLength": 1 }, + { "type": "number", "maximum": 10 } + ] + } + } + } + }); + + let expected = json!({ + "$defs": { + "PositiveInteger": { + "description": "minimum: 1.", + "type": "integer" + } + }, + "additionalProperties": false, + "properties": { + "count": { "$ref": "#/$defs/PositiveInteger" }, + "tags": { + "description": "min items: 2; max items: 4; unique items: true.", + "items": { + "anyOf": [ + { "description": "min length: 1.", "type": "string" }, + { "description": "maximum: 10.", "type": "number" } + ] + }, + "type": "array" + } + }, + "type": "object" + }); + + assert_eq!(sanitize_json_schema(&schema), expected); +} + +/// TS: "converts oneOf to anyOf". +#[test] +fn sanitize_converts_one_of_to_any_of() { + let schema = json!({ + "oneOf": [ + { "type": "string", "minLength": 1 }, + { "type": "number", "minimum": 0 } + ] + }); + + let expected = json!({ + "anyOf": [ + { "description": "min length: 1.", "type": "string" }, + { "description": "minimum: 0.", "type": "number" } + ] + }); + + assert_eq!(sanitize_json_schema(&schema), expected); +} + +/// TS: "does not mutate the input schema". +#[test] +fn sanitize_does_not_mutate_input_schema() { + let schema = json!({ + "type": "object", + "properties": { + "value": { "type": "number", "exclusiveMinimum": 0 } + } + }); + + let snapshot = schema.clone(); + let _ = sanitize_json_schema(&schema); + assert_eq!(schema, snapshot); +} + +/// TS: "should send tool change blocks and the beta header". +#[tokio::test] +async fn sends_tool_change_blocks_and_beta_header() { + let fetch = mock(); + let model = create_anthropic(settings(&fetch)) + .unwrap() + .messages("claude-opus-4-8"); + let mut opts = default_options(vec![ + LanguageModelMessage::user_text("Say OK."), + LanguageModelMessage::System { + content: String::new(), + provider_options: Some( + provider_namespace( + "anthropic", + json!({ + "toolChanges": [{"type": "tool_removal", "toolName": "get_weather"}] + }), + ) + .unwrap(), + ), + }, + ]); + opts.tools = Some(vec![Tool::Function( + FunctionTool::new( + "get_weather", + json!({ + "type": "object", "properties": {"city": {"type": "string"}} + }), + ) + .with_description("Get weather"), + )]); + model.do_generate(&opts).await.unwrap(); + let requests = fetch.seen(); + let body = requests[0].json_body(); + assert!(body["messages"].as_array().unwrap().contains(&json!({ + "role": "system", "content": [{ + "type": "tool_removal", "tool": {"type": "tool_reference", "name": "get_weather"} + }] + }))); + assert!( + requests[0].headers["anthropic-beta"].contains("mid-conversation-tool-changes-2026-07-01") + ); +} diff --git a/aimux-providers/tests/anthropic_upstream_e2e_test.rs b/aimux-providers/tests/anthropic_upstream_e2e_test.rs new file mode 100644 index 000000000..dc4548dda --- /dev/null +++ b/aimux-providers/tests/anthropic_upstream_e2e_test.rs @@ -0,0 +1,253 @@ +//! Small factory-level subset of the ported Anthropic upstream cases. + +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use aimux_core::AiMuxError; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::{CallOptions, ResponseFormat}; +use aimux_core::result::GenerateContent; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::{FunctionTool, Tool}; +use aimux_core::types::FinishReasonUnified; +use aimux_providers::anthropic::{AnthropicProvider, AnthropicProviderSettings, create_anthropic}; +use futures::StreamExt; +use mock_fetch::{Canned, MockFetch}; +use serde_json::{Value, json}; + +fn provider(fetch: &std::sync::Arc) -> AnthropicProvider { + create_anthropic(AnthropicProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some("https://test.invalid/v1".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} + +fn options() -> CallOptions { + CallOptions { + max_output_tokens: Some(4096), + ..CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) + } +} + +fn body(content: Value, stop_reason: &str, usage: Value) -> Value { + json!({ + "id": "msg_017TfcQ4AgGxKyBduUpqYPZn", "type": "message", + "role": "assistant", "content": content, "model": "test-model", + "stop_reason": stop_reason, "stop_sequence": null, "usage": usage + }) +} + +fn request() -> Value { + json!({ + "model": "test-model", "max_tokens": 4096, + "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + }) +} + +/// TS: "should use native structured output without a JSON tool fallback" (anthropic/src/anthropic-language-model.test.ts) +#[tokio::test] +async fn generate_native_structured_output() { + let text = "{\"name\":\"Classic Lasagna\"}"; + let response = body( + json!([{"type": "text", "text": text}]), + "end_turn", + json!({"input_tokens": 371, "output_tokens": 629}), + ); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let schema = json!({"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"], "additionalProperties": false}); + let mut options = options(); + options.max_output_tokens = Some(100); + options.response_format = Some(ResponseFormat::Json { + schema: Some(schema.clone()), + name: None, + description: None, + }); + let result = provider(&fetch) + .messages("claude-future-9") + .do_generate(&options) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "POST"); + assert_eq!(seen[0].url, "https://test.invalid/v1/messages"); + assert_eq!(seen[0].headers["x-api-key"], "test-api-key"); + let mut expected = request(); + expected["model"] = json!("claude-future-9"); + expected["max_tokens"] = json!(100); + expected["output_config"] = json!({"format": {"type": "json_schema", "schema": schema}}); + assert_eq!(seen[0].json_body(), expected); + assert_eq!( + result.content, + vec![GenerateContent::Text { + text: text.into(), + provider_metadata: None + }] + ); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(result.finish_reason.raw.as_deref(), Some("end_turn")); + assert_eq!(result.usage.input_tokens.total, Some(371)); + assert_eq!(result.usage.input_tokens.no_cache, Some(371)); + assert_eq!(result.usage.input_tokens.cache_read, Some(0)); + assert_eq!(result.usage.input_tokens.cache_write, Some(0)); + assert_eq!(result.usage.output_tokens.total, Some(629)); + assert_eq!(result.request.unwrap().body, Some(expected)); + let returned = result.response.unwrap(); + assert_eq!(returned.model_id.as_deref(), Some("test-model")); + assert_eq!(returned.id.as_deref(), Some("msg_017TfcQ4AgGxKyBduUpqYPZn")); + assert_eq!(returned.body, Some(response)); +} + +/// TS: "should stream text deltas" (anthropic/src/anthropic-language-model.test.ts) +#[tokio::test] +async fn stream_text_deltas() { + let events = [ + json!({"type": "message_start", "message": {"id": "msg_01KfpJoAEabmH2iHRRFjQMAG", "type": "message", "role": "assistant", "content": [], "model": "test-model", "stop_reason": null, "stop_sequence": null, "usage": {"input_tokens": 17, "output_tokens": 1}}}), + json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}), + json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": ", "}}), + json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "World!"}}), + json!({"type": "content_block_stop", "index": 0}), + json!({"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": null}, "usage": {"output_tokens": 227}}), + json!({"type": "message_stop"}), + ]; + let fetch = MockFetch::new(vec![Canned { + status: 200, + headers: vec![("content-type".into(), "text/event-stream".into())], + body: events + .iter() + .map(|event| format!("data: {event}\n\n")) + .collect::() + .into_bytes(), + }]); + let result = provider(&fetch) + .messages("test-model") + .do_stream(&options()) + .await + .unwrap(); + let mut expected = request(); + expected["stream"] = json!(true); + assert_eq!(fetch.seen().len(), 1); + assert_eq!(fetch.seen()[0].json_body(), expected); + assert_eq!(result.request.unwrap().body, Some(expected)); + let parts = result + .stream + .collect::>() + .await + .into_iter() + .collect::, _>>() + .unwrap(); + assert_eq!(parts.len(), 8); + assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); + assert!( + matches!(&parts[1], StreamPart::ResponseMetadata(metadata) if metadata.id.as_deref() == Some("msg_01KfpJoAEabmH2iHRRFjQMAG") && metadata.model_id.as_deref() == Some("test-model") && metadata.timestamp.is_none()) + ); + assert!( + matches!(&parts[2], StreamPart::TextStart { id, provider_metadata } if id == "0" && provider_metadata.is_none()) + ); + for (part, expected_delta) in parts[3..6].iter().zip(["Hello", ", ", "World!"]) { + assert!( + matches!(part, StreamPart::TextDelta { id, delta, provider_metadata } if id == "0" && delta == expected_delta && provider_metadata.is_none()) + ); + } + assert!( + matches!(&parts[6], StreamPart::TextEnd { id, provider_metadata } if id == "0" && provider_metadata.is_none()) + ); + let StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } = &parts[7] + else { + panic!("expected finish"); + }; + assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(finish_reason.raw.as_deref(), Some("end_turn")); + assert_eq!(usage.input_tokens.total, Some(17)); + assert_eq!(usage.input_tokens.no_cache, Some(17)); + assert_eq!(usage.input_tokens.cache_read, Some(0)); + assert_eq!(usage.input_tokens.cache_write, Some(0)); + assert_eq!(usage.output_tokens.total, Some(227)); + assert_eq!( + usage.raw, + json!({"input_tokens": 17, "output_tokens": 227}) + .as_object() + .cloned() + ); + assert_eq!(provider_metadata.as_ref().unwrap()["anthropic"], serde_json::from_value::>(json!({"container": null, "contextManagement": null, "iterations": null, "stopSequence": null, "usage": {"input_tokens": 17, "output_tokens": 227}})).unwrap()); +} + +/// TS: "should extract tool calls" (anthropic/src/anthropic-language-model.test.ts) +#[tokio::test] +async fn generate_tool_call() { + let response = body( + json!([ + {"type": "text", "text": "Some text\n\n"}, + {"type": "tool_use", "id": "toolu_1", "name": "test-tool", "input": {"value": "example value"}} + ]), + "tool_use", + json!({"input_tokens": 4, "output_tokens": 30}), + ); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let schema = json!({"type": "object", "properties": {"value": {"type": "string"}}, "required": ["value"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}); + let mut options = options(); + options.tools = Some(vec![Tool::Function(FunctionTool::new( + "test-tool", + schema.clone(), + ))]); + let result = provider(&fetch) + .messages("test-model") + .do_generate(&options) + .await + .unwrap(); + let mut expected = request(); + expected["tools"] = json!([{"name": "test-tool", "input_schema": schema}]); + assert_eq!(fetch.seen().len(), 1); + assert_eq!(fetch.seen()[0].json_body(), expected); + assert_eq!(result.content.len(), 2); + assert_eq!( + result.content[0], + GenerateContent::Text { + text: "Some text\n\n".into(), + provider_metadata: None + } + ); + let GenerateContent::ToolCall(call) = &result.content[1] else { + panic!("expected tool call"); + }; + assert_eq!(call.tool_call_id, "toolu_1"); + assert_eq!(call.tool_name, "test-tool"); + assert_eq!(call.input, "{\"value\":\"example value\"}"); + assert_eq!(call.provider_metadata, None); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); + assert_eq!(result.finish_reason.raw.as_deref(), Some("tool_use")); + assert_eq!(result.usage.input_tokens.total, Some(4)); + assert_eq!(result.usage.output_tokens.total, Some(30)); +} + +/// TS: "should parse overloaded error" (anthropic/src/anthropic-error.test.ts) +#[tokio::test] +async fn overloaded_error_response() { + let mut canned = Canned::json( + &json!({"type": "error", "error": {"details": null, "type": "overloaded_error", "message": "Overloaded"}}), + ); + canned.status = 529; + let fetch = MockFetch::new(vec![canned]); + let error = provider(&fetch) + .messages("test-model") + .do_generate(&options()) + .await + .unwrap_err(); + assert_eq!(fetch.seen().len(), 1); + assert_eq!(fetch.seen()[0].json_body(), request()); + let AiMuxError::ApiCall(error) = error else { + panic!("expected provider API error"); + }; + assert_eq!(error.message, "Overloaded"); + assert_eq!(error.provider_code.as_deref(), Some("overloaded_error")); + assert_eq!(error.status_code, Some(529)); +} diff --git a/aimux-providers/tests/azure_upstream_e2e_test.rs b/aimux-providers/tests/azure_upstream_e2e_test.rs new file mode 100644 index 000000000..923143b45 --- /dev/null +++ b/aimux-providers/tests/azure_upstream_e2e_test.rs @@ -0,0 +1,373 @@ +//! Small end-to-end subset of the pinned Azure provider tests. + +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_core::result::GenerateContent; +use aimux_core::speech_model::{AudioData, SpeechCallOptions, SpeechModel}; +use aimux_core::stream_part::StreamPart; +use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; +use aimux_provider_utils::Resolvable; +use aimux_providers::azure::{AzureOpenAIProvider, AzureOpenAIProviderSettings, create_azure}; +use futures::StreamExt; +use mock_fetch::{Canned, MockFetch}; +use serde_json::{Value, json}; +use std::sync::Arc; + +fn provider(fetch: &Arc) -> AzureOpenAIProvider { + create_azure(AzureOpenAIProviderSettings { + resource_name: Some("test-resource".into()), + api_key: Some(Resolvable::Value("test-api-key".into())), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} + +fn prompt() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +fn assert_request(fetch: &MockFetch, endpoint: &str) -> Value { + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "POST"); + assert_eq!( + seen[0].url, + format!("https://test-resource.openai.azure.com/openai/v1/{endpoint}?api-version=v1") + ); + assert_eq!(seen[0].headers["api-key"], "test-api-key"); + seen[0].json_body() +} +fn chat_response() -> Value { + json!({"id":"chatcmpl-1","object":"chat.completion","created":1711115037,"model":"test-deployment","choices":[{"index":0,"message":{"role":"assistant","content":"Hello World!"},"finish_reason":"stop"}],"usage":{"prompt_tokens":4,"completion_tokens":30,"total_tokens":34}}) +} + +fn text_fixture() -> Value { + json!({"id":"resp_0d6bb044bb6ff37200698c51948054819385e24e2ad931ae6e","object":"response","created_at":1770803604,"status":"completed","model":"test-deployment","output":[{"id":"msg_0d6bb044bb6ff37200698c51952e288193beb9044db4d8c810","type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"Word","annotations":[],"logprobs":[]}]}],"usage":{"input_tokens":11,"input_tokens_details":{"cached_tokens":0},"output_tokens":11,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":22}}) +} + +fn image_response() -> Value { + json!({"created":1733837122,"data":[{"revised_prompt":"A charming visual illustration of a baby sea otter swimming joyously.","b64_json":"base64-image-1"},{"revised_prompt":"A charming visual illustration of a baby sea otter swimming joyously.","b64_json":"base64-image-2"}]}) +} + +fn text_chunks() -> &'static str { + r###"data: {"type":"response.created","sequence_number":0,"response":{"id":"resp_02ce8deeb6197db200698c5196e9588197a572bbea62d38cd1","object":"response","created_at":1770803606,"status":"in_progress","background":false,"completed_at":null,"content_filters":null,"error":null,"incomplete_details":null,"instructions":null,"max_output_tokens":null,"max_tool_calls":null,"model":"test-deployment","output":[],"parallel_tool_calls":true,"previous_response_id":null,"prompt_cache_key":null,"prompt_cache_retention":null,"reasoning":{"effort":"none","summary":null},"safety_identifier":null,"service_tier":"auto","store":true,"temperature":1,"text":{"format":{"type":"text"},"verbosity":"medium"},"tool_choice":"auto","tools":[],"top_logprobs":0,"top_p":1,"truncation":"disabled","usage":null,"user":null,"metadata":{}}} + +data: {"type":"response.in_progress","sequence_number":1,"response":{"id":"resp_02ce8deeb6197db200698c5196e9588197a572bbea62d38cd1","object":"response","created_at":1770803606,"status":"in_progress","background":false,"completed_at":null,"content_filters":null,"error":null,"incomplete_details":null,"instructions":null,"max_output_tokens":null,"max_tool_calls":null,"model":"test-deployment","output":[],"parallel_tool_calls":true,"previous_response_id":null,"prompt_cache_key":null,"prompt_cache_retention":null,"reasoning":{"effort":"none","summary":null},"safety_identifier":null,"service_tier":"auto","store":true,"temperature":1,"text":{"format":{"type":"text"},"verbosity":"medium"},"tool_choice":"auto","tools":[],"top_logprobs":0,"top_p":1,"truncation":"disabled","usage":null,"user":null,"metadata":{}}} + +data: {"type":"response.output_item.added","sequence_number":2,"output_index":0,"item":{"id":"msg_02ce8deeb6197db200698c5198ca0c81979bedbe6c98a8ab93","type":"message","status":"in_progress","content":[],"role":"assistant"}} + +data: {"type":"response.content_part.added","sequence_number":3,"item_id":"msg_02ce8deeb6197db200698c5198ca0c81979bedbe6c98a8ab93","output_index":0,"content_index":0,"part":{"type":"output_text","annotations":[],"logprobs":[],"text":""}} + +data: {"type":"response.output_text.delta","sequence_number":4,"item_id":"msg_02ce8deeb6197db200698c5198ca0c81979bedbe6c98a8ab93","output_index":0,"content_index":0,"delta":"Hello","logprobs":[],"obfuscation":"VLMdSoETMf5"} + +data: {"type":"response.output_text.done","sequence_number":5,"item_id":"msg_02ce8deeb6197db200698c5198ca0c81979bedbe6c98a8ab93","output_index":0,"content_index":0,"text":"Hello","logprobs":[]} + +data: {"type":"response.content_part.done","sequence_number":6,"item_id":"msg_02ce8deeb6197db200698c5198ca0c81979bedbe6c98a8ab93","output_index":0,"content_index":0,"part":{"type":"output_text","annotations":[],"logprobs":[],"text":"Hello"}} + +data: {"type":"response.output_item.done","sequence_number":7,"output_index":0,"item":{"id":"msg_02ce8deeb6197db200698c5198ca0c81979bedbe6c98a8ab93","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":"Hello"}],"role":"assistant"}} + +data: {"type":"response.completed","sequence_number":8,"response":{"id":"resp_02ce8deeb6197db200698c5196e9588197a572bbea62d38cd1","object":"response","created_at":1770803606,"status":"completed","background":false,"completed_at":1770803608,"content_filters":[{"blocked":false,"source_type":"prompt","content_filter_raw":null,"content_filter_results":{"self_harm":{"filtered":false,"severity":"safe"},"sexual":{"filtered":false,"severity":"safe"},"violence":{"filtered":false,"severity":"safe"},"jailbreak":{"filtered":false,"detected":false},"hate":{"filtered":false,"severity":"safe"}},"content_filter_offsets":{"start_offset":1429,"end_offset":1447,"check_offset":0}},{"blocked":false,"source_type":"completion","content_filter_raw":null,"content_filter_results":{"hate":{"filtered":false,"severity":"safe"},"sexual":{"filtered":false,"severity":"safe"},"self_harm":{"filtered":false,"severity":"safe"},"violence":{"filtered":false,"severity":"safe"},"protected_material_text":{"filtered":false,"detected":false},"protected_material_code":{"filtered":false,"detected":false}},"content_filter_offsets":{"start_offset":0,"end_offset":5,"check_offset":0}}],"error":null,"incomplete_details":null,"instructions":null,"max_output_tokens":null,"max_tool_calls":null,"model":"test-deployment","output":[{"id":"msg_02ce8deeb6197db200698c5198ca0c81979bedbe6c98a8ab93","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":"Hello"}],"role":"assistant"}],"parallel_tool_calls":true,"previous_response_id":null,"prompt_cache_key":null,"prompt_cache_retention":null,"reasoning":{"effort":"none","summary":null},"safety_identifier":null,"service_tier":"default","store":true,"temperature":1,"text":{"format":{"type":"text"},"verbosity":"medium"},"tool_choice":"auto","tools":[],"top_logprobs":0,"top_p":1,"truncation":"disabled","usage":{"input_tokens":11,"input_tokens_details":{"cached_tokens":0},"output_tokens":11,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":22},"user":null,"metadata":{}}} + +data: [DONE] + +"### +} + +/// TS: "should use the baseURL correctly" (@ai-sdk/azure/src/azure-openai-provider.test.ts) +#[tokio::test] +async fn chat_generate() { + let fetch = MockFetch::new(vec![Canned::json(&chat_response())]); + let result = provider(&fetch) + .chat("test-deployment") + .do_generate(&prompt()) + .await + .unwrap(); + let body = assert_request(&fetch, "chat/completions"); + assert_eq!(body["model"], "test-deployment"); + assert_eq!(body["messages"], json!([{"role":"user","content":"Hello"}])); + assert_eq!(result.content.len(), 1); + assert!( + matches!(&result.content[0], GenerateContent::Text { text, .. } if text == "Hello World!") + ); + assert_eq!(result.usage.input_tokens.total, Some(4)); + assert_eq!(result.usage.output_tokens.total, Some(30)); + assert_eq!( + result.response.as_ref().unwrap().id.as_deref(), + Some("chatcmpl-1") + ); +} + +/// TS: "should extract text content" (@ai-sdk/azure/src/azure-openai-provider.test.ts) +/// TS: "should extract usage" (@ai-sdk/azure/src/azure-openai-provider.test.ts) +#[tokio::test] +async fn responses_generate() { + let fixture = text_fixture(); + let fetch = MockFetch::new(vec![Canned::json(&fixture)]); + let result = provider(&fetch) + .responses("test-deployment") + .do_generate(&prompt()) + .await + .unwrap(); + let body = assert_request(&fetch, "responses"); + assert_eq!(body["model"], "test-deployment"); + assert_eq!( + body["input"], + json!([{"role":"user","content":[{"type":"input_text","text":"Hello"}]}]) + ); + assert_eq!(result.content.len(), 1); + match &result.content[0] { + GenerateContent::Text { + text, + provider_metadata, + } => { + assert_eq!(text, "Word"); + assert_eq!( + provider_metadata.as_ref().unwrap()["azure"]["itemId"], + "msg_0d6bb044bb6ff37200698c51952e288193beb9044db4d8c810" + ); + } + _ => panic!("expected text"), + } + assert_eq!( + result.provider_metadata.as_ref().unwrap()["azure"]["responseId"], + fixture["id"] + ); + assert_eq!(result.usage.input_tokens.total, Some(11)); + assert_eq!(result.usage.input_tokens.cache_read, Some(0)); + assert_eq!(result.usage.input_tokens.no_cache, Some(11)); + assert_eq!(result.usage.output_tokens.total, Some(11)); + assert_eq!(result.usage.output_tokens.reasoning, Some(0)); + assert_eq!(result.usage.output_tokens.text, Some(11)); + assert_eq!(result.usage.raw, fixture["usage"].as_object().cloned()); + assert_eq!( + result.response.as_ref().unwrap().model_id.as_deref(), + Some("test-deployment") + ); +} + +/// TS: "should extract tool call content" (@ai-sdk/azure/src/azure-openai-provider.test.ts) +#[tokio::test] +async fn responses_tool_call() { + let mut fixture = text_fixture(); + fixture["output"] = json!([{"type":"function_call","id":"fc_0a2fa1b539ba14ba00698c519ebab0819494302fc0b5c31440","call_id":"call_YunNGbIwdVJ2i0y0Mybva4Pw","name":"weather","arguments":"{\"location\":\"San Francisco\"}","status":"completed"}]); + let fetch = MockFetch::new(vec![Canned::json(&fixture)]); + let result = provider(&fetch) + .responses("test-deployment") + .do_generate(&prompt()) + .await + .unwrap(); + assert_eq!( + assert_request(&fetch, "responses")["input"], + json!([{"role":"user","content":[{"type":"input_text","text":"Hello"}]}]) + ); + assert_eq!(result.content.len(), 1); + match &result.content[0] { + GenerateContent::ToolCall(call) => { + assert_eq!(call.tool_call_id, "call_YunNGbIwdVJ2i0y0Mybva4Pw"); + assert_eq!(call.tool_name, "weather"); + assert_eq!(call.input, "{\"location\":\"San Francisco\"}"); + } + _ => panic!("expected tool call"), + } +} + +/// TS: "should stream text content" (@ai-sdk/azure/src/azure-openai-provider.test.ts) +#[tokio::test] +async fn responses_stream() { + let fetch = MockFetch::new(vec![Canned { + status: 200, + headers: vec![("content-type".into(), "text/event-stream".into())], + body: text_chunks().as_bytes().to_vec(), + }]); + let result = provider(&fetch) + .responses("test-deployment") + .do_stream(&prompt()) + .await + .unwrap(); + let body = assert_request(&fetch, "responses"); + assert_eq!(body["stream"], true); + assert_eq!( + body["input"], + json!([{"role":"user","content":[{"type":"input_text","text":"Hello"}]}]) + ); + let parts = result + .stream + .collect::>() + .await + .into_iter() + .collect::, _>>() + .unwrap(); + assert_eq!(parts.len(), 6); + assert!(matches!(parts[0], StreamPart::StreamStart { .. })); + assert!(matches!(parts[1], StreamPart::ResponseMetadata(_))); + assert!(matches!(parts[2], StreamPart::TextStart { .. })); + assert!(matches!(&parts[3], StreamPart::TextDelta { delta, .. } if delta == "Hello")); + match &parts[4] { + StreamPart::TextEnd { + provider_metadata, .. + } => assert_eq!( + provider_metadata.as_ref().unwrap()["azure"]["itemId"], + "msg_02ce8deeb6197db200698c5198ca0c81979bedbe6c98a8ab93" + ), + _ => panic!("expected text end"), + } + match &parts[5] { + StreamPart::Finish { usage, .. } => { + assert_eq!(usage.input_tokens.total, Some(11)); + assert_eq!(usage.output_tokens.total, Some(11)); + } + _ => panic!("expected finish"), + } +} + +/// TS: "should set the correct api version" (@ai-sdk/azure/src/azure-openai-provider.test.ts) +#[tokio::test] +async fn embedding_generate() { + let fetch = MockFetch::new(vec![Canned::json( + &json!({"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2,0.3,0.4,0.5]},{"object":"embedding","index":1,"embedding":[0.6,0.7,0.8,0.9,1.0]}],"model":"my-embedding","usage":{"prompt_tokens":8,"total_tokens":8}}), + )]); + let mut options = EmbeddingCallOptions::new("sunny day at the beach"); + options.values.push("rainy day in the city".into()); + let result = provider(&fetch) + .embedding("my-embedding") + .do_embed(&options) + .await + .unwrap(); + assert_eq!( + assert_request(&fetch, "embeddings"), + json!({"model":"my-embedding","input":options.values,"encoding_format":"float"}) + ); + assert_eq!( + result.embeddings, + vec![vec![0.1, 0.2, 0.3, 0.4, 0.5], vec![0.6, 0.7, 0.8, 0.9, 1.0]] + ); + assert_eq!(result.usage.unwrap().tokens, 8); +} + +/// TS: "should send the correct request body" (@ai-sdk/azure/src/azure-openai-provider.test.ts) +/// TS: "should extract the generated images" (@ai-sdk/azure/src/azure-openai-provider.test.ts) +#[tokio::test] +async fn image_generate() { + let fetch = MockFetch::new(vec![Canned::json(&image_response())]); + let mut options = ImageCallOptions::new("A cute baby sea otter"); + options.n = 2; + options.size = Some("1024x1024".parse().unwrap()); + options.provider_options = + serde_json::from_value(json!({"openai":{"style":"natural"}})).unwrap(); + let result = provider(&fetch) + .image("test-deployment") + .do_generate(&options) + .await + .unwrap(); + assert_eq!( + assert_request(&fetch, "images/generations"), + json!({"model":"test-deployment","prompt":"A cute baby sea otter","n":2,"size":"1024x1024","response_format":"b64_json","style":"natural"}) + ); + match result.images { + ImageOutputs::Base64(images) => assert_eq!(images, ["base64-image-1", "base64-image-2"]), + _ => panic!("expected base64 images"), + } +} + +/// TS: "should use correct URL format" (@ai-sdk/azure/src/azure-openai-provider.test.ts) +#[tokio::test] +async fn transcription_generate() { + let fetch = MockFetch::new(vec![Canned::json( + &json!({"text":"Hello, world!","segments":[],"language":"en","duration":5.0}), + )]); + let result = provider(&fetch) + .transcription("test-deployment") + .do_generate(&TranscriptionCallOptions::new( + AudioInput::Binary(vec![]), + "audio/wav", + )) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "POST"); + assert_eq!( + seen[0].url, + "https://test-resource.openai.azure.com/openai/v1/audio/transcriptions?api-version=v1" + ); + assert_eq!(seen[0].headers["api-key"], "test-api-key"); + assert!(seen[0].headers["content-type"].starts_with("multipart/form-data; boundary=")); + let body = String::from_utf8(seen[0].body.clone()).unwrap(); + assert!(body.contains("name=\"model\"\r\n\r\ntest-deployment")); + assert!(body.contains("filename=\"audio.wav\"")); + assert!(body.contains("Content-Type: audio/wav")); + assert_eq!(result.text, "Hello, world!"); + assert!(result.segments.is_empty()); + assert_eq!(result.duration_in_seconds, Some(5.0)); +} + +/// TS: "should use correct URL format" (@ai-sdk/azure/src/azure-openai-provider.test.ts) +#[tokio::test] +async fn speech_generate() { + let fetch = MockFetch::new(vec![Canned { + status: 200, + headers: vec![("content-type".into(), "audio/mpeg".into())], + body: vec![1, 2, 3], + }]); + let result = provider(&fetch) + .speech("test-deployment") + .do_generate(&SpeechCallOptions::new("Hello, world!")) + .await + .unwrap(); + assert_eq!( + assert_request(&fetch, "audio/speech"), + json!({"model":"test-deployment","input":"Hello, world!","voice":"alloy","response_format":"mp3"}) + ); + assert!(matches!(result.audio, AudioData::Binary(audio) if audio == vec![1,2,3])); +} + +/// TS: "sends audio and the required model definition" (azure-transcription-model.test.ts) +#[tokio::test] +async fn speech_transcription_request() { + for audio in [ + AudioInput::Binary(vec![1, 2, 3]), + AudioInput::Base64("AQID".into()), + ] { + let fetch = MockFetch::new(vec![Canned::json(&json!({ + "combinedPhrases": [{"text": "Hello world."}], "phrases": [], + }))]); + provider(&fetch) + .transcription("mai-transcribe-2") + .do_generate(&TranscriptionCallOptions::new(audio, "audio/wav")) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!( + seen[0].url, + "https://test-resource.cognitiveservices.azure.com/speechtotext/transcriptions:transcribe?api-version=2025-10-15" + ); + assert_eq!(seen[0].headers["ocp-apim-subscription-key"], "test-api-key"); + assert!(!seen[0].headers.contains_key("api-key")); + assert!(seen[0].headers["user-agent"].contains("ai-sdk-azure/")); + let body = String::from_utf8(seen[0].body.clone()).unwrap(); + assert!(body.contains("name=\"audio\"; filename=\"audio.wav\"")); + assert!(body.contains("Content-Type: audio/wav")); + assert!(body.contains("\r\n\r\n\u{1}\u{2}\u{3}\r\n")); + let definition = body + .split("name=\"definition\"\r\n\r\n") + .nth(1) + .unwrap() + .split("\r\n") + .next() + .unwrap(); + assert_eq!( + serde_json::from_str::(definition).unwrap(), + json!({ + "enhancedMode": {"enabled": true, "model": "MAI-Transcribe-2", "modelOptions": {"timestamps": "segment"}}, + }) + ); + } +} diff --git a/aimux-providers/tests/bedrock_convert_test.rs b/aimux-providers/tests/bedrock_convert_test.rs index 270d119ca..eb04ab933 100644 --- a/aimux-providers/tests/bedrock_convert_test.rs +++ b/aimux-providers/tests/bedrock_convert_test.rs @@ -5,21 +5,6 @@ //! - `amazon-bedrock-prepare-tools.test.ts` (23 cases incl. `it.each`) //! - `convert-amazon-bedrock-usage.test.ts` (9 cases) //! -//! Cases the Rust data model cannot express are skipped with an inline -//! comment. The main categories of skips: -//! - **System-after-non-system throw**: the Rust converter currently lifts -//! all system messages into the system array. -//! - **S3 URLs / provider references**: `FileUrl` / `FileReference` are not -//! converted by the Rust Bedrock path. -//! - **Top-level-only mediaType auto-detection from bytes**: the Rust path -//! does not sniff magic bytes. -//! - **Mistral tool-call-id normalization (`isMistral`)**: the Rust -//! `convert_prompt_to_bedrock` has no `isMistral` parameter. The -//! non-Mistral (passthrough) cases ARE covered. -//! - **Provider-defined tools (web_search, anthropic provider tools)** and -//! `additionalTools`/`betas`: the Rust `FunctionTool` has no `type`/`id`. -//! - **`raw` echo on `Usage`**: the Rust `Usage` type has no `raw` field. - use serde_json::{Value, json}; use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions, provider_namespace}; @@ -155,10 +140,6 @@ fn system_combine_multiple_leading() { ); } -// SKIPPED (TS: "should throw an error if a system message is provided after a -// non-system message"): the converter currently lifts all system messages -// into the system array. - /// TS: "should extract the system message" #[test] fn system_extract_single() { @@ -188,9 +169,6 @@ fn user_convert_image_parts() { ); } -// SKIPPED (TS: "should convert image parts with S3 URLs"): FileUrl is not -// converted by the Rust Bedrock path. - /// TS: "should convert messages with document parts" #[test] fn user_convert_document_parts() { @@ -270,11 +248,6 @@ fn user_consistent_document_names() { ); } -// Message-level cache-point cases remain outside this restored test subset. - -// SKIPPED (TS: "should throw for file parts with provider references"): -// FileReference is not converted by the user-file path. - /// TS: "should add cache point to user content part when specified" #[test] fn user_content_part_cache_point() { @@ -369,8 +342,6 @@ fn assistant_combine_sequential() { ); } -// Message-level cache-point cases remain outside this restored test subset. - /// TS: "should add cache point to assistant content part when specified" #[test] fn assistant_content_part_cache_point() { @@ -742,9 +713,6 @@ fn tool_result_content_image() { ); } -// SKIPPED (TS: "should convert tool result images with S3 URLs"): FileUrl is -// not converted by the Rust Bedrock path. - /// TS: "should convert tool result with content array containing PDF" #[test] fn tool_result_content_pdf() { @@ -906,9 +874,6 @@ fn citations_multiple_pdfs() { // ── additional file format tests ──────────────────────────────────────────── -// SKIPPED (TS: "should throw an error for unsupported file mime type in user -// message content"): unknown mimes fall back to a default format (no throw). - /// TS: "should handle xlsx files correctly" #[test] fn file_format_xlsx() { @@ -944,10 +909,6 @@ fn file_format_docx() { // ── Mistral tool call ID normalization ────────────────────────────────────── -// SKIPPED (TS: "should normalize tool call IDs in tool results when isMistral -// is true" and "...in tool calls when isMistral is true"): the Rust -// convert_prompt_to_bedrock has no isMistral parameter. - /// TS: "should not normalize tool call IDs when isMistral is false" #[test] fn mistral_no_normalize_when_false() { @@ -1005,10 +966,6 @@ fn media_type_pass_through_full_image() { ); } -// SKIPPED (TS: "should detect subtype from inline bytes when mediaType is -// top-level-only (image)" and "...(application/pdf)"): the Rust path does not -// sniff magic bytes to resolve a top-level-only mediaType. - /// TS: "should route to document slot for non-image top-level type via getTopLevelMediaType" #[test] fn media_type_route_to_document_text_plain() { @@ -1025,11 +982,6 @@ fn media_type_route_to_document_text_plain() { ); } -// SKIPPED (TS: "should throw UnsupportedFunctionalityError for URL data (File -// URL)", "...for unsupported full image mediaType", and "...when top-level-only -// bytes cannot be detected"): FileUrl is not converted; unknown image mimes -// fall back to "png" instead of throwing; no magic-byte detection. - // ════════════════════════════════════════════════════════════════════════════ // amazon-bedrock-prepare-tools // ════════════════════════════════════════════════════════════════════════════ @@ -1114,10 +1066,8 @@ fn prepare_tools_tool_choice_filters_to_named() { // convert-amazon-bedrock-usage // ════════════════════════════════════════════════════════════════════════════ // -// The TS `convertAmazonBedrockUsage` echoes the input as `result.raw`. The Rust -// `Usage` type has no `raw` field, so only the `inputTokens`/`outputTokens` -// breakdown is asserted; the two `raw`-only cases ("should include totalTokens -// in raw when provided" and "should preserve raw usage data") are SKIPPED. +// The TS `convertAmazonBedrockUsage` echoes the input as `result.raw`. +// The assertions below cover the `inputTokens`/`outputTokens` breakdown. fn bedrock_usage( input: u32, @@ -1131,6 +1081,7 @@ fn bedrock_usage( total_tokens: None, cache_read_input_tokens: cache_read, cache_write_input_tokens: cache_write, + ..Default::default() } } @@ -1217,6 +1168,3 @@ fn usage_undefined() { assert_eq!(u.output_tokens.total, None); assert_eq!(u.output_tokens.text, None); } - -// SKIPPED (TS: "should include totalTokens in raw when provided" and "should -// preserve raw usage data"): the Rust `Usage` type has no `raw` echo field. diff --git a/aimux-providers/tests/bedrock_embedding_test.rs b/aimux-providers/tests/bedrock_embedding_test.rs index e9727dad0..b663074e5 100644 --- a/aimux-providers/tests/bedrock_embedding_test.rs +++ b/aimux-providers/tests/bedrock_embedding_test.rs @@ -73,7 +73,7 @@ fn should_expose_max_embeddings_per_call() { async fn should_handle_titan_single_input() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/model/amazon.titan-embed-text-v2:0/invoke")) + .and(path("/model/amazon.titan-embed-text-v2%3A0/invoke")) .respond_with( ResponseTemplate::new(200) .insert_header("content-type", "application/json") @@ -108,7 +108,7 @@ async fn should_handle_titan_single_input() { async fn should_extract_titan_usage() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/model/amazon.titan-embed-text-v2:0/invoke")) + .and(path("/model/amazon.titan-embed-text-v2%3A0/invoke")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "embedding": mock_embeddings()[0], "inputTextTokenCount": 8 @@ -169,7 +169,7 @@ async fn should_support_cohere_v3() { async fn should_support_cohere_v4() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/model/cohere.embed-v4:0/invoke")) + .and(path("/model/cohere.embed-v4%3A0/invoke")) .respond_with( ResponseTemplate::new(200) .insert_header("x-amzn-bedrock-input-token-count", "6") @@ -198,7 +198,7 @@ async fn should_support_cohere_v4() { async fn should_send_multiple_values_cohere_v4() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/model/cohere.embed-v4:0/invoke")) + .and(path("/model/cohere.embed-v4%3A0/invoke")) .respond_with( ResponseTemplate::new(200) .insert_header("x-amzn-bedrock-input-token-count", "12") @@ -239,7 +239,7 @@ async fn should_send_multiple_values_cohere_v4() { async fn should_support_cross_region_cohere() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/model/us.cohere.embed-v4:0/invoke")) + .and(path("/model/us.cohere.embed-v4%3A0/invoke")) .respond_with( ResponseTemplate::new(200) .insert_header("x-amzn-bedrock-input-token-count", "6") @@ -267,7 +267,7 @@ async fn should_support_nova_embeddings() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path( - "/model/amazon.nova-2-multimodal-embeddings-v1:0/invoke", + "/model/amazon.nova-2-multimodal-embeddings-v1%3A0/invoke", )) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "embeddings": [ diff --git a/aimux-providers/tests/bedrock_remaining_test.rs b/aimux-providers/tests/bedrock_remaining_test.rs new file mode 100644 index 000000000..5c9fa9287 --- /dev/null +++ b/aimux-providers/tests/bedrock_remaining_test.rs @@ -0,0 +1,1215 @@ +//! Remaining Amazon Bedrock upstream tests, restored from the tool-types branch. +//! +//! Pinned translation sources (`amazon-bedrock/src/`): +//! - `normalize-tool-call-id.test.ts` +//! - `convert-amazon-bedrock-usage.test.ts` (raw echo) +//! - `amazon-bedrock-provider.test.ts` (factory/auth variants) +//! - `convert-to-amazon-bedrock-chat-messages.test.ts` (Mistral IDs) +//! - `amazon-bedrock-chat-language-model.test.ts` (request/response cases) + +use aimux_core::tool::RawToolCall; + +use futures::StreamExt; +use serde_json::{Value, json}; +use serial_test::serial; + +use aimux_core::language_model_message::{ + AssistantPart, TextPart, ToolCallPart, ToolPart, ToolResultOutput, ToolResultPart, UserPart, +}; +use aimux_core::shared::provider_namespace; +use aimux_provider_utils::Resolvable; +#[path = "common/mock_fetch.rs"] +mod mock_fetch; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; +use aimux_core::options::CallOptions; +use aimux_core::result::GenerateContent; +use aimux_core::stream_part::StreamPart; +use aimux_core::types::FinishReasonUnified; +use mock_fetch::{Canned, MockFetch}; + +use aimux_providers::bedrock::convert::{build_request_body_checked, normalize_tool_call_id}; +use aimux_providers::bedrock::{ + AmazonBedrockProviderSettings, BedrockModel, create_amazon_bedrock, +}; + +fn test_tool() -> aimux_core::tool::FunctionTool { + aimux_core::tool::FunctionTool { + name: "getWeather".into(), + description: None, + input_schema: json!({"type":"object"}), + strict: None, + provider_options: None, + input_examples: None, + } +} +fn model_normalizes_tool_id(id: &str) -> bool { + let mut opts = default_options(vec![assistant_msg(vec![tool_call( + "tooluse_bpe71yCfRu2b5i-nKGDr5g".into(), + "getWeather".into(), + json!({}), + )])]); + opts.tools = Some(vec![test_tool().into()]); + let (body, _, _, _) = build_request_body_checked(id, &opts, None).unwrap(); + body["messages"][0]["content"][0]["toolUse"]["toolUseId"] == "8eHypBDcw" +} + +fn text(value: &str) -> UserPart { + UserPart::Text(TextPart { + text: value.into(), + provider_options: None, + }) +} +fn tool_call(id: String, name: String, input: Value) -> AssistantPart { + AssistantPart::ToolCall(ToolCallPart { + tool_call_id: id, + tool_name: name, + input, + provider_executed: None, + provider_options: None, + }) +} +fn tool_result(id: String, output: Value) -> ToolPart { + ToolPart::ToolResult(ToolResultPart { + tool_call_id: id, + tool_name: "weather".into(), + output: serde_json::from_value::(output).unwrap(), + provider_options: None, + }) +} + +// ── Shared helpers ─────────────────────────────────────────────────────────── + +fn test_prompt() -> LanguageModelPrompt { + vec![LanguageModelMessage::User { + content: vec![text("Hello")], + provider_options: None, + }] +} + +fn default_options(prompt: LanguageModelPrompt) -> CallOptions { + CallOptions::new(prompt) +} + +fn make_model(server: &std::sync::Arc) -> BedrockModel { + create_amazon_bedrock(AmazonBedrockProviderSettings { + fetch: Some(server.transport()), + base_url: Some("https://bedrock.test".into()), + api_key: Some(Resolvable::Value("test-token".into())), + ..Default::default() + }) + .unwrap() + .chat("anthropic.claude-3-5-sonnet-20240620-v1:0") +} + +fn mock_converse_json(status: u16, body: Value) -> std::sync::Arc { + let mut canned = Canned::json(&body); + canned.status = status; + MockFetch::new(vec![canned]) +} + +fn ok_converse_body() -> Value { + json!({ + "output": { + "message": { + "role": "assistant", + "content": [{ "text": "ok" }] + } + }, + "stopReason": "end_turn", + "usage": { "inputTokens": 4, "outputTokens": 7, "totalTokens": 11 } + }) +} + +fn as_text(item: &GenerateContent) -> &str { + match item { + GenerateContent::Text { text, .. } => text, + _ => panic!("expected Text content, got {item:?}"), + } +} + +fn as_tool_call(item: &GenerateContent) -> (&str, &str, &str) { + match item { + GenerateContent::ToolCall(RawToolCall { + tool_call_id, + tool_name, + input, + .. + }) => (tool_call_id, tool_name, input), + _ => panic!("expected ToolCall content, got {item:?}"), + } +} + +async fn collect_stream(result: aimux_core::result::StreamResult) -> Vec { + let mut parts = Vec::new(); + let mut stream = result.stream; + while let Some(part) = stream.next().await { + match part { + Ok(p) => parts.push(p), + Err(e) => panic!("stream error: {e:?}"), + } + } + parts +} + +/// Clear the authentication environment between serial tests. +fn clear_bedrock_env() { + unsafe { + std::env::remove_var("AWS_BEARER_TOKEN_BEDROCK"); + std::env::remove_var("AWS_ACCESS_KEY_ID"); + std::env::remove_var("AWS_SECRET_ACCESS_KEY"); + std::env::remove_var("AWS_REGION"); + std::env::remove_var("AWS_SESSION_TOKEN"); + } +} + +/// TS: "should return true for mistral models" +#[test] +fn is_mistral_model_true_for_mistral() { + for id in [ + "mistral.mistral-7b-instruct-v0:2", + "mistral.mixtral-8x7b-instruct-v0:1", + "mistral.mistral-large-2402-v1:0", + "mistral.mistral-small-2402-v1:0", + "mistral.mistral-large-2407-v1:0", + "mistral.ministral-3-14b-instruct", + "mistral.ministral-3-8b-instruct", + ] { + assert!(model_normalizes_tool_id(id)); + } +} + +/// TS: "should return true for region-prefixed mistral models" +#[test] +fn is_mistral_model_region_prefixed() { + assert!(model_normalizes_tool_id( + "us.mistral.pixtral-large-2502-v1:0" + )); + assert!(model_normalizes_tool_id( + "eu.mistral.mistral-large-2407-v1:0" + )); +} + +/// TS: "should return false for non-mistral models" +#[test] +fn is_mistral_model_false_for_non_mistral() { + for id in [ + "anthropic.claude-3-5-sonnet-20241022-v2:0", + "amazon.nova-pro-v1:0", + "openai.gpt-4o", + "meta.llama3-70b-instruct-v1:0", + ] { + assert!(!model_normalizes_tool_id(id)); + } +} + +/// TS: "should return the original ID when not a Mistral model" +#[test] +fn normalize_tool_call_id_passthrough_when_not_mistral() { + let id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; + assert_eq!(normalize_tool_call_id(id, false), id); +} + +/// TS: "should hash incompatible IDs deterministically for Mistral models" +#[test] +fn normalize_tool_call_id_hashes_deterministically() { + let id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; + assert_eq!(normalize_tool_call_id(id, true), "8eHypBDcw"); + assert_eq!(normalize_tool_call_id(id, true), "8eHypBDcw"); +} + +/// TS: "should produce 9 alphanumeric characters for incompatible IDs" +#[test] +fn normalize_tool_call_id_special_chars() { + for (id, expected) in [ + ("tool-use_123ABC456", "hvVDqPNyj"), + ("___abc123DEF___", "TnzPqldGU"), + ("abc", "GRuIyUwcV"), + ("12345", "PceAgWDYe"), + ("___---___", "5C589HVqG"), + ] { + let normalized = normalize_tool_call_id(id, true); + assert_eq!(normalized, expected); + assert_eq!(normalized.len(), 9); + assert!(normalized.bytes().all(|c| c.is_ascii_alphanumeric())); + } +} + +/// TS: "should preserve IDs that are already valid Mistral tool call IDs" +#[test] +fn normalize_tool_call_id_already_alphanumeric() { + assert_eq!(normalize_tool_call_id("abcdefghi", true), "abcdefghi"); + assert_eq!(normalize_tool_call_id("abc123XYZ", true), "abc123XYZ"); +} + +use aimux_providers::bedrock::convert::{BedrockUsage, convert_usage}; + +fn bedrock_usage_with_total( + input: u32, + output: u32, + total: Option, + cache_read: Option, + cache_write: Option, +) -> BedrockUsage { + BedrockUsage { + input_tokens: Some(input), + output_tokens: Some(output), + total_tokens: total, + cache_read_input_tokens: cache_read, + cache_write_input_tokens: cache_write, + ..Default::default() + } +} + +/// TS: "should include totalTokens in raw when provided" +#[test] +fn usage_raw_includes_total_tokens() { + let raw = bedrock_usage_with_total(100, 50, Some(150), None, None); + let u = convert_usage(Some(&raw)); + assert_eq!( + u.raw, + Some( + json!({"inputTokens":100,"outputTokens":50,"totalTokens":150}) + .as_object() + .unwrap() + .clone() + ) + ); +} + +/// TS: "should preserve raw usage data" +#[test] +fn usage_raw_preserved() { + let raw = bedrock_usage_with_total(100, 50, Some(150), Some(80), Some(60)); + let u = convert_usage(Some(&raw)); + assert_eq!(u.raw, Some(json!({"inputTokens":100,"outputTokens":50,"totalTokens":150,"cacheReadInputTokens":80,"cacheWriteInputTokens":60}).as_object().unwrap().clone())); +} + +// ════════════════════════════════════════════════════════════════════════════ +// amazon-bedrock-provider.test.ts — provider configuration / auth +// ════════════════════════════════════════════════════════════════════════════ + +/// TS: "should create a provider instance with default options" — the mocked +/// upstream region loader returns `us-east-1`; pass it explicitly here. +#[tokio::test] +#[serial] +async fn provider_default_region_and_base_url() { + let fetch = MockFetch::new(vec![Canned::json(&ok_converse_body())]); + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings { + api_key: Some(Resolvable::Value(String::new())), + region: Some("us-east-1".into()), + access_key_id: Some("ak".into()), + secret_access_key: Some("sk".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + provider + .chat("amazon.nova-pro-v1:0") + .do_generate(&default_options(test_prompt())) + .await + .unwrap(); + let seen = &fetch.seen()[0]; + assert_eq!( + seen.url, + "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.nova-pro-v1%3A0/converse" + ); + assert!(seen.headers["authorization"].starts_with("AWS4-HMAC-SHA256")); +} + +/// TS: "should create a provider instance with custom options" — custom region +/// flows into the base URL. +#[tokio::test] +async fn provider_custom_region_base_url() { + let fetch = MockFetch::new(vec![Canned::json(&ok_converse_body())]); + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings { + api_key: Some(Resolvable::Value(String::new())), + region: Some("eu-west-1".into()), + access_key_id: Some("ak".into()), + secret_access_key: Some("sk".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + provider + .chat("amazon.nova-pro-v1:0") + .do_generate(&default_options(test_prompt())) + .await + .unwrap(); + let seen = &fetch.seen()[0]; + assert!( + seen.url + .starts_with("https://bedrock-runtime.eu-west-1.amazonaws.com/") + ); +} + +/// TS: "should create a provider instance with custom options" — baseURL override. +#[tokio::test] +async fn provider_with_base_url_override() { + let fetch = MockFetch::new(vec![Canned::json(&ok_converse_body())]); + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings { + region: Some("us-east-1".into()), + api_key: Some(Resolvable::Value("tok".into())), + base_url: Some("https://custom.url/".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + provider + .chat("amazon.nova-pro-v1:0") + .do_generate(&default_options(test_prompt())) + .await + .unwrap(); + let seen = &fetch.seen()[0]; + assert!(seen.url.starts_with("https://custom.url/model/")); +} + +/// TS: "should use API key when provided in options" — bearer-token auth path. +#[tokio::test] +async fn provider_bearer_token_auth() { + let fetch = MockFetch::new(vec![Canned::json(&ok_converse_body())]); + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings { + region: Some("us-east-1".into()), + api_key: Some(Resolvable::Value("test-api-key".into())), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + provider + .chat("amazon.nova-pro-v1:0") + .do_generate(&default_options(test_prompt())) + .await + .unwrap(); + let seen = &fetch.seen()[0]; + assert_eq!(seen.headers["authorization"], "Bearer test-api-key"); +} + +/// TS: "should use API key from environment variable" — `AWS_BEARER_TOKEN_BEDROCK` +/// takes precedence over SigV4 environment credentials. +#[tokio::test] +#[serial] +async fn provider_from_env_bearer_token_precedence() { + clear_bedrock_env(); + unsafe { + std::env::set_var("AWS_BEARER_TOKEN_BEDROCK", "env-bearer"); + std::env::set_var("AWS_ACCESS_KEY_ID", "should-not-be-used"); + std::env::set_var("AWS_SECRET_ACCESS_KEY", "should-not-be-used"); + std::env::set_var("AWS_REGION", "us-west-2"); + } + + let fetch = MockFetch::new(vec![Canned::json(&ok_converse_body())]); + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings { + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + provider + .chat("amazon.nova-pro-v1:0") + .do_generate(&default_options(test_prompt())) + .await + .unwrap(); + let seen = &fetch.seen()[0]; + assert_eq!(seen.headers["authorization"], "Bearer env-bearer"); + assert!(seen.url.contains("us-west-2")); + clear_bedrock_env(); +} + +/// TS: "should fall back to SigV4 when no API key provided" — the factory uses +/// SigV4 when `AWS_BEARER_TOKEN_BEDROCK` is absent. +#[tokio::test] +#[serial] +async fn provider_from_env_sigv4_fallback() { + clear_bedrock_env(); + unsafe { + std::env::set_var("AWS_ACCESS_KEY_ID", "test-ak"); + std::env::set_var("AWS_SECRET_ACCESS_KEY", "test-sk"); + std::env::set_var("AWS_REGION", "eu-west-1"); + } + + let fetch = MockFetch::new(vec![Canned::json(&ok_converse_body())]); + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings { + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + provider + .chat("amazon.nova-pro-v1:0") + .do_generate(&default_options(test_prompt())) + .await + .unwrap(); + let seen = &fetch.seen()[0]; + assert!(seen.headers["authorization"].starts_with("AWS4-HMAC-SHA256 Credential=test-ak/")); + assert!(seen.headers["authorization"].contains("/eu-west-1/bedrock/")); + assert!(!seen.headers.contains_key("x-amz-security-token")); + clear_bedrock_env(); +} + +/// TS: "should maintain backward compatibility with existing SigV4 +/// authentication" — `AWS_SESSION_TOKEN` is loaded from the environment. +#[tokio::test] +#[serial] +async fn provider_from_env_session_token() { + clear_bedrock_env(); + unsafe { + std::env::set_var("AWS_ACCESS_KEY_ID", "ak"); + std::env::set_var("AWS_SECRET_ACCESS_KEY", "sk"); + std::env::set_var("AWS_REGION", "us-east-1"); + std::env::set_var("AWS_SESSION_TOKEN", "sts-token"); + } + + let fetch = MockFetch::new(vec![Canned::json(&ok_converse_body())]); + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings { + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + provider + .chat("amazon.nova-pro-v1:0") + .do_generate(&default_options(test_prompt())) + .await + .unwrap(); + let seen = &fetch.seen()[0]; + assert_eq!(seen.headers["x-amz-security-token"], "sts-token"); + clear_bedrock_env(); +} + +/// The `BedrockProvider` exposes its name and can vend language models. +#[test] +fn provider_name_and_language_model() { + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings::default()).unwrap(); + let model = provider.chat("anthropic.claude-3-5-sonnet-20240620-v1:0"); + assert_eq!( + model.model_id(), + "anthropic.claude-3-5-sonnet-20240620-v1:0" + ); + assert_eq!(model.provider(), "amazon-bedrock"); +} + +/// TS: "should prioritize options.apiKey over environment variable" — explicit +/// bearer token construction does not consult env vars. +#[tokio::test] +#[serial] +async fn provider_explicit_bearer_over_env() { + clear_bedrock_env(); + unsafe { + std::env::set_var("AWS_BEARER_TOKEN_BEDROCK", "env-bearer"); + } + // Explicit construction with a different token should win over env. + let fetch = MockFetch::new(vec![Canned::json(&ok_converse_body())]); + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings { + region: Some("us-east-1".into()), + api_key: Some(Resolvable::Value("explicit-tok".into())), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + provider + .chat("amazon.nova-pro-v1:0") + .do_generate(&default_options(test_prompt())) + .await + .unwrap(); + let seen = &fetch.seen()[0]; + assert_eq!(seen.headers["authorization"], "Bearer explicit-tok"); + clear_bedrock_env(); +} + +fn tool_msg(content: Vec) -> LanguageModelMessage { + LanguageModelMessage::Tool { + content, + provider_options: None, + } +} + +fn assistant_msg(content: Vec) -> LanguageModelMessage { + LanguageModelMessage::Assistant { + content, + provider_options: None, + } +} + +/// TS: "should normalize tool call IDs in tool results when isMistral is true" +#[test] +fn convert_mistral_normalize_tool_result_id() { + let original_id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; + let prompt = vec![tool_msg(vec![tool_result( + original_id.to_string(), + json!({ "type": "text", "value": "The result is 42" }), + )])]; + let mut opts = default_options(prompt); + opts.tools = Some(vec![test_tool().into()]); + let (body, _, _, _) = + build_request_body_checked("mistral.mistral-large-2402-v1:0", &opts, None).unwrap(); + let messages = &body["messages"]; + assert_eq!( + messages[0]["content"][0]["toolResult"]["toolUseId"], + json!("8eHypBDcw") + ); +} + +/// TS: "should normalize tool call IDs in tool calls when isMistral is true" +#[test] +fn convert_mistral_normalize_tool_call_id() { + let original_id = "tooluse_xyz123ABC456-def"; + let prompt = vec![assistant_msg(vec![tool_call( + original_id.to_string(), + "test-tool".to_string(), + json!({ "query": "test" }), + )])]; + let mut opts = default_options(prompt); + opts.tools = Some(vec![test_tool().into()]); + let (body, _, _, _) = + build_request_body_checked("mistral.mistral-large-2402-v1:0", &opts, None).unwrap(); + let messages = &body["messages"]; + assert_eq!( + messages[0]["content"][0]["toolUse"]["toolUseId"], + json!("De61BW1Dz") + ); +} + +// ════════════════════════════════════════════════════════════════════════════ +// amazon-bedrock-chat-language-model.test.ts — request URL / body / response +// ════════════════════════════════════════════════════════════════════════════ + +// ── supportedUrls ──────────────────────────────────────────────────────────── + +/// TS: "should support S3 URLs for image and video parts". +#[test] +fn model_supported_urls_s3() { + let model = create_amazon_bedrock(AmazonBedrockProviderSettings::default()) + .unwrap() + .chat("amazon.nova-pro-v1:0"); + let urls = model.supported_urls(); + assert_eq!(urls.0.len(), 2); + for kind in ["image/*", "video/*"] { + assert_eq!(urls.0[kind].len(), 1); + assert_eq!(urls.0[kind][0].as_str(), "^s3://"); + } +} + +// ── ARN model IDs containing a slash ───────────────────────────────────────── + +/// TS: "should generate text through the encoded Converse route". +/// +/// ARN inference-profile IDs contain a `/` which must be percent-encoded in the +/// URL path. The upstream uses `encodeURIComponent(modelId)`. +#[tokio::test] +async fn arn_model_id_encoded_generate_route() { + let arn = "arn:aws:bedrock:eu-west-1:474668406012:inference-profile/eu.amazon.nova-lite-v1:0"; + let encoded = "arn%3Aaws%3Abedrock%3Aeu-west-1%3A474668406012%3Ainference-profile%2Feu.amazon.nova-lite-v1%3A0"; + + let server = mock_converse_json( + 200, + json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "Hello!" }] } }, + "stopReason": "end_turn", + "usage": { "inputTokens": 1, "outputTokens": 1 } + }), + ); + + let model = create_amazon_bedrock(AmazonBedrockProviderSettings { + fetch: Some(server.transport()), + base_url: Some("https://bedrock.test".into()), + api_key: Some(Resolvable::Value("tok".into())), + ..Default::default() + }) + .unwrap() + .chat(arn); + + let result = model + .do_generate(&default_options(test_prompt())) + .await + .expect("should succeed via the encoded route"); + assert_eq!(as_text(&result.content[0]), "Hello!"); + assert_eq!( + server.seen()[0].url, + format!("https://bedrock.test/model/{encoded}/converse") + ); +} + +/// TS: "should stream text through the encoded Converse route". +#[tokio::test] +async fn arn_model_id_encoded_stream_route() { + let arn = "arn:aws:bedrock:eu-west-1:474668406012:inference-profile/eu.amazon.nova-lite-v1:0"; + let encoded = "arn%3Aaws%3Abedrock%3Aeu-west-1%3A474668406012%3Ainference-profile%2Feu.amazon.nova-lite-v1%3A0"; + + let events: Vec<(&str, &str, &str)> = vec![ + ("event", "messageStart", r#"{"role":"assistant"}"#), + ( + "event", + "contentBlockDelta", + r#"{"contentBlockIndex":0,"delta":{"text":"Hello!"}}"#, + ), + ("event", "contentBlockStop", r#"{"contentBlockIndex":0}"#), + ("event", "messageStop", r#"{"stopReason":"end_turn"}"#), + ( + "event", + "metadata", + r#"{"usage":{"inputTokens":1,"outputTokens":1}}"#, + ), + ]; + let body_bytes = aimux_providers::bedrock::event_stream::encode_messages(&events); + + let server = MockFetch::new(vec![Canned { + status: 200, + headers: vec![( + "content-type".into(), + "application/vnd.amazon.eventstream".into(), + )], + body: body_bytes, + }]); + + let model = create_amazon_bedrock(AmazonBedrockProviderSettings { + fetch: Some(server.transport()), + base_url: Some("https://bedrock.test".into()), + api_key: Some(Resolvable::Value("tok".into())), + ..Default::default() + }) + .unwrap() + .chat(arn); + + let result = model + .do_stream(&default_options(test_prompt())) + .await + .expect("should succeed via the encoded stream route"); + let parts = collect_stream(result).await; + + let text: String = parts + .iter() + .filter_map(|p| match p { + StreamPart::TextDelta { delta, .. } => Some(delta.clone()), + _ => None, + }) + .collect(); + assert_eq!(text, "Hello!"); + assert_eq!( + server.seen()[0].url, + format!("https://bedrock.test/model/{encoded}/converse-stream") + ); +} + +// ── temperature clamping ───────────────────────────────────────────────────── + +/// TS: "should clamp temperature above 1 to 1 and add warning". +#[tokio::test] +async fn temperature_clamped_above_1() { + let server = mock_converse_json(200, ok_converse_body()); + + let model = make_model(&server); + let mut opts = default_options(test_prompt()); + opts.temperature = Some(1.5); + + let result = model.do_generate(&opts).await.expect("should succeed"); + + let body = result.request.unwrap().body.expect("request body"); + assert_eq!( + body["inferenceConfig"]["temperature"].as_f64().unwrap(), + 1.0 + ); + assert!( + result.warnings.iter().any( + |w| matches!(w, aimux_core::types::Warning::Unsupported { feature, .. } + if feature == "temperature") + ), + "expected an unsupported-temperature warning, got {:?}", + result.warnings + ); +} + +/// TS: "should clamp temperature below 0 to 0 and add warning". +#[tokio::test] +async fn temperature_clamped_below_0() { + let server = mock_converse_json(200, ok_converse_body()); + + let model = make_model(&server); + let mut opts = default_options(test_prompt()); + opts.temperature = Some(-0.5); + + let result = model.do_generate(&opts).await.expect("should succeed"); + + let body = result.request.unwrap().body.expect("request body"); + assert_eq!( + body["inferenceConfig"]["temperature"].as_f64().unwrap(), + 0.0 + ); + assert!( + result.warnings.iter().any( + |w| matches!(w, aimux_core::types::Warning::Unsupported { feature, .. } + if feature == "temperature") + ), + "expected an unsupported-temperature warning" + ); +} + +/// TS: "should not clamp valid temperature between 0 and 1". +#[tokio::test] +async fn temperature_not_clamped_in_range() { + let server = mock_converse_json(200, ok_converse_body()); + + let model = make_model(&server); + let mut opts = default_options(test_prompt()); + opts.temperature = Some(0.7); + + let result = model.do_generate(&opts).await.expect("should succeed"); + + let body = result.request.unwrap().body.expect("request body"); + let temp = body["inferenceConfig"]["temperature"].as_f64().unwrap(); + assert!( + (temp - 0.7).abs() < 1e-6, + "temperature should be 0.7, got {temp}" + ); + assert!(result.warnings.is_empty(), "no warnings expected"); +} + +// ── guardrails ─────────────────────────────────────────────────────────────── + +/// TS: "should support guardrails" — `providerOptions.bedrock.guardrailConfig` +/// is forwarded as a top-level `guardrailConfig` in the request body. +#[tokio::test] +async fn guardrail_config_in_request_body() { + let server = mock_converse_json(200, ok_converse_body()); + + let model = make_model(&server); + let mut opts = default_options(test_prompt()); + let po = provider_namespace( + "bedrock", + json!({ + "guardrailConfig": { + "guardrailIdentifier": "-1", + "guardrailVersion": "1", + "trace": "enabled" + } + }), + ) + .unwrap(); + opts.provider_options = Some(po); + + let result = model.do_generate(&opts).await.expect("should succeed"); + let body = result.request.unwrap().body.expect("request body"); + assert_eq!( + body["guardrailConfig"], + json!({ + "guardrailIdentifier": "-1", + "guardrailVersion": "1", + "trace": "enabled" + }) + ); +} + +// ── trace in providerMetadata ──────────────────────────────────────────────── + +/// TS: "should include trace information in providerMetadata" — the response +/// `trace` field is surfaced as `providerMetadata.bedrock.trace`. +#[tokio::test] +async fn trace_in_provider_metadata() { + let trace = json!({ + "guardrail": { + "inputAssessment": { + "1abcd2ef34gh": { + "contentPolicy": { + "filters": [{ + "action": "BLOCKED", + "confidence": "LOW", + "type": "INSULTS" + }] + } + } + } + } + }); + let server = mock_converse_json( + 200, + json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "Hello, World!" }] } }, + "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, + "stopReason": "stop_sequence", + "trace": trace + }), + ); + + let model = make_model(&server); + let result = model + .do_generate(&default_options(test_prompt())) + .await + .expect("should succeed"); + + let pm = result + .provider_metadata + .as_ref() + .expect("provider_metadata should be Some"); + assert_eq!(pm["bedrock"]["trace"], trace); +} + +// ── stop_sequence in providerMetadata ──────────────────────────────────────── + +/// TS: "should include stop_sequence in provider metadata" — when the response +/// carries `additionalModelResponseFields.delta.stop_sequence`, it is surfaced +/// as `providerMetadata.bedrock.stopSequence` (and `amazonBedrock.stopSequence`). +#[tokio::test] +async fn stop_sequence_in_provider_metadata() { + let server = mock_converse_json( + 200, + json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "Hello, World!" }] } }, + "stopReason": "stop_sequence", + "additionalModelResponseFields": { "delta": { "stop_sequence": "STOP" } }, + "usage": { "inputTokens": 4, "outputTokens": 30, "totalTokens": 34 } + }), + ); + + let model = make_model(&server); + let mut opts = default_options(test_prompt()); + opts.stop_sequences = Some(vec!["STOP".to_string()]); + + let result = model.do_generate(&opts).await.expect("should succeed"); + + let pm = result + .provider_metadata + .as_ref() + .expect("provider_metadata should be Some"); + assert_eq!(pm["bedrock"]["stopSequence"], json!("STOP")); + assert_eq!(pm["amazonBedrock"]["stopSequence"], json!("STOP")); +} + +// ── tool calls with empty input ────────────────────────────────────────────── + +/// TS: "should support tool calls with empty input (no arguments)" (stream). +/// +/// When a `toolUse` block carries no input deltas, the stream event loop +/// accumulates an empty string and yields `input: "{}"`. +#[tokio::test] +async fn stream_tool_call_empty_input() { + let events: Vec<(&str, &str, &str)> = vec![ + ("event", "messageStart", r#"{"role":"assistant"}"#), + ( + "event", + "contentBlockStart", + r#"{"contentBlockIndex":0,"start":{"toolUse":{"name":"updateIssueList","toolUseId":"tool_1"}}}"#, + ), + // No input deltas — the tool call has no arguments. + ("event", "contentBlockStop", r#"{"contentBlockIndex":0}"#), + ("event", "messageStop", r#"{"stopReason":"tool_use"}"#), + ]; + let body_bytes = aimux_providers::bedrock::event_stream::encode_messages(&events); + + let server = MockFetch::new(vec![Canned { + status: 200, + headers: vec![( + "content-type".into(), + "application/vnd.amazon.eventstream".into(), + )], + body: body_bytes, + }]); + + let model = make_model(&server); + let result = model + .do_stream(&default_options(test_prompt())) + .await + .expect("do_stream should succeed"); + let parts = collect_stream(result).await; + + let tool_calls: Vec<_> = parts + .iter() + .filter_map(|p| match p { + StreamPart::ToolCall(RawToolCall { + tool_call_id, + tool_name, + input, + .. + }) => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), + _ => None, + }) + .collect(); + assert_eq!(tool_calls.len(), 1); + assert_eq!(tool_calls[0].0, "tool_1"); + assert_eq!(tool_calls[0].1, "updateIssueList"); + assert_eq!(tool_calls[0].2, "{}"); +} + +// ── omit toolConfig ────────────────────────────────────────────────────────── + +/// TS: "should omit toolConfig and filter tool content when conversation has +/// tool calls but no active tools". +/// +#[tokio::test] +async fn omit_tool_config_when_no_active_tools() { + let server = mock_converse_json(200, ok_converse_body()); + + let prompt: LanguageModelPrompt = vec![ + LanguageModelMessage::User { + content: vec![text("What is the weather in Toronto?")], + provider_options: None, + }, + LanguageModelMessage::Assistant { + content: vec![tool_call( + "tool-call-1".to_string(), + "weather".to_string(), + json!({ "city": "Toronto" }), + )], + provider_options: None, + }, + LanguageModelMessage::Tool { + content: vec![tool_result( + "tool-call-1".to_string(), + json!({ "type": "text", "value": "The weather in Toronto is 20°C." }), + )], + provider_options: None, + }, + LanguageModelMessage::User { + content: vec![text("Now give me a summary.")], + provider_options: None, + }, + ]; + + let model = make_model(&server); + let opts = CallOptions { + tools: Some(vec![]), + ..default_options(prompt) + }; + + let result = model.do_generate(&opts).await.expect("should succeed"); + let body = result.request.unwrap().body.expect("request body"); + assert!( + body.get("toolConfig").is_none(), + "toolConfig should be absent when tools list is empty" + ); +} + +/// TS: (same test) the assistant's `toolUse` block should be filtered out of +/// the messages when no tools are active. +#[tokio::test] +async fn filter_tool_content_when_no_active_tools() { + let server = mock_converse_json(200, ok_converse_body()); + + let prompt: LanguageModelPrompt = vec![ + LanguageModelMessage::User { + content: vec![text("What is the weather in Toronto?")], + provider_options: None, + }, + LanguageModelMessage::Assistant { + content: vec![tool_call( + "tool-call-1".to_string(), + "weather".to_string(), + json!({ "city": "Toronto" }), + )], + provider_options: None, + }, + LanguageModelMessage::Tool { + content: vec![tool_result( + "tool-call-1".to_string(), + json!({ "type": "text", "value": "The weather in Toronto is 20°C." }), + )], + provider_options: None, + }, + LanguageModelMessage::User { + content: vec![text("Now give me a summary.")], + provider_options: None, + }, + ]; + + let model = make_model(&server); + let opts = CallOptions { + tools: Some(vec![]), + ..default_options(prompt) + }; + + let result = model.do_generate(&opts).await.expect("should succeed"); + let body = result.request.unwrap().body.expect("request body"); + assert_eq!( + body["messages"], + json!([{ "role": "user", "content": [{ "text": "What is the weather in Toronto?" }, { "text": "Now give me a summary." }] }]) + ); +} + +// ── doGenerate: tool call with empty input (non-streaming) ─────────────────── + +/// TS: "should support tool calls with empty input (no arguments)" (generate). +#[tokio::test] +async fn generate_tool_call_empty_input() { + let server = mock_converse_json( + 200, + json!({ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tool_1", + "name": "updateIssueList" + // no "input" field + } + } + ] + } + }, + "stopReason": "tool_use", + "usage": { "inputTokens": 5, "outputTokens": 5 } + }), + ); + + let model = make_model(&server); + let result = model + .do_generate(&default_options(test_prompt())) + .await + .expect("should succeed"); + + assert_eq!(result.content.len(), 1); + let (id, name, input) = as_tool_call(&result.content[0]); + assert_eq!(id, "tool_1"); + assert_eq!(name, "updateIssueList"); + assert_eq!(input, "{}"); +} + +// ── doGenerate: basic text + finish reason (sanity, already covered but +// exercised here against the remaining-test helper set) ──────────────────── + +/// TS: "should pass the model and the messages" — the request body carries the +/// model id implicitly via the URL path and the messages array. +#[tokio::test] +async fn request_body_messages_shape() { + let server = mock_converse_json(200, ok_converse_body()); + + let prompt: LanguageModelPrompt = vec![ + LanguageModelMessage::System { + content: "System Prompt".into(), + provider_options: None, + }, + LanguageModelMessage::User { + content: vec![text("Hello")], + provider_options: None, + }, + ]; + + let model = make_model(&server); + let result = model + .do_generate(&default_options(prompt)) + .await + .expect("ok"); + let body = result.request.unwrap().body.expect("request body"); + assert_eq!(body["system"], json!([{ "text": "System Prompt" }])); + assert_eq!(body["messages"][0]["role"], "user"); + assert_eq!(body["messages"][0]["content"][0]["text"], "Hello"); +} + +/// TS: "should extract finish reason" — `guardrail_intervened` → ContentFilter. +#[tokio::test] +async fn finish_reason_guardrail_intervened() { + let server = mock_converse_json( + 200, + json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "" }] } }, + "stopReason": "guardrail_intervened", + "usage": { "inputTokens": 4, "outputTokens": 1 } + }), + ); + + let model = make_model(&server); + let result = model + .do_generate(&default_options(test_prompt())) + .await + .expect("should succeed"); + + assert_eq!( + result.finish_reason.unified, + FinishReasonUnified::ContentFilter + ); + assert_eq!( + result.finish_reason.raw.as_deref(), + Some("guardrail_intervened") + ); +} + +/// TS: "should handle throttlingException error" (stream). +#[tokio::test] +async fn stream_throttling_error() { + let body_bytes = aimux_providers::bedrock::event_stream::encode_messages(&[( + "exception", + "throttlingException", + r#"{"message":"Throttling Error"}"#, + )]); + let server = MockFetch::new(vec![Canned { + status: 200, + headers: vec![( + "content-type".into(), + "application/vnd.amazon.eventstream".into(), + )], + body: body_bytes, + }]); + let model = make_model(&server); + let result = model + .do_stream(&default_options(test_prompt())) + .await + .unwrap(); + let parts = collect_stream(result).await; + assert!(parts.iter().any( + |part| matches!(part, StreamPart::Error { error } if error.status_code() == Some(429)) + )); + assert!(parts.iter().any(|part| matches!(part, StreamPart::Finish { finish_reason, .. } if finish_reason.unified == FinishReasonUnified::Error))); +} + +/// TS: "should handle validationException error" (stream). +#[tokio::test] +async fn stream_validation_error() { + let body_bytes = aimux_providers::bedrock::event_stream::encode_messages(&[( + "exception", + "validationException", + r#"{"message":"Validation Error"}"#, + )]); + let server = MockFetch::new(vec![Canned { + status: 200, + headers: vec![( + "content-type".into(), + "application/vnd.amazon.eventstream".into(), + )], + body: body_bytes, + }]); + let model = make_model(&server); + let result = model + .do_stream(&default_options(test_prompt())) + .await + .unwrap(); + let parts = collect_stream(result).await; + assert!(parts.iter().any( + |part| matches!(part, StreamPart::Error { error } if error.status_code() == Some(400)) + )); + assert!(parts.iter().any(|part| matches!(part, StreamPart::Finish { finish_reason, .. } if finish_reason.unified == FinishReasonUnified::Error))); +} + +// ── doStream: error handling ───────────────────────────────────────────────── + +// ── doStream: request body parity ──────────────────────────────────────────── + +/// TS: "should return the request body" (stream) — the stream result carries +/// the request body for debugging. +#[tokio::test] +async fn stream_request_body_available() { + let events: Vec<(&str, &str, &str)> = vec![ + ("event", "messageStart", r#"{"role":"assistant"}"#), + ( + "event", + "contentBlockDelta", + r#"{"contentBlockIndex":0,"delta":{"text":"hi"}}"#, + ), + ("event", "contentBlockStop", r#"{"contentBlockIndex":0}"#), + ("event", "messageStop", r#"{"stopReason":"end_turn"}"#), + ]; + let body_bytes = aimux_providers::bedrock::event_stream::encode_messages(&events); + + let server = MockFetch::new(vec![Canned { + status: 200, + headers: vec![( + "content-type".into(), + "application/vnd.amazon.eventstream".into(), + )], + body: body_bytes, + }]); + + let model = make_model(&server); + let mut opts = default_options(test_prompt()); + opts.max_output_tokens = Some(256); + + let result = model.do_stream(&opts).await.expect("should succeed"); + let body = result.request.unwrap().body.expect("stream request body"); + assert_eq!(body["messages"][0]["role"], "user"); + assert_eq!(body["inferenceConfig"]["maxTokens"], 256); +} diff --git a/aimux-providers/tests/bedrock_upstream_e2e_test.rs b/aimux-providers/tests/bedrock_upstream_e2e_test.rs new file mode 100644 index 000000000..c778c2af7 --- /dev/null +++ b/aimux-providers/tests/bedrock_upstream_e2e_test.rs @@ -0,0 +1,469 @@ +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use aimux_core::AiMuxError; +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_core::reranking_model::{RerankingCallOptions, RerankingDocuments, RerankingModel}; +use aimux_core::result::GenerateContent; +use aimux_core::shared::Size; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::FunctionTool; +use aimux_core::types::FinishReasonUnified; +use aimux_provider_utils::Resolvable; +use aimux_providers::bedrock::event_stream::encode_message; +use aimux_providers::{AmazonBedrockProviderSettings, create_amazon_bedrock}; +use futures::StreamExt; +use mock_fetch::{Canned, MockFetch}; +use serde_json::{Value, json}; + +fn provider(fetch: &std::sync::Arc) -> aimux_providers::bedrock::AmazonBedrockProvider { + create_amazon_bedrock(AmazonBedrockProviderSettings { + api_key: Some(Resolvable::Value("test-token".into())), + region: Some("us-west-2".into()), + base_url: Some("https://bedrock.test".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} + +fn options() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +fn chat_body() -> Value { + json!({"messages":[{"role":"user","content":[{"text":"Hello"}]}], + "additionalModelResponseFieldPaths":["/delta/stop_sequence"]}) +} + +fn stream_response(events: Value) -> Canned { + let mut body = Vec::new(); + for event in events.as_array().unwrap() { + let (kind, payload) = event.as_object().unwrap().iter().next().unwrap(); + body.extend(encode_message("event", kind, &payload.to_string())); + } + Canned { + status: 200, + headers: vec![( + "content-type".into(), + "application/vnd.amazon.eventstream".into(), + )], + body, + } +} + +/// TS: "should extract reasoning text with signature" (amazon-bedrock/src/amazon-bedrock-chat-language-model.test.ts) +#[tokio::test] +async fn generate_reasoning_signature_and_usage() { + let response = json!({"output":{"message":{"role":"assistant","content":[ + {"reasoningContent":{"reasoningText":{"text":"I need to think about this problem carefully...","signature":"abc123signature"}}}, + {"text":"The answer is 42."}]}},"usage":{"inputTokens":4,"outputTokens":34,"totalTokens":38},"stopReason":"stop_sequence"}); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let result = provider(&fetch) + .chat("anthropic.model") + .do_generate(&options()) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "POST"); + assert_eq!( + seen[0].url, + "https://bedrock.test/model/anthropic.model/converse" + ); + assert_eq!(seen[0].headers["authorization"], "Bearer test-token"); + assert_eq!(seen[0].json_body(), chat_body()); + assert_eq!(result.request.unwrap().body.unwrap(), chat_body()); + assert_eq!(result.content.len(), 2); + let GenerateContent::Reasoning(reasoning) = &result.content[0] else { + panic!("expected reasoning") + }; + assert_eq!( + reasoning.text, + "I need to think about this problem carefully..." + ); + assert_eq!( + serde_json::to_value(&reasoning.provider_metadata).unwrap(), + json!({"amazonBedrock":{"signature":"abc123signature"},"bedrock":{"signature":"abc123signature"}}) + ); + assert!( + matches!(&result.content[1], GenerateContent::Text { text, .. } if text == "The answer is 42.") + ); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(result.finish_reason.raw.as_deref(), Some("stop_sequence")); + assert_eq!(result.usage.input_tokens.total, Some(4)); + assert_eq!(result.usage.input_tokens.no_cache, Some(4)); + assert_eq!(result.usage.output_tokens.total, Some(34)); + assert_eq!(result.usage.output_tokens.text, Some(34)); + assert_eq!(result.usage.raw, response["usage"].as_object().cloned()); +} + +/// TS: "should stream text deltas with metadata and usage" (amazon-bedrock/src/amazon-bedrock-chat-language-model.test.ts) +#[tokio::test] +async fn stream_text_and_usage() { + let fetch = MockFetch::new(vec![stream_response(json!([ + {"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"Hello"}}}, + {"contentBlockDelta":{"contentBlockIndex":1,"delta":{"text":", "}}}, + {"contentBlockDelta":{"contentBlockIndex":2,"delta":{"text":"World!"}}}, + {"metadata":{"usage":{"inputTokens":4,"outputTokens":34,"totalTokens":38},"metrics":{"latencyMs":10}}}, + {"messageStop":{"stopReason":"stop_sequence"}} + ]))]); + let result = provider(&fetch) + .chat("anthropic.model") + .do_stream(&options()) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].url, + "https://bedrock.test/model/anthropic.model/converse-stream" + ); + assert_eq!(fetch.seen()[0].json_body(), chat_body()); + assert_eq!(result.request.unwrap().body.unwrap(), chat_body()); + let parts = result + .stream + .collect::>() + .await + .into_iter() + .collect::, _>>() + .unwrap(); + assert_eq!(parts.len(), 9); + assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); + assert!( + matches!(&parts[1], StreamPart::ResponseMetadata(meta) if meta.model_id.as_deref() == Some("anthropic.model")) + ); + for (offset, id, expected) in [(2, "0", "Hello"), (4, "1", ", "), (6, "2", "World!")] { + assert!(matches!(&parts[offset], StreamPart::TextStart { id: actual, .. } if actual == id)); + assert!( + matches!(&parts[offset + 1], StreamPart::TextDelta { id: actual, delta, .. } if actual == id && delta == expected) + ); + } + let StreamPart::Finish { + finish_reason, + usage, + .. + } = &parts[8] + else { + panic!("expected finish") + }; + assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(finish_reason.raw.as_deref(), Some("stop_sequence")); + assert_eq!(usage.input_tokens.total, Some(4)); + assert_eq!(usage.input_tokens.no_cache, Some(4)); + assert_eq!(usage.input_tokens.cache_read, Some(0)); + assert_eq!(usage.input_tokens.cache_write, Some(0)); + assert_eq!(usage.output_tokens.total, Some(34)); + assert_eq!(usage.output_tokens.text, Some(34)); + assert_eq!( + usage.raw, + json!({"inputTokens":4,"outputTokens":34,"totalTokens":38}) + .as_object() + .cloned() + ); +} + +/// TS: "should stream tool deltas" (amazon-bedrock/src/amazon-bedrock-chat-language-model.test.ts) +#[tokio::test] +async fn stream_partial_tool_arguments() { + let fetch = MockFetch::new(vec![stream_response(json!([ + {"contentBlockStart":{"contentBlockIndex":0,"start":{"toolUse":{"toolUseId":"tool-use-id","name":"test-tool"}}}}, + {"contentBlockDelta":{"contentBlockIndex":0,"delta":{"toolUse":{"input":"{\"value\":"}}}}, + {"contentBlockDelta":{"contentBlockIndex":0,"delta":{"toolUse":{"input":"\"Sparkle Day\"}"}}}}, + {"contentBlockStop":{"contentBlockIndex":0}}, + {"messageStop":{"stopReason":"tool_use"}} + ]))]); + let schema = + json!({"type":"object","properties":{"value":{"type":"string"}},"required":["value"]}); + let mut opts = options(); + opts.tools = Some(vec![FunctionTool::new("test-tool", schema.clone()).into()]); + let result = provider(&fetch) + .chat("anthropic.model") + .do_stream(&opts) + .await + .unwrap(); + let mut expected = chat_body(); + expected["toolConfig"] = + json!({"tools":[{"toolSpec":{"name":"test-tool","inputSchema":{"json":schema}}}]}); + assert_eq!(fetch.seen()[0].json_body(), expected); + let parts = result + .stream + .collect::>() + .await + .into_iter() + .collect::, _>>() + .unwrap(); + assert_eq!(parts.len(), 8); + assert!( + matches!(&parts[2], StreamPart::ToolInputStart { id, tool_name, .. } if id == "tool-use-id" && tool_name == "test-tool") + ); + assert!( + matches!(&parts[3], StreamPart::ToolInputDelta { id, delta, .. } if id == "tool-use-id" && delta == "{\"value\":") + ); + assert!( + matches!(&parts[4], StreamPart::ToolInputDelta { id, delta, .. } if id == "tool-use-id" && delta == "\"Sparkle Day\"}") + ); + assert!(matches!(&parts[5], StreamPart::ToolInputEnd { id, .. } if id == "tool-use-id")); + let StreamPart::ToolCall(call) = &parts[6] else { + panic!("expected tool call") + }; + assert_eq!(call.tool_call_id, "tool-use-id"); + assert_eq!(call.tool_name, "test-tool"); + assert_eq!(call.input, "{\"value\":\"Sparkle Day\"}"); + assert!( + matches!(&parts[7], StreamPart::Finish { finish_reason, .. } if finish_reason.unified == FinishReasonUnified::ToolCalls && finish_reason.raw.as_deref() == Some("tool_use")) + ); +} + +/// TS: "prefixes the provider message when the error type is present" (amazon-bedrock/src/amazon-bedrock-error.test.ts) +#[tokio::test] +async fn generate_error_response() { + let response = json!({"type":"ValidationException","message":"boom"}); + let mut canned = Canned::json(&response); + canned.status = 400; + let fetch = MockFetch::new(vec![canned]); + let error = provider(&fetch) + .chat("test") + .do_generate(&options()) + .await + .unwrap_err(); + assert_eq!( + fetch.seen()[0].json_body(), + json!({"messages":[{"role":"user","content":[{"text":"Hello"}]}]}) + ); + let AiMuxError::ApiCall(error) = error else { + panic!("expected API error") + }; + assert_eq!(error.message, "ValidationException: boom"); + assert_eq!(error.status_code, Some(400)); + assert!(!error.is_retryable); + assert_eq!(error.data, Some(response)); +} + +/// TS: "should handle single input value and return embeddings" (amazon-bedrock/src/amazon-bedrock-embedding-model.test.ts) +#[tokio::test] +async fn embed_single_value_and_usage() { + let response = json!({"embedding":[-0.09,0.05,-0.02,0.01,0.04],"inputTextTokenCount":8}); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let result = provider(&fetch) + .embedding("amazon.titan-embed-text-v2:0") + .do_embed(&EmbeddingCallOptions::new("sunny day at the beach")) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].url, + "https://bedrock.test/model/amazon.titan-embed-text-v2%3A0/invoke" + ); + assert_eq!( + fetch.seen()[0].json_body(), + json!({"inputText":"sunny day at the beach"}) + ); + assert_eq!( + result.embeddings, + vec![vec![-0.09, 0.05, -0.02, 0.01, 0.04]] + ); + assert_eq!(result.usage.unwrap().tokens, 8); + assert!(result.response.is_none()); +} + +/// TS: "should pass the model and the settings" (amazon-bedrock/src/amazon-bedrock-image-model.test.ts) +#[tokio::test] +async fn generate_images_with_settings() { + let fetch = MockFetch::new(vec![Canned::json( + &json!({"images":["base64-image-1","base64-image-2"]}), + )]); + let mut opts = ImageCallOptions::new("A cute baby sea otter"); + opts.size = Some(Size::new(1024, 1024)); + opts.seed = Some(1234); + opts.provider_options = serde_json::from_value( + json!({"bedrock":{"negativeText":"bad","quality":"premium","cfgScale":1.2}}), + ) + .unwrap(); + let result = provider(&fetch) + .image("amazon.nova-canvas-v1:0") + .do_generate(&opts) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].url, + "https://bedrock.test/model/amazon.nova-canvas-v1%3A0/invoke" + ); + assert_eq!( + fetch.seen()[0].json_body(), + json!({"imageGenerationConfig":{"cfgScale":1.2,"height":1024,"numberOfImages":1,"quality":"premium","seed":1234,"width":1024},"taskType":"TEXT_IMAGE","textToImageParams":{"negativeText":"bad","text":"A cute baby sea otter"}}) + ); + let ImageOutputs::Base64(images) = result.images else { + panic!("expected base64 images") + }; + assert_eq!(images, vec!["base64-image-1", "base64-image-2"]); +} + +/// TS: "should throw error when request is moderated" (amazon-bedrock/src/amazon-bedrock-image-model.test.ts) +#[tokio::test] +async fn image_moderation_error_response() { + let fetch = MockFetch::new(vec![Canned::json( + &json!({"id":"fe7256d1-50d9-4663-8592-85eaf002e80c","status":"Request Moderated","result":null,"progress":null,"details":{"Moderation Reasons":["Derivative Works Filter"]},"preview":null}), + )]); + let error = provider(&fetch) + .image("amazon.nova-canvas-v1:0") + .do_generate(&ImageCallOptions::new( + "Generate something that triggers moderation", + )) + .await + .unwrap_err(); + assert_eq!( + fetch.seen()[0].json_body(), + json!({"imageGenerationConfig":{"numberOfImages":1},"taskType":"TEXT_IMAGE","textToImageParams":{"text":"Generate something that triggers moderation"}}) + ); + assert!( + error + .to_string() + .contains("Amazon Bedrock request was moderated: Derivative Works Filter") + ); +} + +/// TS: "should send request with stringified json documents" (amazon-bedrock/src/reranking/amazon-bedrock-reranking-model.test.ts) +#[tokio::test] +async fn rerank_json_documents_and_parse_ranking() { + let response = json!({"results":[{"index":0,"relevanceScore":0.5110583305358887},{"index":5,"relevanceScore":0.30241215229034424}]}); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let mut opts = RerankingCallOptions::new( + "rainy day", + RerankingDocuments::Object { + values: vec![ + json!({"example":"sunny day at the beach"}), + json!({"example":"rainy day in the city"}), + ], + }, + ); + opts.top_n = Some(2); + opts.provider_options = Some(serde_json::from_value(json!({"bedrock":{"nextToken":"test-token","additionalModelRequestFields":{"test":"test-value"}}})).unwrap()); + let result = provider(&fetch) + .reranking("cohere.rerank-v3-5:0") + .do_rerank(&opts) + .await + .unwrap(); + assert_eq!(fetch.seen()[0].url, "https://bedrock.test/rerank"); + assert_eq!( + fetch.seen()[0].json_body(), + json!({"nextToken":"test-token","queries":[{"textQuery":{"text":"rainy day"},"type":"TEXT"}],"rerankingConfiguration":{"bedrockRerankingConfiguration":{"modelConfiguration":{"modelArn":"arn:aws:bedrock:us-west-2::foundation-model/cohere.rerank-v3-5:0","additionalModelRequestFields":{"test":"test-value"}},"numberOfResults":2},"type":"BEDROCK_RERANKING_MODEL"},"sources":[{"type":"INLINE","inlineDocumentSource":{"type":"JSON","jsonDocument":{"example":"sunny day at the beach"}}},{"type":"INLINE","inlineDocumentSource":{"type":"JSON","jsonDocument":{"example":"rainy day in the city"}}}]}) + ); + assert_eq!(result.ranking.len(), 2); + assert_eq!(result.ranking[0].index, 0); + assert_eq!(result.ranking[0].relevance_score, 0.5110583305358887); + assert_eq!(result.ranking[1].index, 5); + assert_eq!(result.ranking[1].relevance_score, 0.30241215229034424); + assert!(result.provider_metadata.is_none()); + assert_eq!(result.response.unwrap().body.unwrap(), response); +} + +/// TS: "should handle JSON response format in streaming" (amazon-bedrock-chat-language-model.test.ts). +#[tokio::test] +async fn stream_json_response_tool_as_text() { + let fetch = MockFetch::new(vec![stream_response(json!([ + {"contentBlockStart":{"contentBlockIndex":0,"start":{"toolUse":{"toolUseId":"tool-use-id","name":"json"}}}}, + {"contentBlockDelta":{"contentBlockIndex":0,"delta":{"toolUse":{"input":"{\"value\":\"test\"}"}}}}, + {"contentBlockStop":{"contentBlockIndex":0}}, + {"messageStop":{"stopReason":"tool_use"}} + ]))]); + let mut call = CallOptions::new(vec![LanguageModelMessage::user_text("Generate JSON")]); + call.response_format = Some(aimux_core::options::ResponseFormat::Json { + schema: Some( + json!({"type":"object","properties":{"value":{"type":"string"}},"required":["value"]}), + ), + name: None, + description: None, + }); + let result = provider(&fetch) + .chat("anthropic.model") + .do_stream(&call) + .await + .unwrap(); + let parts = result + .stream + .collect::>() + .await + .into_iter() + .collect::, _>>() + .unwrap(); + assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); + assert!( + matches!(&parts[1], StreamPart::ResponseMetadata(meta) if meta.model_id.as_deref() == Some("anthropic.model")) + ); + assert_eq!(parts.len(), 6); + assert!(matches!(&parts[2], StreamPart::TextStart { id, .. } if id == "0")); + assert!( + matches!(&parts[3], StreamPart::TextDelta { id, delta, .. } if id == "0" && delta == "{\"value\":\"test\"}") + ); + assert!(matches!(&parts[4], StreamPart::TextEnd { id, .. } if id == "0")); + let StreamPart::Finish { + finish_reason, + provider_metadata, + .. + } = &parts[5] + else { + panic!("expected finish") + }; + assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(finish_reason.raw.as_deref(), Some("tool_use")); + assert_eq!( + serde_json::to_value(provider_metadata.as_ref().unwrap()).unwrap(), + json!({ + "amazonBedrock":{"isJsonResponseFromTool":true,"stopSequence":null}, + "bedrock":{"isJsonResponseFromTool":true,"stopSequence":null} + }) + ); + assert_eq!( + fetch.seen()[0].json_body(), + json!({ + "messages":[{"role":"user","content":[{"text":"Generate JSON"}]}], + "additionalModelResponseFieldPaths":["/delta/stop_sequence"], + "toolConfig":{"tools":[{"toolSpec":{"name":"json","description":"Respond with a JSON object.", + "inputSchema":{"json":{"type":"object","properties":{"value":{"type":"string"}},"required":["value"]}}}}],"toolChoice":{"any":{}}} + }) + ); +} + +/// TS: "should handle Anthropic provider-defined tools" (amazon-bedrock-chat-language-model.test.ts). +#[tokio::test] +async fn generate_provider_defined_tool_with_input_schema() { + let fetch = MockFetch::new(vec![Canned::json(&json!({ + "output":{"message":{"role":"assistant","content":[{"toolUse":{"toolUseId":"tool-use-id","name":"bash","input":{"command":"ls -l"}}}]}}, + "usage":{"inputTokens":10,"outputTokens":20,"totalTokens":30},"stopReason":"tool_use" + }))]); + let mut call = options(); + call.tools = Some(vec![aimux_core::tool::Tool::Provider( + aimux_core::tool::ProviderTool { + id: "anthropic.bash_20241022".into(), + name: "bash".into(), + args: serde_json::Map::new(), + }, + )]); + call.tool_choice = Some(aimux_core::options::ToolChoice::Auto); + let result = provider(&fetch) + .chat("anthropic.model") + .do_generate(&call) + .await + .unwrap(); + let body = fetch.seen()[0].json_body(); + assert_eq!( + body["additionalModelRequestFields"], + json!({"tool_choice":{"type":"auto"},"anthropic_beta":["computer-use-2024-10-22"]}) + ); + assert_eq!( + body["toolConfig"]["tools"], + json!([{"toolSpec":{"name":"bash","inputSchema":{"json":{ + "$schema":"http://json-schema.org/draft-07/schema#","type":"object", + "properties":{"command":{"type":"string"},"restart":{"type":"boolean"}}, + "required":["command"],"additionalProperties":false + }}}}]) + ); + assert!(body["toolConfig"].get("toolChoice").is_none()); + assert_eq!(result.content.len(), 1); + assert!(result.warnings.is_empty()); + assert!( + matches!(&result.content[0], GenerateContent::ToolCall(tool) if tool.tool_call_id == "tool-use-id" && tool.tool_name == "bash" && tool.input == "{\"command\":\"ls -l\"}") + ); +} diff --git a/aimux-providers/tests/cartesia_transcription_test.rs b/aimux-providers/tests/cartesia_transcription_test.rs index 6be74b644..88058d533 100644 --- a/aimux-providers/tests/cartesia_transcription_test.rs +++ b/aimux-providers/tests/cartesia_transcription_test.rs @@ -231,3 +231,62 @@ async fn should_use_real_date() { assert!(result.response.timestamp.is_some()); assert_eq!(result.response.model_id, Some("best".to_string())); } + +#[tokio::test] +async fn should_warn_about_streaming_options() { + let server = MockServer::start().await; + mock_response(&server, &fixture_response(), &[]).await; + + let config = CartesiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_cartesia(config).unwrap(); + let model = provider.transcription("best"); + + let mut opts = options(mock_audio(), "audio/wav"); + let po = + aimux_core::shared::provider_namespace("cartesia", json!({"streaming": {"delay": "low"}})) + .unwrap(); + opts.provider_options = Some(po); + + let result = model.do_generate(&opts).await.unwrap(); + + assert_eq!(result.warnings.len(), 1); + match &result.warnings[0] { + aimux_core::types::Warning::Unsupported { feature, .. } => { + assert_eq!(feature, "providerOptions.cartesia.streaming"); + } + w => panic!("expected Unsupported warning, got: {w:?}"), + } +} + +#[tokio::test] +async fn should_pass_language_and_timestamp_granularities() { + let server = MockServer::start().await; + mock_response(&server, &fixture_response(), &[]).await; + + let config = CartesiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_cartesia(config).unwrap(); + let model = provider.transcription("best"); + + let mut opts = options(mock_audio(), "audio/wav"); + let po = aimux_core::shared::provider_namespace( + "cartesia", + json!({"language": "en", "timestampGranularities": ["word"]}), + ) + .unwrap(); + opts.provider_options = Some(po); + + model.do_generate(&opts).await.unwrap(); + + let requests = server.received_requests().await.unwrap(); + let form = parse_multipart(&requests[0].body); + assert_eq!(form.get("language").unwrap(), "en"); + assert_eq!(form.get("timestamp_granularities[]").unwrap(), "word"); +} diff --git a/aimux-providers/tests/cassette_multimodal_test.rs b/aimux-providers/tests/cassette_multimodal_test.rs index 49750418d..46fc782d3 100644 --- a/aimux-providers/tests/cassette_multimodal_test.rs +++ b/aimux-providers/tests/cassette_multimodal_test.rs @@ -74,51 +74,6 @@ async fn mount_single(cass: &Value) -> (MockServer, String) { // ── Embedding cassettes ───────────────────────────────────────────────────── -#[tokio::test] -async fn cassette_openai_embedding_query() { - let cass = load_cassette("openai", "TestOpenAI.test_query.json").expect("cassette should load"); - let (_server, base_url) = mount_single(&cass).await; - - let provider = native_provider(base_url); - let model = provider.embedding("text-embedding-3-small"); - - let opts = EmbeddingCallOptions::new("hello"); - let result = model.do_embed(&opts).await.expect("embed should succeed"); - - // OpenAI cassette vectors are stored as strings (base64-encoded floats), - // so provider may return empty Vec or string-as-float parse failures. - // Hard assert: provider must return at least the embeddings array structure. - assert!( - !result.embeddings.is_empty(), - "should have at least one embedding vector" - ); -} - -#[tokio::test] -async fn cassette_openai_embedding_documents() { - let cass = - load_cassette("openai", "TestOpenAI.test_documents.json").expect("cassette should load"); - let (_server, base_url) = mount_single(&cass).await; - - let provider = native_provider(base_url); - let model = provider.embedding("text-embedding-3-small"); - - let opts = EmbeddingCallOptions { - values: vec!["doc1".into(), "doc2".into()], - abort_signal: None, - provider_options: None, - headers: None, - max_retries: None, - timeout: None, - }; - let result = model.do_embed(&opts).await.expect("embed should succeed"); - - assert!( - !result.embeddings.is_empty(), - "should have at least one embedding vector" - ); -} - #[tokio::test] async fn cassette_openai_embedding_error() { let cass = diff --git a/aimux-providers/tests/codex_test.rs b/aimux-providers/tests/codex_test.rs index 7c354b56f..1b69e49a6 100644 --- a/aimux-providers/tests/codex_test.rs +++ b/aimux-providers/tests/codex_test.rs @@ -240,8 +240,6 @@ async fn api_key_streams_text() { } } -// Path B: subscription mode - #[tokio::test] async fn subscription_generate_forces_stream_and_assembles_result() { let server = MockServer::start().await; diff --git a/aimux-providers/tests/cohere_upstream_e2e_test.rs b/aimux-providers/tests/cohere_upstream_e2e_test.rs new file mode 100644 index 000000000..3e67f898c --- /dev/null +++ b/aimux-providers/tests/cohere_upstream_e2e_test.rs @@ -0,0 +1,480 @@ +//! Selected end-to-end cases from @ai-sdk/cohere 4.0.52. + +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::{CallOptions, Tool}; +use aimux_core::reranking_model::{RerankingCallOptions, RerankingDocuments, RerankingModel}; +use aimux_core::result::GenerateContent; +use aimux_core::shared::Warning; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::FunctionTool; +use aimux_core::types::{FinishReasonUnified, Usage}; +use aimux_providers::cohere::{CohereProvider, CohereProviderSettings, create_cohere}; +use futures::StreamExt; +use mock_fetch::{Canned, MockFetch}; +use serde_json::{Value, json}; + +fn provider(fetch: &std::sync::Arc) -> CohereProvider { + create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} + +fn options() -> CallOptions { + CallOptions::new(vec![ + LanguageModelMessage::System { + content: "you are a friendly bot!".into(), + provider_options: None, + }, + LanguageModelMessage::user_text("Hello"), + ]) +} + +fn request(streaming: bool) -> Value { + let mut body = json!({"model":"command-r-plus","messages":[ + {"role":"system","content":"you are a friendly bot!"}, + {"role":"user","content":"Hello"} + ]}); + if streaming { + body["stream"] = json!(true); + } + body +} + +fn assert_request(fetch: &MockFetch, endpoint: &str, body: &Value) { + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "POST"); + assert_eq!(seen[0].url, format!("https://api.cohere.com/v2/{endpoint}")); + assert_eq!(seen[0].json_body(), *body); + assert_eq!(seen[0].headers["authorization"], "Bearer test-api-key"); + assert_eq!(seen[0].headers["content-type"], "application/json"); +} + +fn sse(body: String) -> Canned { + Canned { + status: 200, + headers: vec![("content-type".into(), "text/event-stream".into())], + body: body.into_bytes(), + } +} + +fn assert_usage(usage: &Usage, raw: Value) { + assert_eq!(usage.input_tokens.total, Some(507)); + assert_eq!(usage.input_tokens.no_cache, Some(507)); + assert_eq!(usage.input_tokens.cache_read, None); + assert_eq!(usage.input_tokens.cache_write, None); + assert_eq!(usage.output_tokens.total, Some(10)); + assert_eq!(usage.output_tokens.text, Some(10)); + assert_eq!(usage.output_tokens.reasoning, None); + assert_eq!(usage.raw, raw.as_object().cloned()); +} + +fn cohere_text() -> Value { + serde_json::from_str( + r###"{ + "id": "e7592632-1e3d-424f-b129-bd5f9f980f7b", + "message": { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of France is Paris." + } + ] + }, + "finish_reason": "COMPLETE", + "usage": { + "billed_units": { + "input_tokens": 12, + "output_tokens": 7 + }, + "tokens": { + "input_tokens": 507, + "output_tokens": 10 + }, + "cached_tokens": 448 + } +}"###, + ) + .unwrap() +} +fn cohere_tool_call() -> Value { + serde_json::from_str(r###"{ + "id": "f201af17-e24a-4396-8f6a-98e8bf9c3432", + "message": { + "role": "assistant", + "tool_plan": "I will use the weather tool to find out the weather in San Francisco. I will also use the cityAttractions tool to find out what attractions are in San Francisco.", + "tool_calls": [ + { + "id": "weather_dqgshstja6p9", + "type": "function", + "function": { + "name": "weather", + "arguments": "{\"location\":\"San Francisco\"}" + } + }, + { + "id": "cityAttractions_dcxfx4myvx68", + "type": "function", + "function": { + "name": "cityAttractions", + "arguments": "{\"city\":\"San Francisco\"}" + } + } + ] + }, + "finish_reason": "TOOL_CALL", + "usage": { + "billed_units": { + "input_tokens": 119, + "output_tokens": 52 + }, + "tokens": { + "input_tokens": 1549, + "output_tokens": 103 + }, + "cached_tokens": 992 + } +}"###).unwrap() +} +fn cohere_text_chunks() -> Vec { + r###"{"id":"321d178c-2c12-44d3-ae42-2f5510f6b1cc","type":"message-start","delta":{"message":{"role":"assistant","content":[],"tool_plan":"","tool_calls":[],"citations":[]}}} +{"type":"content-start","index":0,"delta":{"message":{"content":{"type":"text","text":""}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":"The"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":" capital"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":" of"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":" France"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":" is"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":" Paris"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":"."}}}} +{"type":"content-end","index":0} +{"type":"message-end","delta":{"finish_reason":"COMPLETE","usage":{"billed_units":{"input_tokens":12,"output_tokens":7},"tokens":{"input_tokens":507,"output_tokens":10},"cached_tokens":448}}}"###.lines().filter(|line| !line.trim().is_empty()).map(|line| serde_json::from_str(line).unwrap()).collect() +} + +fn embedding_fixture() -> Value { + serde_json::from_str( + r#"{ + "id": "f5aa3e7b-f011-4c5c-a825-f94669f760e5", + "texts": ["sunny day at the beach", "rainy day in the city"], + "embeddings": { + "float": [ + [0.03302002, 0.020904541, -0.019744873, -0.0625, 0.04437256], + [-0.04660034, 0.00037765503, -0.061157227, -0.08239746, -0.010360718] + ] + }, + "meta": { + "api_version": { + "version": "2" + }, + "billed_units": { + "input_tokens": 10 + } + }, + "response_type": "embeddings_by_type" +} +"#, + ) + .unwrap() +} + +fn reranking_fixture() -> Value { + serde_json::from_str( + r#"{ + "id": "b44fe75b-e3d3-489a-b61e-1a1aede3ef72", + "results": [ + { + "index": 1, + "relevance_score": 0.10183054 + }, + { + "index": 0, + "relevance_score": 0.03762639 + } + ], + "meta": { + "api_version": { + "version": "2" + }, + "billed_units": { + "search_units": 1 + } + } +} +"#, + ) + .unwrap() +} + +/// TS: "should extract text response" (cohere/src/cohere-chat-language-model.test.ts) +#[tokio::test] +async fn generate_text() { + let fixture = cohere_text(); + let fetch = MockFetch::new(vec![Canned::json(&fixture)]); + let result = provider(&fetch) + .chat("command-r-plus") + .do_generate(&options()) + .await + .unwrap(); + assert_request(&fetch, "chat", &request(false)); + assert_eq!(result.request.unwrap().body, Some(request(false))); + assert_eq!( + result.content, + vec![GenerateContent::Text { + text: "The capital of France is Paris.".into(), + provider_metadata: None, + }] + ); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(result.finish_reason.raw.as_deref(), Some("COMPLETE")); + assert_usage(&result.usage, fixture["usage"].clone()); + assert!(result.warnings.is_empty()); + assert!(result.provider_metadata.is_none()); + let response = result.response.unwrap(); + assert_eq!(response.body, Some(fixture)); + assert_eq!(response.id, None); + assert_eq!(response.timestamp, None); + assert_eq!(response.model_id, None); +} + +/// TS: "should stream text deltas" (cohere/src/cohere-chat-language-model.test.ts) +#[tokio::test] +async fn stream_text() { + let chunks = cohere_text_chunks(); + let body = chunks + .iter() + .map(|chunk| { + format!( + "event: {}\ndata: {chunk}\n\n", + chunk["type"].as_str().unwrap() + ) + }) + .collect(); + let fetch = MockFetch::new(vec![sse(body)]); + let result = provider(&fetch) + .chat("command-r-plus") + .do_stream(&options()) + .await + .unwrap(); + assert_request(&fetch, "chat", &request(true)); + assert_eq!(result.request.unwrap().body, Some(request(true))); + let parts: Vec<_> = result.stream.map(|part| part.unwrap()).collect().await; + assert_eq!(parts.len(), 12); + assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); + assert!( + matches!(&parts[1], StreamPart::ResponseMetadata(meta) if meta.id.as_deref() == Some("321d178c-2c12-44d3-ae42-2f5510f6b1cc")) + ); + assert!(matches!(&parts[2], StreamPart::TextStart { id, .. } if id == "0")); + let deltas: Vec<_> = parts + .iter() + .filter_map(|part| match part { + StreamPart::TextDelta { id, delta, .. } => { + assert_eq!(id, "0"); + Some(delta.as_str()) + } + _ => None, + }) + .collect(); + assert_eq!( + deltas, + vec!["The", " capital", " of", " France", " is", " Paris", "."] + ); + assert!(matches!(&parts[10], StreamPart::TextEnd { id, .. } if id == "0")); + match &parts[11] { + StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } => { + assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(finish_reason.raw.as_deref(), Some("COMPLETE")); + assert_usage(usage, chunks.last().unwrap()["delta"]["usage"].clone()); + assert!(provider_metadata.is_none()); + } + part => panic!("expected finish, got {part:?}"), + } +} + +/// TS: "should extract tool calls" (cohere/src/cohere-chat-language-model.test.ts) +#[tokio::test] +async fn generate_tool_calls() { + let fixture = cohere_tool_call(); + let fetch = MockFetch::new(vec![Canned::json(&fixture)]); + let schema = json!({"type":"object","properties":{"value":{"type":"string"}},"required":["value"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}); + let mut opts = options(); + opts.tools = Some(vec![Tool::Function(FunctionTool::new( + "test-tool", + schema.clone(), + ))]); + let result = provider(&fetch) + .chat("command-r-plus") + .do_generate(&opts) + .await + .unwrap(); + let mut expected = request(false); + expected["tools"] = + json!([{"type":"function","function":{"name":"test-tool","parameters":schema}}]); + assert_request(&fetch, "chat", &expected); + assert_eq!(result.content.len(), 2); + for (content, wire) in result + .content + .iter() + .zip(fixture["message"]["tool_calls"].as_array().unwrap()) + { + match content { + GenerateContent::ToolCall(call) => { + assert_eq!(call.tool_call_id, wire["id"]); + assert_eq!(call.tool_name, wire["function"]["name"]); + assert_eq!(call.input, wire["function"]["arguments"]); + assert!(call.provider_metadata.is_none()); + } + other => panic!("expected tool call, got {other:?}"), + } + } + assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); + assert_eq!(result.finish_reason.raw.as_deref(), Some("TOOL_CALL")); + assert_eq!(result.usage.input_tokens.total, Some(1549)); + assert_eq!(result.usage.output_tokens.total, Some(103)); + assert_eq!(result.usage.raw, fixture["usage"].as_object().cloned()); +} + +/// TS: "should handle unparsable stream parts" (cohere/src/cohere-chat-language-model.test.ts) +#[tokio::test] +async fn malformed_stream_response() { + let fetch = MockFetch::new(vec![sse( + "event: foo-message\ndata: {unparsable}\n\n".into() + )]); + let result = provider(&fetch) + .chat("command-r-plus") + .do_stream(&options()) + .await + .unwrap(); + assert_request(&fetch, "chat", &request(true)); + let parts: Vec<_> = result.stream.map(|part| part.unwrap()).collect().await; + assert_eq!(parts.len(), 3); + assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); + assert!(matches!(&parts[1], StreamPart::Error { .. })); + match &parts[2] { + StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } => { + assert_eq!(finish_reason.unified, FinishReasonUnified::Error); + assert_eq!(finish_reason.raw, None); + assert_eq!( + serde_json::to_value(usage).unwrap(), + serde_json::to_value(Usage::default()).unwrap() + ); + assert!(provider_metadata.is_none()); + } + part => panic!("expected finish, got {part:?}"), + } +} + +/// TS: "should extract embedding" (cohere/src/cohere-embedding-model.test.ts) +#[tokio::test] +async fn embed_values() { + let fixture = embedding_fixture(); + let fetch = MockFetch::new(vec![Canned::json(&fixture)]); + let mut opts = EmbeddingCallOptions::new("sunny day at the beach"); + opts.values.push("rainy day in the city".into()); + let result = provider(&fetch) + .embedding("embed-english-v3.0") + .do_embed(&opts) + .await + .unwrap(); + assert_request( + &fetch, + "embed", + &json!({"model":"embed-english-v3.0","embedding_types":["float"],"texts":opts.values,"input_type":"search_query"}), + ); + let expected: Vec> = + serde_json::from_value(fixture["embeddings"]["float"].clone()).unwrap(); + assert_eq!(result.embeddings, expected); + assert_eq!(result.usage.unwrap().tokens, 10); + assert_eq!(result.response.unwrap().body, Some(fixture)); + assert!(result.provider_metadata.is_none()); + assert!(result.warnings.is_empty()); +} + +async fn rerank(object_documents: bool) { + let fixture = reranking_fixture(); + let fetch = MockFetch::new(vec![Canned::json(&fixture)]); + let texts = vec![ + "sunny day at the beach".to_string(), + "rainy day in the city".to_string(), + ]; + let documents = if object_documents { + RerankingDocuments::Object { + values: texts.iter().map(|text| json!({"example":text})).collect(), + } + } else { + RerankingDocuments::Text { + values: texts.clone(), + } + }; + let mut opts = RerankingCallOptions::new("rainy day", documents); + opts.top_n = Some(2); + opts.provider_options = Some( + serde_json::from_value(json!({"cohere":{"maxTokensPerDoc":1000,"priority":1}})).unwrap(), + ); + let result = provider(&fetch) + .reranking("rerank-english-v3.0") + .do_rerank(&opts) + .await + .unwrap(); + let wire_documents = if object_documents { + texts + .iter() + .map(|text| json!({"example":text}).to_string()) + .collect::>() + } else { + texts + }; + assert_request( + &fetch, + "rerank", + &json!({"documents":wire_documents,"model":"rerank-english-v3.0","query":"rainy day","top_n":2,"max_tokens_per_doc":1000,"priority":1}), + ); + assert_eq!(result.ranking.len(), 2); + assert_eq!(result.ranking[0].index, 1); + assert_eq!(result.ranking[0].relevance_score, 0.10183054); + assert_eq!(result.ranking[1].index, 0); + assert_eq!(result.ranking[1].relevance_score, 0.03762639); + assert!(result.provider_metadata.is_none()); + let warnings = result.warnings.unwrap(); + if object_documents { + assert_eq!(warnings.len(), 1); + assert!( + matches!(&warnings[0], Warning::Compatibility { feature, details } if feature == "object documents" && details.as_deref() == Some("Object documents are converted to strings.")) + ); + } else { + assert!(warnings.is_empty()); + } + let response = result.response.unwrap(); + assert_eq!( + response.id.as_deref(), + Some("b44fe75b-e3d3-489a-b61e-1a1aede3ef72") + ); + assert_eq!(response.body, Some(fixture)); +} + +/// TS: "should send request with stringified json documents" (cohere/src/reranking/cohere-reranking-model.test.ts) +#[tokio::test] +async fn rerank_object_documents() { + rerank(true).await; +} + +/// TS: "should send request with text documents" (cohere/src/reranking/cohere-reranking-model.test.ts) +#[tokio::test] +async fn rerank_text_documents() { + rerank(false).await; +} diff --git a/aimux-providers/tests/common/replay.rs b/aimux-providers/tests/common/replay.rs index e56bc1449..a52a70718 100644 --- a/aimux-providers/tests/common/replay.rs +++ b/aimux-providers/tests/common/replay.rs @@ -84,7 +84,7 @@ use std::path::Path; use base64::Engine as _; use base64::engine::general_purpose::STANDARD as BASE64; use serde_json::Value; -use wiremock::matchers::{method, path}; +use wiremock::matchers::method; use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; /// Replay the ordered HTTP exchanges of a unified runtime Recording. Request @@ -387,7 +387,7 @@ async fn mount_entries(server: &MockServer, cassettes: Vec) -> usize { for ((m, p), entries) in groups { Mock::given(method(m)) - .and(path(p)) + .and(move |req: &Request| percent_decode(req.url.path()) == p) .respond_with(CassetteRespond { entries }) .mount(server) .await; diff --git a/aimux-providers/tests/conformance_test.rs b/aimux-providers/tests/conformance_test.rs index f4def891f..a05fe1e76 100644 --- a/aimux-providers/tests/conformance_test.rs +++ b/aimux-providers/tests/conformance_test.rs @@ -872,34 +872,6 @@ mod bedrock_conformance { } } } - - #[tokio::test] - async fn do_stream_returns_parts() { - let server = MockServer::start().await; - mount_cassettes(&server, "tests/cassettes/bedrock").await; - - let provider = make_provider(&server); - let model = provider.chat("amazon.nova-lite-v1:0"); - - let result = model.do_stream(&default_options(test_prompt())).await; - - match result { - Ok(stream_result) => { - let parts = collect_stream(stream_result).await; - assert!( - has_finish(&parts) || !parts.is_empty(), - "stream should produce some parts" - ); - } - Err(e) => { - if is_infrastructure_error(&e) { - panic!("request did not match any cassette (404): {e:?}"); - } - let msg = e.to_string(); - assert!(!msg.contains("panic"), "unexpected error: {msg}"); - } - } - } } // ════════════════════════════════════════════════════════════════════════════ diff --git a/aimux-providers/tests/data_loss_regression_test.rs b/aimux-providers/tests/data_loss_regression_test.rs index b880d2524..7bd988e72 100644 --- a/aimux-providers/tests/data_loss_regression_test.rs +++ b/aimux-providers/tests/data_loss_regression_test.rs @@ -30,13 +30,14 @@ use std::path::PathBuf; use futures::StreamExt; use serde_json::{Value, json}; -use wiremock::matchers::{method, path as path_matcher}; +use wiremock::matchers::{method, path as path_matcher, query_param}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; -use aimux_core::result::{GenerateContent, ReasoningOutput, Source, StreamResult}; +use aimux_core::result::{GenerateContent, GeneratedFile, ReasoningOutput, Source, StreamResult}; +use aimux_core::shared::{FileBytes, GeneratedFileData}; use aimux_core::stream_part::StreamPart; use aimux_core::tool::{ProviderTool, Tool}; use aimux_core::types::ProviderMetadata; @@ -101,7 +102,12 @@ fn cassette(rel: &str) -> Cassette { /// Mount a cassette's response on `server` at its recorded path. async fn mount(server: &MockServer, c: &Cassette) { Mock::given(method("POST")) - .and(path_matcher(c.request_path.clone())) + .and({ + // Model ids reach the wire percent-encoded; recordings differ in + // whether they kept the encoding. + let recorded = c.request_path.replace("%3A", ":"); + move |req: &wiremock::Request| req.url.path().replace("%3A", ":") == recorded + }) .respond_with( ResponseTemplate::new(c.status) .insert_header("content-type", "application/json") @@ -154,6 +160,26 @@ async fn collect(result: StreamResult) -> Vec { parts } +fn file_base64(content: &GenerateContent) -> &str { + match content { + GenerateContent::File(GeneratedFile { + data: + GeneratedFileData::Data { + data: FileBytes::Base64(data), + }, + .. + }) => data, + other => panic!("expected File with base64 data, got {other:?}"), + } +} + +fn files(content: &[GenerateContent]) -> Vec<&GenerateContent> { + content + .iter() + .filter(|part| matches!(part, GenerateContent::File(_))) + .collect() +} + fn texts(content: &[GenerateContent]) -> Vec<&str> { content .iter() @@ -302,6 +328,9 @@ fn sources(content: &[GenerateContent]) -> Vec> { // gemini/test_google_image_and_text_output.json // ═══════════════════════════════════════════════════════════════════════════ +const PNG_BASE64_PREFIX: &str = "iVBORw0KGgoAAAANSUhEUgAA"; +const NANO_BANANA_BASE64_LEN: usize = 258_820; + fn google_at(uri: &str) -> GoogleProvider { create_google(GoogleProviderSettings { api_key: Some("test-api-key".to_string()), @@ -311,6 +340,138 @@ fn google_at(uri: &str) -> GoogleProvider { .unwrap() } +#[tokio::test] +async fn finding_1_gemini_inline_data_surfaces_as_file() { + let c = cassette("gemini/nano_banana_image_generation_smoke.json"); + let server = MockServer::start().await; + mount(&server, &c).await; + + let result = google_at(&server.uri()) + .chat("gemini-2.5-flash-image") + .do_generate(&opts()) + .await + .expect("do_generate should succeed"); + + let f = files(&result.content); + assert_eq!(f.len(), 1, "exactly one inlineData part → one File"); + match f[0] { + GenerateContent::File(GeneratedFile { media_type, .. }) => { + assert_eq!(media_type, "image/png"); + } + other => panic!("expected File, got {other:?}"), + } + let b64 = file_base64(f[0]); + assert!( + b64.starts_with(PNG_BASE64_PREFIX), + "image bytes must be the recorded PNG, got prefix {:?}", + &b64[..b64.len().min(32)] + ); + assert_eq!( + b64.len(), + NANO_BANANA_BASE64_LEN, + "the whole base64 payload must survive, not a truncated head" + ); +} + +#[tokio::test] +async fn finding_1_gemini_image_and_text_output_both_survive() { + let c = cassette("gemini/test_google_image_and_text_output.json"); + let server = MockServer::start().await; + mount(&server, &c).await; + + let result = google_at(&server.uri()) + .chat("gemini-2.5-flash-image") + .do_generate(&opts()) + .await + .expect("do_generate should succeed"); + + // Text part first, then the inlineData part — order matters, upstream + // walks `parts` in order. + assert_eq!(result.content.len(), 2, "one text + one file"); + assert!( + matches!(result.content[0], GenerateContent::Text { .. }), + "text part comes first" + ); + assert!( + matches!(result.content[1], GenerateContent::File(_)), + "inlineData part comes second" + ); + + let t = texts(&result.content); + assert_eq!(t.len(), 1); + assert!( + t[0].starts_with("Once, in a hidden cenote, lived an axolotl named Pip"), + "recorded story text must survive verbatim, got {:?}", + &t[0][..t[0].len().min(60)] + ); + + let f = files(&result.content); + assert_eq!(f.len(), 1); + match f[0] { + GenerateContent::File(GeneratedFile { media_type, .. }) => { + assert_eq!(media_type, "image/png") + } + other => panic!("expected File, got {other:?}"), + } + let b64 = file_base64(f[0]); + assert!(b64.starts_with(PNG_BASE64_PREFIX)); + assert_eq!(b64.len(), 2_580_504); +} + +/// The streaming path had the same hole. The chunk is built from the cassette's +/// own `parts` array so the bytes are the recorded ones. +#[tokio::test] +async fn finding_1_gemini_inline_data_streams_as_file_part() { + let body = cassette_json("gemini/nano_banana_image_generation_smoke.json"); + let candidate = &body["candidates"][0]; + let sse = format!( + "data: {}\n\n", + json!({ "candidates": [{ "content": candidate["content"].clone(), "finishReason": "STOP" }] }) + ); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_matcher( + "/v1beta/models/gemini-2.5-flash-image:streamGenerateContent", + )) + .and(query_param("alt", "sse")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse), + ) + .mount(&server) + .await; + + let parts = collect( + google_at(&server.uri()) + .chat("gemini-2.5-flash-image") + .do_stream(&opts()) + .await + .expect("do_stream should succeed"), + ) + .await; + + let file_parts: Vec<_> = parts + .iter() + .filter_map(|p| match p { + StreamPart::File(GeneratedFile { + data: + GeneratedFileData::Data { + data: FileBytes::Base64(b64), + }, + media_type, + .. + }) => Some((b64.as_str(), media_type.as_str())), + _ => None, + }) + .collect(); + assert_eq!(file_parts.len(), 1, "expected one StreamPart::File"); + assert_eq!(file_parts[0].1, "image/png"); + assert!(file_parts[0].0.starts_with(PNG_BASE64_PREFIX)); + assert_eq!(file_parts[0].0.len(), NANO_BANANA_BASE64_LEN); +} + // ═══════════════════════════════════════════════════════════════════════════ // Finding 3 — Gemini `part.thought` never read (thinking merged into the answer) // Finding 5 — Gemini `thoughtSignature` dropped @@ -638,6 +799,9 @@ async fn finding_2_anthropic_web_search_result_mapped_and_sources_emitted() { let by_url: BTreeMap<&str, &ProviderMetadata> = s .iter() .filter_map(|(_, _, url, _, m)| Some(((*url)?, (*m)?))) + // Citations of the same pages are sources too; they carry the cited + // text instead of a page age. + .filter(|(_, m)| m["anthropic"].contains_key("pageAge")) .collect(); for (url, expected_page_age) in [ ( @@ -1223,10 +1387,6 @@ async fn finding_10_bedrock_generate_reasoning_signature_in_provider_metadata() &sig[..sig.len().min(24)] ); assert_eq!(sig.len(), 496, "the signature must not be truncated"); - assert!( - m.get("bedrock").is_none(), - "the legacy provider key is not written" - ); // The visible answer is separate from the reasoning. let t = texts(&result.content); @@ -1245,7 +1405,10 @@ fn encode_events(events: &[(&str, &str, Value)]) -> Vec { async fn bedrock_stream_parts(model: &str, events: &[(&str, &str, Value)]) -> Vec { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path_matcher(format!("/model/{model}/converse-stream"))) + .and(path_matcher(format!( + "/model/{}/converse-stream", + model.replace(':', "%3A") + ))) .respond_with( ResponseTemplate::new(200) .insert_header("content-type", "application/vnd.amazon.eventstream") @@ -1309,15 +1472,14 @@ async fn finding_10_bedrock_stream_reasoning_signature_emitted_as_metadata_delta let end_meta = parts .iter() .find_map(|p| match p { - StreamPart::ReasoningEnd { + StreamPart::ReasoningDelta { provider_metadata: Some(m), .. } => Some(m), _ => None, }) - .expect("the ReasoningEnd must carry the signature"); + .expect("a reasoning delta must carry the signature"); assert_eq!(end_meta["amazonBedrock"]["signature"], json!(signature)); - assert!(end_meta.get("bedrock").is_none()); } #[tokio::test] @@ -1370,10 +1532,6 @@ async fn finding_27_bedrock_stream_performance_config_and_service_tier_reach_fin meta["amazonBedrock"]["serviceTier"], json!({ "type": "flex" }) ); - assert!( - meta.get("bedrock").is_none(), - "the legacy provider key is not written" - ); } #[tokio::test] @@ -1430,7 +1588,6 @@ async fn finding_27_bedrock_stream_guardrail_trace_reaches_finish() { json!(397), "nested metrics must not be flattened away" ); - assert!(meta.get("bedrock").is_none()); } #[tokio::test] @@ -1471,5 +1628,4 @@ async fn finding_26_bedrock_stream_stop_sequence_reaches_finish() { }) .expect("Finish must carry provider_metadata"); assert_eq!(meta["amazonBedrock"]["stopSequence"], json!("STOP")); - assert!(meta.get("bedrock").is_none()); } diff --git a/aimux-providers/tests/dataforseo_test.rs b/aimux-providers/tests/dataforseo_test.rs index 34a7930f4..bb8c09422 100644 --- a/aimux-providers/tests/dataforseo_test.rs +++ b/aimux-providers/tests/dataforseo_test.rs @@ -222,5 +222,3 @@ async fn status_401_maps_to_auth_error() { "expected Auth error, got {result:?}" ); } - -// -- Provider trait ---------------------------------------------------------- diff --git a/aimux-providers/tests/deepgram_transcription_test.rs b/aimux-providers/tests/deepgram_transcription_test.rs index d52f99e60..bd56ed904 100644 --- a/aimux-providers/tests/deepgram_transcription_test.rs +++ b/aimux-providers/tests/deepgram_transcription_test.rs @@ -152,6 +152,53 @@ async fn should_extract_text() { assert_eq!(result.text, "galileo was an american robotic space program"); } +#[tokio::test] +async fn should_pass_detect_language_query_param() { + let server = MockServer::start().await; + mock_response(&server, &fixture_response()).await; + + let config = DeepgramProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_deepgram(config).unwrap(); + let model = provider.transcription("nova-3"); + + let mut opts = options(mock_audio(), "audio/wav"); + let po = aimux_core::shared::provider_namespace("deepgram", json!({"detectLanguage": true})) + .unwrap(); + opts.provider_options = Some(po); + + model.do_generate(&opts).await.unwrap(); + + let requests = server.received_requests().await.unwrap(); + assert!(requests[0].url.as_str().contains("detect_language=true")); +} + +#[tokio::test] +async fn should_return_detected_language() { + let server = MockServer::start().await; + mock_response(&server, &fixture_response()).await; + + let config = DeepgramProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_deepgram(config).unwrap(); + let model = provider.transcription("nova-3"); + + let mut opts = options(mock_audio(), "audio/wav"); + let po = aimux_core::shared::provider_namespace("deepgram", json!({"detectLanguage": true})) + .unwrap(); + opts.provider_options = Some(po); + + let result = model.do_generate(&opts).await.unwrap(); + + assert_eq!(result.language, Some("en".to_string())); +} + #[tokio::test] async fn should_include_response_data() { let server = MockServer::start().await; @@ -207,6 +254,41 @@ async fn should_use_real_date() { assert_eq!(result.response.model_id, Some("nova-3".to_string())); } +#[tokio::test] +async fn should_return_language_from_inline_response() { + let server = MockServer::start().await; + mock_response( + &server, + &json!({ + "metadata": {"duration": 1.0}, + "results": { + "channels": [{ + "detected_language": "sv", + "alternatives": [{"transcript": "hej", "words": []}] + }] + } + }), + ) + .await; + + let config = DeepgramProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_deepgram(config).unwrap(); + let model = provider.transcription("nova-3"); + + let mut opts = options(mock_audio(), "audio/wav"); + let po = aimux_core::shared::provider_namespace("deepgram", json!({"detectLanguage": true})) + .unwrap(); + opts.provider_options = Some(po); + + let result = model.do_generate(&opts).await.unwrap(); + + assert_eq!(result.language, Some("sv".to_string())); +} + #[tokio::test] async fn should_return_none_language_when_not_detected() { let server = MockServer::start().await; diff --git a/aimux-providers/tests/deepseek_chat_test.rs b/aimux-providers/tests/deepseek_chat_test.rs index 995e628b3..d2ad07937 100644 --- a/aimux-providers/tests/deepseek_chat_test.rs +++ b/aimux-providers/tests/deepseek_chat_test.rs @@ -13,19 +13,18 @@ //! Left out: the `json response format with structured outputs` block (it needs //! a model built with `supportsStructuredOutputs`, which only the Azure wiring //! sets and this package does not expose), the tool-result `content` output -//! cases, -//! and the snapshot fixtures (the responses are inlined here). +//! cases, and the snapshot fixtures (the responses are inlined here). //! //! HTTP is a local `wiremock` server; every model is built with //! `create_deepseek(..).chat(..)`. +use aimux_core::error::AiMuxError; use aimux_core::tool::RawToolCall; use futures::StreamExt; use serde_json::{Value, json}; use wiremock::matchers::{method, path_regex}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{ AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, @@ -47,20 +46,18 @@ use aimux_providers::deepseek::{ // --------------------------------------------------------------------------- fn with_options(mut message: LanguageModelMessage, options: Value) -> LanguageModelMessage { - let provider_options = match &mut message { - LanguageModelMessage::System { - provider_options, .. - } - | LanguageModelMessage::User { - provider_options, .. - } - | LanguageModelMessage::Assistant { - provider_options, .. - } - | LanguageModelMessage::Tool { - provider_options, .. - } => provider_options, - }; + let (LanguageModelMessage::System { + provider_options, .. + } + | LanguageModelMessage::User { + provider_options, .. + } + | LanguageModelMessage::Assistant { + provider_options, .. + } + | LanguageModelMessage::Tool { + provider_options, .. + }) = &mut message; *provider_options = Some(serde_json::from_value(options).unwrap()); message } @@ -204,7 +201,7 @@ async fn generate_body(model_id: &str, options: &CallOptions) -> Value { .await .unwrap() .request - .and_then(|request| request.body) + .and_then(|r| r.body) .unwrap() } @@ -458,6 +455,17 @@ async fn text_should_preserve_supported_sampling_options_when_v4_thinking_is_dis ); } +#[tokio::test] +async fn text_should_send_max_tokens_and_stop_sequences() { + let mut options = options(); + options.max_output_tokens = Some(64); + options.stop_sequences = Some(vec!["END".to_string()]); + let body = generate_body("deepseek-chat", &options).await; + assert_eq!(body["max_tokens"], 64); + assert_eq!(body["stop"], json!(["END"])); + assert!(body.get("max_completion_tokens").is_none()); +} + #[tokio::test] async fn text_should_send_message_names() { let name = |value: &str| json!({ "deepseek": { "name": value } }); @@ -621,7 +629,7 @@ async fn text_should_send_all_strict_tools_on_the_beta_endpoint() { .await .unwrap() .request - .and_then(|request| request.body) + .and_then(|r| r.body) .unwrap(); assert_eq!(body["tools"][0]["type"], "function"); assert_eq!(body["tools"][0]["function"]["name"], "getWeather"); @@ -684,17 +692,11 @@ async fn text_should_extract_text_content() { assert_eq!(result.usage.input_tokens.total, Some(13)); assert_eq!(result.usage.output_tokens.total, Some(300)); assert_eq!( - result - .response - .as_ref() - .and_then(|response| response.id.as_deref()), + result.response.as_ref().and_then(|r| r.id.as_deref()), Some("00f10ecd-60b3-4707-b5db-e4bcadf7aea1") ); assert_eq!( - result - .response - .as_ref() - .and_then(|response| response.model_id.as_deref()), + result.response.as_ref().and_then(|r| r.model_id.as_deref()), Some("deepseek-chat") ); } @@ -938,7 +940,7 @@ async fn top_level_reasoning_none_should_keep_the_temperature() { options.temperature = Some(0.4); let result = generate_result("deepseek-reasoner", &options).await; assert_eq!( - result.request.and_then(|request| request.body).unwrap()["temperature"], + result.request.and_then(|r| r.body).unwrap()["temperature"], 0.4 ); assert!(result.warnings.is_empty()); @@ -952,7 +954,7 @@ async fn top_level_reasoning_xhigh_should_map_to_reasoning_effort_max() { ) .await; assert_eq!( - result.request.and_then(|request| request.body).unwrap()["reasoning_effort"], + result.request.and_then(|r| r.body).unwrap()["reasoning_effort"], "max" ); assert!(warning_values(&result.warnings).contains(&serde_json::to_value(compatibility( @@ -971,7 +973,7 @@ async fn top_level_reasoning_low_should_map_to_reasoning_effort_low_without_a_co ) .await; assert_eq!( - result.request.and_then(|request| request.body).unwrap()["reasoning_effort"], + result.request.and_then(|r| r.body).unwrap()["reasoning_effort"], "low" ); assert!( @@ -991,7 +993,7 @@ async fn top_level_reasoning_medium_should_map_to_reasoning_effort_high_with_a_c ) .await; assert_eq!( - result.request.and_then(|request| request.body).unwrap()["reasoning_effort"], + result.request.and_then(|r| r.body).unwrap()["reasoning_effort"], "high" ); assert_eq!( @@ -1012,7 +1014,7 @@ async fn top_level_reasoning_minimal_should_map_to_reasoning_effort_low_with_com ) .await; assert_eq!( - result.request.and_then(|request| request.body).unwrap()["reasoning_effort"], + result.request.and_then(|r| r.body).unwrap()["reasoning_effort"], "low" ); assert_eq!( @@ -1037,7 +1039,7 @@ async fn top_level_reasoning_should_map_provider_options_reasoning_effort() { ); let result = generate_result("deepseek-reasoner", &options).await; assert_eq!( - result.request.and_then(|request| request.body).unwrap()["reasoning_effort"], + result.request.and_then(|r| r.body).unwrap()["reasoning_effort"], output ); let expected = if warns { @@ -1063,7 +1065,7 @@ async fn top_level_reasoning_should_map_legacy_thinking_type_adaptive_to_enabled ); let result = generate_result("deepseek-reasoner", &options).await; assert_eq!( - result.request.and_then(|request| request.body).unwrap()["thinking"], + result.request.and_then(|r| r.body).unwrap()["thinking"], json!({ "type": "enabled" }) ); assert!(warning_values(&result.warnings).contains(&serde_json::to_value(compatibility( @@ -1189,6 +1191,49 @@ async fn tool_call_should_send_correct_request_body() { ); } +#[tokio::test] +async fn tool_call_should_send_the_tool_choice() { + use aimux_core::options::ToolChoice; + let mut options = tool_options(); + options.tool_choice = Some(ToolChoice::Tool { + tool_name: "weather".to_string(), + }); + assert_eq!( + generate_body("deepseek-reasoner", &options).await["tool_choice"], + json!({ "type": "function", "function": { "name": "weather" } }) + ); + options.tool_choice = Some(ToolChoice::Required); + assert_eq!( + generate_body("deepseek-reasoner", &options).await["tool_choice"], + "required" + ); +} + +#[tokio::test] +async fn tool_call_should_warn_about_provider_defined_tools() { + use aimux_core::tool::ProviderTool; + let mut options = options(); + options.tools = Some(vec![Tool::Provider(ProviderTool { + id: "test.search".to_string(), + name: "search".to_string(), + args: serde_json::Map::new(), + })]); + let result = generate_result("deepseek-chat", &options).await; + assert_eq!( + warning_values(&result.warnings), + warning_values(&[Warning::Unsupported { + feature: "provider-defined tool test.search".to_string(), + details: None, + }]) + ); + assert!( + result.request.and_then(|r| r.body).unwrap()["tools"] + .as_array() + .unwrap() + .is_empty() + ); +} + #[tokio::test] async fn tool_call_should_extract_tool_call_content() { let server = json_server(tool_call_response()).await; @@ -1269,7 +1314,7 @@ async fn json_response_format_should_send_correct_request_body_with_schema() { "$schema": "http://json-schema.org/draft-07/schema#" }); let result = generate_result("deepseek-reasoner", &json_options(Some(schema.clone()))).await; - let body = result.request.and_then(|request| request.body).unwrap(); + let body = result.request.and_then(|r| r.body).unwrap(); assert_eq!( body["messages"][0], json!({ @@ -1335,7 +1380,7 @@ async fn prefix_should_send_name_and_prefix_on_the_final_assistant_message() { .await .unwrap() .request - .and_then(|request| request.body) + .and_then(|r| r.body) .unwrap(); assert_eq!( body["messages"], @@ -1371,17 +1416,14 @@ async fn prefix_should_reject_prefix_completion_with_the_default_base_url() { ); } -// =========================================================================== -// describe('doStream') -// =========================================================================== - +// Upstream emits provider error envelopes as error parts, including the first event. async fn stream_error(data: &Value) -> aimux_core::error::ApiCallError { let server = sse_server(vec![data_event(data)]).await; - let stream = chat(&server, "deepseek-chat") + let result = chat(&server, "deepseek-chat") .do_stream(&options()) .await .unwrap(); - collect(stream) + collect(result) .await .into_iter() .find_map(|part| match part { @@ -1390,7 +1432,7 @@ async fn stream_error(data: &Value) -> aimux_core::error::ApiCallError { } => Some(*error), _ => None, }) - .expect("an API call error part") + .expect("expected a provider error part") } #[tokio::test] @@ -1430,8 +1472,6 @@ async fn stream_should_preserve_the_provider_type_when_code_is_an_http_status() assert!(error.is_retryable); } -// ---- describe('text') ------------------------------------------------------------ - async fn stream_parts( model_id: &str, options: &CallOptions, @@ -1439,11 +1479,7 @@ async fn stream_parts( ) -> (Value, Vec) { let server = sse_server(chunks).await; let result = chat(&server, model_id).do_stream(options).await.unwrap(); - let body = result - .request - .clone() - .and_then(|request| request.body) - .unwrap(); + let body = result.request.clone().and_then(|r| r.body).unwrap(); (body, collect(result).await) } @@ -1922,7 +1958,7 @@ async fn stream_prefix_should_send_prefix_true_on_the_final_assistant_message() .do_stream(&options_for(prefix_prompt(json!({ "prefix": true })))) .await .unwrap(); - let body = result.request.and_then(|request| request.body).unwrap(); + let body = result.request.and_then(|r| r.body).unwrap(); assert_eq!( body["messages"], json!([ @@ -1968,7 +2004,7 @@ async fn convert_should_ignore_a_name_on_a_tool_message_with_an_unsupported_warn value: "sunny".into(), provider_options: None, }, - tool_name: "getWeather".into(), + tool_name: "weather_tool".into(), provider_options: None, })], provider_options: None, @@ -1977,7 +2013,7 @@ async fn convert_should_ignore_a_name_on_a_tool_message_with_an_unsupported_warn )]); let result = generate_result("deepseek-chat", &options).await; assert_eq!( - result.request.and_then(|request| request.body).unwrap()["messages"], + result.request.and_then(|r| r.body).unwrap()["messages"], json!([{ "role": "tool", "tool_call_id": "call-1", "content": "sunny" }]) ); assert_eq!( @@ -2154,7 +2190,7 @@ async fn convert_should_reject_image_urls_longer_than_8192_characters() { &options_for(vec![LanguageModelMessage::User { content: vec![UserPart::File(FilePart { data: FileData::Url { - url, + url: (url), original_url: None, }, media_type: ("image/png").into(), @@ -2273,7 +2309,7 @@ async fn convert_should_warn_about_unsupported_non_image_file_parts() { ) .await; assert_eq!( - result.request.and_then(|request| request.body).unwrap()["messages"], + result.request.and_then(|r| r.body).unwrap()["messages"], json!([{ "role": "user", "content": "Hello" }]) ); assert_eq!( @@ -2333,7 +2369,7 @@ fn wire_tool_call() -> Value { async fn convert_should_stringify_arguments_to_tool_calls() { let result = generate_result("deepseek-chat", &options_for(tool_turn(false))).await; assert_eq!( - result.request.and_then(|request| request.body).unwrap()["messages"], + result.request.and_then(|r| r.body).unwrap()["messages"], json!([ // upstream sends the empty text as `content: ""`, not null. { "role": "assistant", "content": "", "tool_calls": wire_tool_call() }, diff --git a/aimux-providers/tests/google_factory_test.rs b/aimux-providers/tests/google_factory_test.rs index 1d4210d88..422859766 100644 --- a/aimux-providers/tests/google_factory_test.rs +++ b/aimux-providers/tests/google_factory_test.rs @@ -227,12 +227,6 @@ fn assert_generate_result(fixture: &Fixture, result: &GenerateResult) { "{}: response id", fixture.name ); - assert_eq!( - result.response.as_ref().and_then(|r| r.model_id.as_deref()), - recorded["response"]["modelId"].as_str(), - "{}: response model", - fixture.name - ); assert_eq!( result.usage.input_tokens.total, recorded["usage"]["inputTokens"].as_u64().map(|v| v as _), diff --git a/aimux-providers/tests/google_pse_test.rs b/aimux-providers/tests/google_pse_test.rs index 760e6ec73..d54a84d60 100644 --- a/aimux-providers/tests/google_pse_test.rs +++ b/aimux-providers/tests/google_pse_test.rs @@ -224,8 +224,6 @@ async fn status_403_maps_to_provider_error() { ); } -// -- Provider trait ---------------------------------------------------------- - #[test] fn model_id_is_google_pse_search() { let config = GooglePseProviderSettings { diff --git a/aimux-providers/tests/google_upstream_e2e_test.rs b/aimux-providers/tests/google_upstream_e2e_test.rs new file mode 100644 index 000000000..96e84e729 --- /dev/null +++ b/aimux-providers/tests/google_upstream_e2e_test.rs @@ -0,0 +1,404 @@ +//! Small public-factory end-to-end subset of the pinned Google upstream tests. +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use aimux_core::AiMuxError; +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_core::result::GenerateContent; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::{FunctionTool, Tool}; +use aimux_core::types::FinishReasonUnified; +use aimux_core::video_model::{VideoCallOptions, VideoData, VideoModel, VideoOperationStatus}; +use aimux_providers::google::{GoogleProviderSettings, create_google}; +use futures::StreamExt; +use mock_fetch::{Canned, MockFetch}; +use serde_json::{Value, json}; +use std::sync::Arc; + +const BASE: &str = "https://api.example.com/v1beta"; + +fn settings(mock: &Arc) -> GoogleProviderSettings { + GoogleProviderSettings { + base_url: Some(BASE.into()), + api_key: Some("test-api-key".into()), + fetch: Some(mock.transport()), + generate_id: Some(Arc::new(|| "test-id".into())), + ..Default::default() + } +} + +fn options() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +fn basic_request() -> Value { + json!({"contents":[{"role":"user","parts":[{"text":"Hello"}]}],"generationConfig":{}}) +} + +/// TS: "should preserve complete raw usage metadata" (google/src/google-language-model.test.ts) +#[tokio::test] +async fn generate_preserves_usage_and_metadata() { + let usage: Value = serde_json::from_str(r##"{"promptTokenCount": 12, "cachedContentTokenCount": 4, "candidatesTokenCount": 71, "toolUsePromptTokenCount": 65, "thoughtsTokenCount": 89, "totalTokenCount": 237, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12, "nestedSentinel": "prompt"}], "cacheTokensDetails": [{"modality": "TEXT", "tokenCount": 4, "nestedSentinel": "cache"}], "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 71, "nestedSentinel": "candidate"}], "toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 65, "nestedSentinel": "tool"}], "serviceTier": "standard", "topLevelSentinel": "preserve-me"}"##).unwrap(); + let body = json!({"candidates":[{"content":{"parts":[{"text":"Blue."}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":usage}); + let mock = MockFetch::new(vec![Canned::json(&body)]); + let result = create_google(settings(&mock)) + .unwrap() + .chat("gemini-pro") + .do_generate(&options()) + .await + .unwrap(); + assert_eq!(mock.seen().len(), 1); + let sent = &mock.seen()[0]; + assert_eq!(sent.method, "POST"); + assert_eq!( + sent.url, + format!("{BASE}/models/gemini-pro:generateContent") + ); + assert_eq!(sent.headers["x-goog-api-key"], "test-api-key"); + assert_eq!(sent.json_body(), basic_request()); + assert_eq!(result.request.as_ref().unwrap().body, Some(basic_request())); + assert_eq!(result.response.as_ref().unwrap().body, Some(body)); + assert!(matches!(&result.content[..], [GenerateContent::Text { text, .. }] if text == "Blue.")); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(result.usage.input_tokens.total, Some(77)); + assert_eq!(result.usage.input_tokens.no_cache, Some(73)); + assert_eq!(result.usage.input_tokens.cache_read, Some(4)); + assert_eq!(result.usage.output_tokens.total, Some(160)); + assert_eq!(result.usage.output_tokens.text, Some(71)); + assert_eq!(result.usage.output_tokens.reasoning, Some(89)); + assert_eq!(result.usage.raw, usage.as_object().cloned()); + let metadata = result.provider_metadata.unwrap(); + assert_eq!(metadata["google"]["usageMetadata"], usage); + assert_eq!(metadata["google"]["serviceTier"], "standard"); +} + +/// TS: "should stream text deltas" (google/src/google-language-model.test.ts) +#[tokio::test] +async fn stream_text_and_finish_usage() { + let chunks = r##"{"candidates":[{"content":{"parts":[{"text":"There are **3**"}],"role":"model"},"index":0}],"usageMetadata":{"promptTokenCount":9,"candidatesTokenCount":5,"totalTokenCount":199,"promptTokensDetails":[{"modality":"TEXT","tokenCount":9}],"thoughtsTokenCount":185},"modelVersion":"gemini-3-pro-preview","responseId":"bH6LaZW8Fp_3nsEPqtaSwQ4"} +{"candidates":[{"content":{"parts":[{"text":" \"r\"s in strawberry.\n\nst**r**awbe**rr**y"}],"role":"model"},"index":0}],"usageMetadata":{"promptTokenCount":9,"candidatesTokenCount":23,"totalTokenCount":217,"promptTokensDetails":[{"modality":"TEXT","tokenCount":9}],"thoughtsTokenCount":185},"modelVersion":"gemini-3-pro-preview","responseId":"bH6LaZW8Fp_3nsEPqtaSwQ4"} +{"candidates":[{"content":{"parts":[{"text":"","thoughtSignature":"EqsFCqgFAb4+9vvtAF5n87lB4OGDOoTRMOqp35jW65XsYXh6BySMwl9nvrbAvPcl2U0xITaYUyV4CmREEDB1z0ZPpCg7iEwiZcj40Eh1jXoL8Y/BbPqxdgZKvKxdBsJx92y2ML5ytajQHVFQb9ohEMMnjs9uNadLAhDEsOU1nC5tl3FQkx94uaGfWvg61bJT3Y9OxFdo/kbpm4RBngvYhVkBzHKkHBj72T2bUd8J4HPssi7ORC5iPosPRIOyH/CAVHEtMzFYMwb7OhRu+CW8Z9u7gDieME5iJjXtJtLrNGDxgR7XtWfRRyGjsj6uDS+KvjR3SUSWPdn5eeH6w+LXZm1X///Hvhhcx+NHxsuGjF3fGhyzTVAoIzk0lxyB4+/A9I4Xa0o/T4coVDiewMzGZDwmket//ig8x9UC8cyWr/hy1joZWUO7ooJlLncv8gy4Ng+y1JdievZokSFDNWfMMNAQr3kgUwJDucqDp44C1xMtgR3lhJ75IBBnprHCE/ThgvNXujmqNkwAjp5dS4PjVbrw8fqSylfE80tvU0g9dXqg4pEyG+hGIxbANLhsWjAKLqh69hyqvVLg2Ds3wppphf61IfC4VoeLWj85CjBZMf+k85NsUIJQ6+DQS9IPNbM29ZOzpUbHoWKJB6VzNCSJse7Pi07L+pd6skl77km00y4lJdHIGHfEgi8PaOonakBcxbRqKzGJAA/urlP0tiWya2fTWrvNZOybJHyyofNNSI4s5y76yKEjP1wnPqC7ujrQk6xb7eyCeqH9ekByy3vv0JfgERFptoSUoG2toIr9M3lS/LKpnwfCvZh+z3J0iMb83d4MaPKhGhE49J4660XUsEmjygAZNi9HnjfC3KtaU/07Sx4JCezMtpsLKUxBgy4xaNqwew3FwAG37eeWcow="}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":9,"candidatesTokenCount":23,"totalTokenCount":217,"promptTokensDetails":[{"modality":"TEXT","tokenCount":9}],"thoughtsTokenCount":185},"modelVersion":"gemini-3-pro-preview","responseId":"bH6LaZW8Fp_3nsEPqtaSwQ4"}"##; + let mock = MockFetch::new(vec![Canned { + status: 200, + headers: vec![("content-type".into(), "text/event-stream".into())], + body: chunks + .lines() + .map(|line| format!("data: {line}\n\n")) + .collect::() + .into_bytes(), + }]); + let result = create_google(settings(&mock)) + .unwrap() + .chat("gemini-pro") + .do_stream(&options()) + .await + .unwrap(); + assert_eq!(result.request.as_ref().unwrap().body, Some(basic_request())); + assert_eq!( + result.response.as_ref().unwrap().headers.as_ref().unwrap()["content-type"], + "text/event-stream" + ); + let events: Vec<_> = result + .stream + .collect::>() + .await + .into_iter() + .map(Result::unwrap) + .collect(); + assert_eq!(mock.seen().len(), 1); + let sent = &mock.seen()[0]; + assert_eq!( + sent.url, + format!("{BASE}/models/gemini-pro:streamGenerateContent?alt=sse") + ); + assert_eq!(sent.json_body(), basic_request()); + let deltas: Vec<_> = events + .iter() + .filter_map(|event| match event { + StreamPart::TextDelta { delta, .. } => Some(delta.as_str()), + _ => None, + }) + .collect(); + assert_eq!( + deltas, + vec![ + "There are **3**", + " \"r\"s in strawberry.\n\nst**r**awbe**rr**y", + "" + ] + ); + let last_chunk: Value = serde_json::from_str(chunks.lines().last().unwrap()).unwrap(); + assert!(events.iter().any(|event| matches!(event, + StreamPart::TextDelta { delta, provider_metadata: Some(metadata), .. } + if delta.is_empty() && metadata["google"]["thoughtSignature"] == last_chunk["candidates"][0]["content"]["parts"][0]["thoughtSignature"] + ))); + let (finish, usage, metadata) = events + .iter() + .find_map(|event| match event { + StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } => Some((finish_reason, usage, provider_metadata.as_ref().unwrap())), + _ => None, + }) + .unwrap(); + assert_eq!(finish.unified, FinishReasonUnified::Stop); + assert_eq!(usage.input_tokens.total, Some(9)); + assert_eq!(usage.output_tokens.total, Some(208)); + assert_eq!(usage.output_tokens.text, Some(23)); + assert_eq!(usage.output_tokens.reasoning, Some(185)); + assert_eq!(metadata["google"]["usageMetadata"]["totalTokenCount"], 217); + assert_eq!( + events + .iter() + .filter(|event| matches!(event, StreamPart::TextStart { .. })) + .count(), + 1 + ); + assert_eq!( + events + .iter() + .filter(|event| matches!(event, StreamPart::TextEnd { .. })) + .count(), + 1 + ); +} + +/// TS: "should extract tool calls" (google/src/google-language-model.test.ts) +#[tokio::test] +async fn generate_function_call_with_signature() { + let body: Value = serde_json::from_str(r##"{ + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "weather", + "args": { + "location": "San Francisco" + } + }, + "thoughtSignature": "EskgCsYgAb4+9vtF7/499YQS2bjZs3xcQI+iAl+ILn29nK1j0Kg6su7QsUUUk3nrAAfnS2w5WiVvlcCqu9fAebJ2cvfaEyBahEt5" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s)." + } + ], + "usageMetadata": { + "promptTokenCount": 29, + "candidatesTokenCount": 15, + "totalTokenCount": 937, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 29 + } + ], + "thoughtsTokenCount": 893 + }, + "modelVersion": "gemini-3-pro-preview", + "responseId": "m36LaZGyCLz1xs0PtNSB-QU" +} +"##).unwrap(); + let mock = MockFetch::new(vec![Canned::json(&body)]); + let schema = json!({"type":"object","properties":{"location":{"type":"string"}},"required":["location"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}); + let mut opts = options(); + opts.tools = Some(vec![Tool::Function(FunctionTool::new( + "weather", + schema.clone(), + ))]); + let result = create_google(settings(&mock)) + .unwrap() + .chat("gemini-pro") + .do_generate(&opts) + .await + .unwrap(); + let mut expected = basic_request(); + expected["tools"] = json!([{"functionDeclarations":[{"name":"weather","description":"","parametersJsonSchema":schema}]}]); + assert_eq!(mock.seen()[0].json_body(), expected); + let [GenerateContent::ToolCall(call)] = &result.content[..] else { + panic!("expected one tool call") + }; + assert_eq!(call.tool_call_id, "test-id"); + assert_eq!(call.tool_name, "weather"); + assert_eq!(call.input, r#"{"location":"San Francisco"}"#); + assert_eq!( + call.provider_metadata.as_ref().unwrap()["google"]["thoughtSignature"], + body["candidates"][0]["content"]["parts"][0]["thoughtSignature"] + ); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); + assert_eq!(result.usage.input_tokens.total, Some(29)); + assert_eq!(result.usage.output_tokens.total, Some(908)); +} + +/// TS: "preserves google.rpc.RetryInfo in APICallError.data" (google/src/google-error.test.ts) +#[tokio::test] +async fn error_preserves_retry_info() { + let body = json!({"error":{"code":429,"message":"You exceeded your current quota, please check your plan.","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.QuotaFailure","violations":[{"quotaId":"GenerateRequestsPerMinutePerProjectPerModel-FreeTier"}]},{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"34.4s"}]}}); + let mut canned = Canned::json(&body); + canned.status = 429; + let mock = MockFetch::new(vec![canned]); + let error = create_google(settings(&mock)) + .unwrap() + .chat("gemini-2.5-flash") + .do_generate(&options()) + .await + .unwrap_err(); + assert_eq!(mock.seen().len(), 1); + assert_eq!(mock.seen()[0].json_body(), basic_request()); + let AiMuxError::ApiCall(error) = error else { + panic!("expected API call error") + }; + assert_eq!(error.status_code, Some(429)); + assert_eq!( + error.message, + "You exceeded your current quota, please check your plan." + ); + assert_eq!(error.data, Some(body)); +} + +/// TS: "should extract embedding" (google/src/google-embedding-model.test.ts) +/// TS: "should pass the model and the values" (google/src/google-embedding-model.test.ts) +#[tokio::test] +async fn embed_batch_request_and_vectors() { + let body = + json!({"embeddings":[{"values":[0.1,0.2,0.3,0.4,0.5]},{"values":[0.6,0.7,0.8,0.9,1.0]}]}); + let mock = MockFetch::new(vec![Canned::json(&body)]); + let mut opts = EmbeddingCallOptions::new("sunny day at the beach"); + opts.values.push("rainy day in the city".into()); + let result = create_google(settings(&mock)) + .unwrap() + .embedding("gemini-embedding-001") + .do_embed(&opts) + .await + .unwrap(); + let sent = &mock.seen()[0]; + assert_eq!( + sent.url, + format!("{BASE}/models/gemini-embedding-001:batchEmbedContents") + ); + assert_eq!( + sent.json_body(), + json!({"requests":[{"model":"models/gemini-embedding-001","content":{"role":"user","parts":[{"text":"sunny day at the beach"}]}},{"model":"models/gemini-embedding-001","content":{"role":"user","parts":[{"text":"rainy day in the city"}]}}]}) + ); + assert_eq!( + result.embeddings, + vec![vec![0.1, 0.2, 0.3, 0.4, 0.5], vec![0.6, 0.7, 0.8, 0.9, 1.0]] + ); + assert!(result.usage.is_none()); + assert_eq!(result.response.unwrap().body, Some(body)); +} + +/// TS: "should use the language model endpoint and extract generated images" (google/src/google-image-model.test.ts) +#[tokio::test] +async fn image_request_and_metadata() { + let usage = json!({"promptTokenCount":10,"candidatesTokenCount":100,"totalTokenCount":110}); + let body = json!({"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"base64-generated-image"}}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":usage}); + let mock = MockFetch::new(vec![Canned::json(&body)]); + let result = create_google(settings(&mock)) + .unwrap() + .image("gemini-2.5-flash-image") + .do_generate(&ImageCallOptions::new("A beautiful sunset")) + .await + .unwrap(); + let sent = &mock.seen()[0]; + assert_eq!( + sent.url, + format!("{BASE}/models/gemini-2.5-flash-image:generateContent") + ); + assert_eq!( + sent.json_body(), + json!({"contents":[{"role":"user","parts":[{"text":"A beautiful sunset"}]}],"generationConfig":{"responseModalities":["IMAGE"]}}) + ); + let ImageOutputs::Base64(images) = result.images else { + panic!("expected base64 images") + }; + assert_eq!(images, vec!["base64-generated-image"]); + assert_eq!( + serde_json::to_value(result.provider_metadata.unwrap()).unwrap(), + json!({"google":{"finishMessage":null,"finishReason":"STOP","groundingMetadata":null,"images":[{}],"promptFeedback":null,"safetyRatings":null,"serviceTier":null,"urlContextMetadata":null,"usageMetadata":usage}}) + ); + let usage = result.usage.unwrap(); + assert_eq!(usage.input_tokens, Some(10)); + assert_eq!(usage.output_tokens, Some(100)); + assert_eq!(usage.total_tokens, Some(110)); +} + +/// TS: "should return operation with operationName" (google/src/google-video-model.test.ts) +/// TS: "should pass correct request body" (google/src/google-video-model.test.ts) +/// TS: "should return completed with video data when done" (google/src/google-video-model.test.ts) +#[tokio::test] +async fn video_start_and_completed_status() { + let mock = MockFetch::new(vec![ + Canned::json(&json!({"name":"operations/start-test-op","done":false})), + Canned::json( + &json!({"done":true,"response":{"generateVideoResponse":{"generatedSamples":[{"video":{"uri":"https://api.example.com/files/video-456.mp4"}}]}}}), + ), + ]); + let model = create_google(settings(&mock)) + .unwrap() + .video("veo-3.1-generate-preview"); + let opts = VideoCallOptions::new("A futuristic city with flying cars"); + let start = model.do_start(&opts).await.unwrap(); + assert_eq!( + start.operation, + json!({"operationName":"operations/start-test-op"}) + ); + let sent = &mock.seen()[0]; + assert_eq!(sent.method, "POST"); + assert_eq!( + sent.url, + format!("{BASE}/models/veo-3.1-generate-preview:predictLongRunning") + ); + assert_eq!( + sent.json_body(), + json!({"instances":[{"prompt":"A futuristic city with flying cars"}],"parameters":{"sampleCount":1}}) + ); + let status = model.do_status(&start.operation, &opts).await.unwrap(); + assert_eq!(mock.seen().len(), 2); + assert_eq!(mock.seen()[1].method, "GET"); + assert_eq!( + mock.seen()[1].url, + format!("{BASE}/operations/start-test-op") + ); + let VideoOperationStatus::Completed(result) = status else { + panic!("expected completed video") + }; + assert_eq!(result.videos.len(), 1); + assert!( + matches!(&result.videos[0], VideoData::Url { url, media_type } if url == "https://api.example.com/files/video-456.mp4?key=test-api-key" && media_type == "video/mp4") + ); +} + +/// TS: "should throw when no videos in response" (google/src/google-video-model.test.ts) +#[tokio::test] +async fn video_completed_response_without_videos_errors() { + let mock = MockFetch::new(vec![Canned::json( + &json!({"done":true,"response":{"generateVideoResponse":{"generatedSamples":[]}}}), + )]); + let error = create_google(settings(&mock)) + .unwrap() + .video("veo-3.1-generate-preview") + .do_status( + &json!({"operationName":"operations/status-test-op"}), + &VideoCallOptions::new("A futuristic city with flying cars"), + ) + .await + .unwrap_err(); + assert_eq!(mock.seen()[0].method, "GET"); + assert_eq!( + mock.seen()[0].url, + format!("{BASE}/operations/status-test-op") + ); + assert!(error.to_string().contains("No videos in response")); +} diff --git a/aimux-providers/tests/groq_test.rs b/aimux-providers/tests/groq_test.rs index ededb90a4..98b65fad9 100644 --- a/aimux-providers/tests/groq_test.rs +++ b/aimux-providers/tests/groq_test.rs @@ -1261,10 +1261,10 @@ mod do_generate { let model = make_provider(&server); - let provider_opts = provider_namespace("groq", json!({"reasoningEffort": "high"})); + let provider_opts = provider_namespace("groq", json!({"reasoningEffort": "high"})).unwrap(); let options = CallOptions { reasoning: Some(ReasoningEffort::Medium), - provider_options: Some(provider_opts.unwrap()), + provider_options: Some(provider_opts), ..default_options(test_prompt()) }; model.do_generate(&options).await.unwrap(); @@ -1426,9 +1426,10 @@ mod do_generate { "user": "test-user-id", "parallelToolCalls": false }), - ); + ) + .unwrap(); let options = CallOptions { - provider_options: Some(provider_opts.unwrap()), + provider_options: Some(provider_opts), ..default_options(test_prompt()) }; model.do_generate(&options).await.unwrap(); @@ -1447,9 +1448,9 @@ mod do_generate { let model = make_provider(&server); - let provider_opts = provider_namespace("groq", json!({"serviceTier": "flex"})); + let provider_opts = provider_namespace("groq", json!({"serviceTier": "flex"})).unwrap(); let options = CallOptions { - provider_options: Some(provider_opts.unwrap()), + provider_options: Some(provider_opts), ..default_options(test_prompt()) }; model.do_generate(&options).await.unwrap(); @@ -1466,9 +1467,10 @@ mod do_generate { let model = make_provider(&server); - let provider_opts = provider_namespace("groq", json!({"serviceTier": "performance"})); + let provider_opts = + provider_namespace("groq", json!({"serviceTier": "performance"})).unwrap(); let options = CallOptions { - provider_options: Some(provider_opts.unwrap()), + provider_options: Some(provider_opts), ..default_options(test_prompt()) }; model.do_generate(&options).await.unwrap(); @@ -1550,9 +1552,10 @@ mod do_generate { let model = make_provider(&server); - let provider_opts = provider_namespace("groq", json!({"structuredOutputs": false})); + let provider_opts = + provider_namespace("groq", json!({"structuredOutputs": false})).unwrap(); let options = CallOptions { - provider_options: Some(provider_opts.unwrap()), + provider_options: Some(provider_opts), response_format: Some(ResponseFormat::Json { schema: Some(json!({ "type": "object", @@ -1585,9 +1588,9 @@ mod do_generate { let model = make_provider(&server); - let provider_opts = provider_namespace("groq", json!({"strictJsonSchema": false})); + let provider_opts = provider_namespace("groq", json!({"strictJsonSchema": false})).unwrap(); let options = CallOptions { - provider_options: Some(provider_opts.unwrap()), + provider_options: Some(provider_opts), response_format: Some(ResponseFormat::Json { schema: Some(json!({ "type": "object", @@ -1814,17 +1817,44 @@ mod do_stream { async fn handles_error_stream() { let server = MockServer::start().await; let body = sse_body(&[&sse_event( - r#"{"error":{"message":"The server had an error processing your request. Sorry about that!","type":"invalid_request_error"}}"#, + r#"{"error":{"message":"Rate limit reached","type":"rate_limit_error"}}"#, )]); mock_sse(&server, body).await; let model = make_provider(&server); - let result = model.do_stream(&default_options(test_prompt())).await; - // upstream emits an error part and finishes; this crate's convention - // (as in the Mistral, Cohere and OpenAI-compatible models) is to reject - // when the very first event is an error, inside Core's retry boundary. - assert!(result.is_err()); + let result = model + .do_stream(&default_options(test_prompt())) + .await + .unwrap(); + let parts = collect_stream(result).await; + assert_eq!(parts.len(), 3); + assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); + match &parts[1] { + StreamPart::Error { + error: aimux_core::AiMuxError::ApiCall(error), + } => { + assert_eq!(error.message, "Rate limit reached"); + assert_eq!(error.status_code, Some(429)); + assert!(error.is_retryable); + } + other => panic!("expected Error, got {other:?}"), + } + match &parts[2] { + StreamPart::Finish { + finish_reason, + usage, + .. + } => { + assert_eq!(finish_reason.unified, FinishReasonUnified::Error); + assert_eq!(finish_reason.raw, None); + assert_eq!( + serde_json::to_value(usage).unwrap(), + serde_json::to_value(aimux_core::types::Usage::default()).unwrap() + ); + } + other => panic!("expected Finish, got {other:?}"), + } } /// upstream: "should handle error stream parts" (an error chunk after the @@ -2189,7 +2219,7 @@ mod package { .await .unwrap(); - let body = result.request.and_then(|request| request.body).unwrap(); + let body = result.request.and_then(|r| r.body).unwrap(); assert_eq!(body["user"], "u1"); assert_eq!(body["reasoning_format"], "parsed"); assert_eq!(body["service_tier"], "flex"); @@ -2220,7 +2250,7 @@ mod package { result .request .as_ref() - .and_then(|request| request.body.as_ref()) + .and_then(|r| r.body.as_ref()) .unwrap() .get("stream_options") .is_none() @@ -2255,7 +2285,7 @@ mod package { options.max_output_tokens = Some(64); options.top_k = Some(40.0); let result = model.do_generate(&options).await.unwrap(); - let body = result.request.and_then(|request| request.body).unwrap(); + let body = result.request.and_then(|r| r.body).unwrap(); // upstream: getArgs sends `max_tokens: maxOutputTokens` assert_eq!(body["max_tokens"], 64); assert!(body.get("max_completion_tokens").is_none()); diff --git a/aimux-providers/tests/huggingface_responses_test.rs b/aimux-providers/tests/huggingface_responses_test.rs index 15ff20404..648194dae 100644 --- a/aimux-providers/tests/huggingface_responses_test.rs +++ b/aimux-providers/tests/huggingface_responses_test.rs @@ -16,8 +16,8 @@ //! overridden with the mock server's root URI (no `/v1` suffix), so the //! resulting request path is `/responses`. -use aimux_core::tool::RawToolCall; use aimux_core::tool::ToolResult; +use aimux_core::tool::{FunctionTool, RawToolCall}; use std::collections::HashMap; use futures::StreamExt; @@ -36,7 +36,6 @@ use aimux_core::options::{CallOptions, ResponseFormat, Tool, ToolChoice}; use aimux_core::result::{GenerateContent, ReasoningOutput, Source}; use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; -use aimux_core::tool::FunctionTool; use aimux_core::types::{FinishReasonUnified, ResponseMetadata}; use aimux_providers::huggingface::responses::convert_to_huggingface_responses_messages; @@ -1483,7 +1482,7 @@ async fn should_handle_reasoning_content_in_responses() { "huggingface", json!({ "itemId": "reasoning_1" }) ) - .expect("provider metadata object") + .unwrap() ) ); } @@ -1574,7 +1573,7 @@ async fn should_stream_reasoning_content() { "huggingface", json!({ "itemId": "reasoning_stream" }) ) - .expect("provider metadata object") + .unwrap() ) ); } @@ -1906,8 +1905,8 @@ fn passes_through_url_source_for_top_level_only_image() { let prompt = vec![LanguageModelMessage::User { content: vec![UserPart::File(FilePart { data: FileData::Url { - url: "https://example.com/x.png".into(), original_url: None, + url: "https://example.com/x.png".into(), }, media_type: "image".into(), filename: None, diff --git a/aimux-providers/tests/jina_ai_test.rs b/aimux-providers/tests/jina_ai_test.rs index 46270f6af..89ca0e666 100644 --- a/aimux-providers/tests/jina_ai_test.rs +++ b/aimux-providers/tests/jina_ai_test.rs @@ -278,5 +278,3 @@ async fn status_401_maps_to_auth_error() { "expected Auth error, got {result:?}" ); } - -// -- Provider trait ---------------------------------------------------------- diff --git a/aimux-providers/tests/linkup_test.rs b/aimux-providers/tests/linkup_test.rs index 6b71e5f07..a89d6daa8 100644 --- a/aimux-providers/tests/linkup_test.rs +++ b/aimux-providers/tests/linkup_test.rs @@ -233,8 +233,6 @@ async fn status_403_maps_to_provider_error() { ); } -// -- Provider trait ---------------------------------------------------------- - #[test] fn model_id_is_linkup_search() { let config = LinkupProviderSettings { diff --git a/aimux-providers/tests/mistral_upstream_e2e_test.rs b/aimux-providers/tests/mistral_upstream_e2e_test.rs new file mode 100644 index 000000000..768f8e82d --- /dev/null +++ b/aimux-providers/tests/mistral_upstream_e2e_test.rs @@ -0,0 +1,253 @@ +//! Small end-to-end subset of the pinned Mistral upstream tests. + +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_core::result::GenerateContent; +use aimux_core::stream_part::StreamPart; +use aimux_core::types::FinishReasonUnified; +use aimux_providers::mistral::{MistralProvider, MistralProviderSettings, create_mistral}; +use futures::StreamExt; +use mock_fetch::{Canned, MockFetch}; +use serde_json::{Value, json}; + +fn provider(fetch: &std::sync::Arc) -> MistralProvider { + create_mistral(MistralProviderSettings { + api_key: Some("test-api-key".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} + +fn options() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +fn expected_request(stream: bool) -> Value { + let mut body = json!({ + "model": "mistral-small-latest", + "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + }); + if stream { + body["stream"] = json!(true); + } + body +} + +fn assert_chat_request(fetch: &MockFetch, stream: bool) { + let requests = fetch.seen(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].method, "POST"); + assert_eq!( + requests[0].url, + "https://api.mistral.ai/v1/chat/completions" + ); + assert_eq!(requests[0].headers["authorization"], "Bearer test-api-key"); + assert_eq!(requests[0].json_body(), expected_request(stream)); +} +fn usage_details_response() -> Value { + serde_json::from_str(r###"{"id":"usage-details","created":1787866434,"model":"mistral-small-latest","usage":{"prompt_tokens":20,"completion_tokens":2,"total_tokens":22,"prompt_audio_seconds":1,"request_count":1,"service_tier":"standard","num_cached_tokens":0,"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":1,"messages":[{"role":"user","total_tokens":20,"settings_tokens":null,"truncated":false,"usage_count":1}],"additional_prompt_detail":{"value":true}},"prompt_token_details":{"cached_tokens":0,"audio_tokens":1,"additional_prompt_token_detail":["value"]},"completion_tokens_details":{"reasoning_tokens":7,"additional_completion_detail":"value"},"additional_usage_field":{"nested":true}},"object":"chat.completion","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","tool_calls":null,"content":"OK"}}]}"###).unwrap() +} + +fn tool_call_response() -> Value { + serde_json::from_str(r###"{"id":"b3999b8c93e04e11bcbff7bcab829667","created":1769088854,"model":"mistral-small-latest","usage":{"prompt_tokens":124,"total_tokens":146,"completion_tokens":22},"object":"chat.completion","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","tool_calls":[{"id":"gSIMJiOkT","function":{"name":"weather","arguments":"{\"location\": \"San Francisco\"}"}}]}}]}"###).unwrap() +} + +fn text_chunks() -> Value { + serde_json::from_str(r###"[{"id":"5319bd0299614c679a0068a4f2c8ffd0","object":"chat.completion.chunk","created":1769088720,"model":"mistral-small-latest","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null,"logprobs":null}]},{"id":"5319bd0299614c679a0068a4f2c8ffd0","object":"chat.completion.chunk","created":1769088720,"model":"mistral-small-latest","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null,"logprobs":null}]},{"id":"5319bd0299614c679a0068a4f2c8ffd0","object":"chat.completion.chunk","created":1769088720,"model":"mistral-small-latest","choices":[{"index":0,"delta":{"content":", "},"finish_reason":null,"logprobs":null}]},{"id":"5319bd0299614c679a0068a4f2c8ffd0","object":"chat.completion.chunk","created":1769088720,"model":"mistral-small-latest","choices":[{"index":0,"delta":{"content":"world!"},"finish_reason":null,"logprobs":null}]},{"id":"5319bd0299614c679a0068a4f2c8ffd0","object":"chat.completion.chunk","created":1769088720,"model":"mistral-small-latest","choices":[{"index":0,"delta":{"content":" This"},"finish_reason":null,"logprobs":null}]},{"id":"5319bd0299614c679a0068a4f2c8ffd0","object":"chat.completion.chunk","created":1769088720,"model":"mistral-small-latest","choices":[{"index":0,"delta":{"content":" is a test"},"finish_reason":null,"logprobs":null}]},{"id":"5319bd0299614c679a0068a4f2c8ffd0","object":"chat.completion.chunk","created":1769088720,"model":"mistral-small-latest","choices":[{"index":0,"delta":{"content":" response."},"finish_reason":null,"logprobs":null}]},{"id":"5319bd0299614c679a0068a4f2c8ffd0","object":"chat.completion.chunk","created":1769088720,"model":"mistral-small-latest","choices":[{"index":0,"delta":{"content":""},"finish_reason":"stop","logprobs":null}],"usage":{"prompt_tokens":13,"total_tokens":21,"completion_tokens":8}}]"###).unwrap() +} + +/// TS: "should preserve the complete raw usage object without changing normalized usage" (mistral/src/mistral-chat-language-model.test.ts) +#[tokio::test] +async fn generate_preserves_complete_usage() { + let body = usage_details_response(); + let fetch = MockFetch::new(vec![Canned::json(&body)]); + let result = provider(&fetch) + .chat("mistral-small-latest") + .do_generate(&options()) + .await + .unwrap(); + assert_chat_request(&fetch, false); + assert_eq!(result.request.unwrap().body, Some(expected_request(false))); + assert_eq!( + result.content, + vec![GenerateContent::Text { + text: "OK".into(), + provider_metadata: None + }] + ); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(result.finish_reason.raw.as_deref(), Some("stop")); + assert_eq!( + serde_json::to_value(result.usage).unwrap(), + json!({ + "input_tokens": {"total": 20, "no_cache": 20}, + "output_tokens": {"total": 2, "text": 2}, + "raw": body["usage"] + }) + ); + assert!(result.warnings.is_empty()); + assert!(result.provider_metadata.is_none()); + let response = result.response.unwrap(); + assert_eq!(response.id.as_deref(), Some("usage-details")); + assert_eq!(response.model_id.as_deref(), Some("mistral-small-latest")); + assert_eq!(response.body, Some(body)); +} + +/// TS: "should stream text" (mistral/src/mistral-chat-language-model.test.ts) +#[tokio::test] +async fn stream_text_and_usage() { + let chunks = text_chunks(); + let mut body: String = chunks + .as_array() + .unwrap() + .iter() + .map(|chunk| format!("data: {chunk}\n\n")) + .collect(); + body.push_str("data: [DONE]\n\n"); + let fetch = MockFetch::new(vec![Canned { + status: 200, + headers: vec![("content-type".into(), "text/event-stream".into())], + body: body.into_bytes(), + }]); + let result = provider(&fetch) + .chat("mistral-small-latest") + .do_stream(&options()) + .await + .unwrap(); + assert_chat_request(&fetch, true); + assert_eq!(result.request.unwrap().body, Some(expected_request(true))); + let parts = result + .stream + .collect::>() + .await + .into_iter() + .collect::, _>>() + .unwrap(); + assert_eq!(parts.len(), 11); + assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); + match &parts[1] { + StreamPart::ResponseMetadata(metadata) => { + assert_eq!( + metadata.id.as_deref(), + Some("5319bd0299614c679a0068a4f2c8ffd0") + ); + assert_eq!(metadata.model_id.as_deref(), Some("mistral-small-latest")); + assert_eq!( + metadata.timestamp.as_deref(), + Some("2026-01-22T13:32:00+00:00") + ); + } + part => panic!("unexpected metadata: {part:?}"), + } + assert!( + matches!(&parts[2], StreamPart::TextStart { id, provider_metadata } if id == "0" && provider_metadata.is_none()) + ); + for (part, expected) in + parts[3..9] + .iter() + .zip(["Hello", ", ", "world!", " This", " is a test", " response."]) + { + assert!( + matches!(part, StreamPart::TextDelta { id, delta, provider_metadata } if id == "0" && delta == expected && provider_metadata.is_none()) + ); + } + assert!( + matches!(&parts[9], StreamPart::TextEnd { id, provider_metadata } if id == "0" && provider_metadata.is_none()) + ); + match &parts[10] { + StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } => { + assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(finish_reason.raw.as_deref(), Some("stop")); + assert_eq!( + serde_json::to_value(usage).unwrap(), + json!({ + "input_tokens": {"total": 13, "no_cache": 13}, + "output_tokens": {"total": 8, "text": 8}, + "raw": {"prompt_tokens": 13, "completion_tokens": 8, "total_tokens": 21} + }) + ); + assert!(provider_metadata.is_none()); + } + part => panic!("unexpected finish: {part:?}"), + } +} + +/// TS: "should extract tool call content" (mistral/src/mistral-chat-language-model.test.ts) +#[tokio::test] +async fn generate_tool_call_keeps_argument_text() { + let body = tool_call_response(); + let fetch = MockFetch::new(vec![Canned::json(&body)]); + let result = provider(&fetch) + .chat("mistral-small-latest") + .do_generate(&options()) + .await + .unwrap(); + assert_chat_request(&fetch, false); + assert_eq!(result.content.len(), 1); + match &result.content[0] { + GenerateContent::ToolCall(call) => { + assert_eq!(call.tool_call_id, "gSIMJiOkT"); + assert_eq!(call.tool_name, "weather"); + assert_eq!(call.input, "{\"location\": \"San Francisco\"}"); + assert!(call.provider_executed.is_none()); + assert!(call.dynamic.is_none()); + assert!(call.provider_metadata.is_none()); + } + part => panic!("unexpected tool content: {part:?}"), + } + assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); + assert_eq!(result.finish_reason.raw.as_deref(), Some("tool_calls")); + assert_eq!(result.usage.input_tokens.total, Some(124)); + assert_eq!(result.usage.output_tokens.total, Some(22)); + assert_eq!(result.usage.raw, body["usage"].as_object().cloned()); + assert!(result.warnings.is_empty()); + assert!(result.provider_metadata.is_none()); +} + +/// TS: "should extract embedding" (mistral/src/mistral-embedding-model.test.ts) +#[tokio::test] +async fn embed_request_vectors_and_usage() { + let body = json!({ + "id": "b322cfc2b9d34e2f8e14fc99874faee5", "object": "list", + "data": [ + {"object": "embedding", "embedding": [0.1, 0.2, 0.3, 0.4, 0.5], "index": 0}, + {"object": "embedding", "embedding": [0.6, 0.7, 0.8, 0.9, 1], "index": 1} + ], + "model": "mistral-embed", "usage": {"prompt_tokens": 8, "total_tokens": 8} + }); + let fetch = MockFetch::new(vec![Canned::json(&body)]); + let mut options = EmbeddingCallOptions::new("sunny day at the beach"); + options.values.push("rainy day in the city".into()); + let result = provider(&fetch) + .embedding("mistral-embed") + .do_embed(&options) + .await + .unwrap(); + let requests = fetch.seen(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].method, "POST"); + assert_eq!(requests[0].url, "https://api.mistral.ai/v1/embeddings"); + assert_eq!(requests[0].headers["authorization"], "Bearer test-api-key"); + assert_eq!( + requests[0].json_body(), + json!({"model": "mistral-embed", "input": options.values, "encoding_format": "float"}) + ); + assert_eq!( + result.embeddings, + vec![vec![0.1, 0.2, 0.3, 0.4, 0.5], vec![0.6, 0.7, 0.8, 0.9, 1.0]] + ); + assert_eq!(result.usage.unwrap().tokens, 8); + assert_eq!(result.response.unwrap().body, Some(body)); + assert!(result.provider_metadata.is_none()); + assert!(result.warnings.is_empty()); +} diff --git a/aimux-providers/tests/moonshotai_test.rs b/aimux-providers/tests/moonshotai_test.rs index 0895f28be..3febe6505 100644 --- a/aimux-providers/tests/moonshotai_test.rs +++ b/aimux-providers/tests/moonshotai_test.rs @@ -302,20 +302,6 @@ async fn converts_basic_usage() { assert_eq!(usage.output_tokens.reasoning, Some(0)); } -/// TS: "should convert usage with top-level cached_tokens (Moonshot format)". -#[tokio::test] -async fn converts_top_level_cached_tokens() { - let usage = usage_for(json!({ - "prompt_tokens": 100, "completion_tokens": 50, "cached_tokens": 30 - })) - .await; - assert_eq!(usage.input_tokens.total, Some(100)); - assert_eq!(usage.input_tokens.cache_read, Some(30)); - assert_eq!(usage.input_tokens.no_cache, Some(70)); - assert_eq!(usage.output_tokens.total, Some(50)); - assert_eq!(usage.output_tokens.text, Some(50)); -} - /// TS: "should convert usage with nested cached_tokens (OpenAI format)". #[tokio::test] async fn converts_nested_cached_tokens() { @@ -329,20 +315,6 @@ async fn converts_nested_cached_tokens() { assert_eq!(usage.input_tokens.no_cache, Some(75)); } -/// TS: "should prioritize top-level cached_tokens over nested". -#[tokio::test] -async fn prioritizes_top_level_cached_tokens() { - let usage = usage_for(json!({ - "prompt_tokens": 100, - "completion_tokens": 50, - "cached_tokens": 40, - "prompt_tokens_details": { "cached_tokens": 25 } - })) - .await; - assert_eq!(usage.input_tokens.cache_read, Some(40)); - assert_eq!(usage.input_tokens.no_cache, Some(60)); -} - /// TS: "should convert usage with reasoning tokens". #[tokio::test] async fn converts_reasoning_tokens() { @@ -357,22 +329,6 @@ async fn converts_reasoning_tokens() { assert_eq!(usage.output_tokens.text, Some(50)); } -/// TS: "should convert usage with both cached and reasoning tokens". -#[tokio::test] -async fn converts_cached_and_reasoning_tokens() { - let usage = usage_for(json!({ - "prompt_tokens": 100, - "completion_tokens": 80, - "cached_tokens": 35, - "completion_tokens_details": { "reasoning_tokens": 30 } - })) - .await; - assert_eq!(usage.input_tokens.cache_read, Some(35)); - assert_eq!(usage.input_tokens.no_cache, Some(65)); - assert_eq!(usage.output_tokens.reasoning, Some(30)); - assert_eq!(usage.output_tokens.text, Some(50)); -} - /// TS: "should handle null values in usage fields". #[tokio::test] async fn handles_null_usage_field_values() { diff --git a/aimux-providers/tests/open_responses_test.rs b/aimux-providers/tests/open_responses_test.rs index 6ffb9a3f4..2bc6c929e 100644 --- a/aimux-providers/tests/open_responses_test.rs +++ b/aimux-providers/tests/open_responses_test.rs @@ -19,7 +19,7 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{ AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolCallPart, - ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, + ToolPart, ToolResultPart, UserPart, }; use aimux_core::options::{CallOptions, ResponseFormat, Tool, ToolChoice}; use aimux_core::result::{GenerateContent, ReasoningOutput}; @@ -205,7 +205,7 @@ fn openai_pdf_json() -> Value { "created_at": 1775566496, "status": "completed", "incomplete_details": null, - "model": "test-pdf-model", + "model": lmstudio_basic_json()["model"], "output": [ { "id": "msg_048edf44633e41ae0069d4fea0d1a08194af1e491c093df1d9", @@ -446,8 +446,8 @@ mod convert_tests { let prompt = vec![LanguageModelMessage::User { content: vec![UserPart::File(FilePart { data: FileData::Url { - url: "https://example.com/image.png".into(), original_url: None, + url: "https://example.com/image.png".into(), }, media_type: "image/png".into(), filename: None, @@ -505,8 +505,8 @@ mod convert_tests { let prompt = vec![LanguageModelMessage::User { content: vec![UserPart::File(FilePart { data: FileData::Url { - url: "https://example.com/document.pdf".into(), original_url: None, + url: "https://example.com/document.pdf".into(), }, media_type: "application/pdf".into(), filename: None, @@ -740,12 +740,11 @@ mod convert_tests { let prompt = vec![LanguageModelMessage::Tool { content: vec![ToolPart::ToolResult(ToolResultPart { tool_call_id: "call_123".into(), - output: ToolResultOutput::Json { - value: json!({"temperature": 72, "condition": "sunny"}), - provider_options: None, - }, - tool_name: "test".into(), - + output: serde_json::from_value( + json!({"type": "json", "value": {"temperature": 72, "condition": "sunny"}}), + ) + .unwrap(), + tool_name: "get_weather".into(), provider_options: None, })], provider_options: None, @@ -766,12 +765,11 @@ mod convert_tests { let prompt = vec![LanguageModelMessage::Tool { content: vec![ToolPart::ToolResult(ToolResultPart { tool_call_id: "call_456".into(), - output: ToolResultOutput::Text { - value: "Search results: Found 5 items".into(), - provider_options: None, - }, - tool_name: "test".into(), - + output: serde_json::from_value( + json!({"type": "text", "value": "Search results: Found 5 items"}), + ) + .unwrap(), + tool_name: "get_weather".into(), provider_options: None, })], provider_options: None, @@ -792,12 +790,11 @@ mod convert_tests { let prompt = vec![LanguageModelMessage::Tool { content: vec![ToolPart::ToolResult(ToolResultPart { tool_call_id: "call_789".into(), - output: ToolResultOutput::ErrorText { - value: "API request failed: timeout".into(), - provider_options: None, - }, - tool_name: "test".into(), - + output: serde_json::from_value( + json!({"type": "error-text", "value": "API request failed: timeout"}), + ) + .unwrap(), + tool_name: "get_weather".into(), provider_options: None, })], provider_options: None, @@ -818,12 +815,11 @@ mod convert_tests { let prompt = vec![LanguageModelMessage::Tool { content: vec![ToolPart::ToolResult(ToolResultPart { tool_call_id: "call_denied".into(), - output: ToolResultOutput::ExecutionDenied { - reason: Some("User declined the action".into()), - provider_options: None, - }, - tool_name: "test".into(), - + output: serde_json::from_value( + json!({"type": "execution-denied", "reason": "User declined the action"}), + ) + .unwrap(), + tool_name: "get_weather".into(), provider_options: None, })], provider_options: None, @@ -844,20 +840,15 @@ mod convert_tests { let prompt = vec![LanguageModelMessage::Tool { content: vec![ToolPart::ToolResult(ToolResultPart { tool_call_id: "call_content".into(), - output: ToolResultOutput::Content { - value: vec![ - ToolResultContent::Text(TextPart { - text: "First result".into(), - provider_options: None, - }), - ToolResultContent::Text(TextPart { - text: "Second result".into(), - provider_options: None, - }), - ], - }, - tool_name: "test".into(), - + output: serde_json::from_value(json!({ + "type": "content", + "value": [ + {"type": "text", "text": "First result"}, + {"type": "text", "text": "Second result"} + ] + })) + .unwrap(), + tool_name: "get_weather".into(), provider_options: None, })], provider_options: None, @@ -882,22 +873,16 @@ mod convert_tests { content: vec![ ToolPart::ToolResult(ToolResultPart { tool_call_id: "call_001".into(), - output: ToolResultOutput::Json { - value: json!({"temp": 72}), - provider_options: None, - }, - tool_name: "test".into(), - + output: serde_json::from_value(json!({"type": "json", "value": {"temp": 72}})) + .unwrap(), + tool_name: "get_weather".into(), provider_options: None, }), ToolPart::ToolResult(ToolResultPart { tool_call_id: "call_002".into(), - output: ToolResultOutput::Text { - value: "3:00 PM".into(), - provider_options: None, - }, - tool_name: "test".into(), - + output: serde_json::from_value(json!({"type": "text", "value": "3:00 PM"})) + .unwrap(), + tool_name: "get_weather".into(), provider_options: None, }), ], @@ -964,12 +949,8 @@ mod convert_tests { LanguageModelMessage::Tool { content: vec![ToolPart::ToolResult(ToolResultPart { tool_call_id: "call_weather".into(), - output: ToolResultOutput::Json { - value: json!({"temperature": 25, "condition": "cloudy"}), - provider_options: None, - }, - tool_name: "test".into(), - + output: serde_json::from_value(json!({"type": "json", "value": {"temperature": 25, "condition": "cloudy"}})).unwrap(), + tool_name: "get_weather".into(), provider_options: None, })], provider_options: None, @@ -1003,12 +984,8 @@ mod convert_tests { LanguageModelMessage::Tool { content: vec![ToolPart::ToolResult(ToolResultPart { tool_call_id: "call_weather".into(), - output: ToolResultOutput::Json { - value: json!({"temperature": 25, "condition": "cloudy"}), - provider_options: None, - }, - tool_name: "test".into(), - + output: serde_json::from_value(json!({"type": "json", "value": {"temperature": 25, "condition": "cloudy"}})).unwrap(), + tool_name: "get_weather".into(), provider_options: None, })], provider_options: None, @@ -1423,12 +1400,10 @@ mod do_generate_tests { // -- ProviderOptions reasoning tests -- fn lmstudio_opts(value: Value) -> Option { - Some( - aimux_core::shared::provider_namespace("lmstudio", value) - .expect("provider options object"), - ) + Some(aimux_core::shared::provider_namespace("lmstudio", value).unwrap()) } + /// TS: `packages/open-responses/src/responses/open-responses-language-model.test.ts:594` #[tokio::test] async fn provider_options_reasoning_summary_detailed() { let server = MockServer::start().await; @@ -1446,6 +1421,7 @@ mod do_generate_tests { assert_eq!(body["reasoning"], json!({"summary": "detailed"})); } + /// TS: `packages/open-responses/src/responses/open-responses-language-model.test.ts:623` #[tokio::test] async fn provider_options_combines_effort_with_summary() { let server = MockServer::start().await; @@ -1467,6 +1443,7 @@ mod do_generate_tests { ); } + /// TS: `packages/open-responses/src/responses/open-responses-language-model.test.ts:655` #[tokio::test] async fn provider_options_reasoning_summary_concise() { let server = MockServer::start().await; @@ -1484,6 +1461,7 @@ mod do_generate_tests { assert_eq!(body["reasoning"], json!({"summary": "concise"})); } + /// TS: `packages/open-responses/src/responses/open-responses-language-model.test.ts:685` #[tokio::test] async fn provider_options_no_reasoning_fields() { let server = MockServer::start().await; @@ -1503,6 +1481,7 @@ mod do_generate_tests { // -- Tool call parsing tests -- + /// TS: `packages/open-responses/src/responses/open-responses-language-model.test.ts:742` #[tokio::test] async fn parse_tool_call_from_response() { let server = MockServer::start().await; @@ -1547,6 +1526,7 @@ mod do_generate_tests { } } + /// TS: `packages/open-responses/src/responses/open-responses-language-model.test.ts:746` #[tokio::test] async fn tool_call_finish_reason() { let server = MockServer::start().await; @@ -1567,6 +1547,7 @@ mod do_generate_tests { assert_eq!(result.finish_reason.raw, None); } + /// TS: `packages/open-responses/src/responses/open-responses-language-model.test.ts:753` #[tokio::test] async fn tool_call_usage() { let server = MockServer::start().await; @@ -1607,6 +1588,7 @@ mod do_generate_tests { ) } + /// TS: `packages/open-responses/src/responses/open-responses-language-model.test.ts:772` #[tokio::test] async fn tool_choice_auto() { let server = MockServer::start().await; @@ -1626,6 +1608,7 @@ mod do_generate_tests { assert!(body.get("tools").is_some()); } + /// TS: `packages/open-responses/src/responses/open-responses-language-model.test.ts:784` #[tokio::test] async fn tool_choice_none() { let server = MockServer::start().await; @@ -1644,6 +1627,7 @@ mod do_generate_tests { assert_eq!(body["tool_choice"], json!("none")); } + /// TS: `packages/open-responses/src/responses/open-responses-language-model.test.ts:796` #[tokio::test] async fn tool_choice_required() { let server = MockServer::start().await; @@ -1662,6 +1646,7 @@ mod do_generate_tests { assert_eq!(body["tool_choice"], json!("required")); } + /// TS: `packages/open-responses/src/responses/open-responses-language-model.test.ts:808` #[tokio::test] async fn tool_choice_specific_tool() { let server = MockServer::start().await; @@ -1757,12 +1742,12 @@ mod do_generate_tests { LanguageModelMessage::Tool { content: vec![ToolPart::ToolResult(ToolResultPart { tool_call_id: "call_weather_123".into(), - output: ToolResultOutput::Json { - value: json!({"temperature": 22, "condition": "sunny", "humidity": 65}), - provider_options: None, - }, - tool_name: "test".into(), - + output: serde_json::from_value(json!({ + "type": "json", + "value": {"temperature": 22, "condition": "sunny", "humidity": 65} + })) + .unwrap(), + tool_name: "get_weather".into(), provider_options: None, })], provider_options: None, @@ -1812,7 +1797,7 @@ mod do_generate_tests { let server = MockServer::start().await; mock_json(&server, openai_pdf_json()).await; - let model = make_model(&server, "test-pdf-model"); + let model = make_model(&server, "gpt-4.1-nano"); let prompt = vec![LanguageModelMessage::User { content: vec![ UserPart::Text(TextPart { @@ -1821,9 +1806,9 @@ provider_options: None }), UserPart::File(FilePart { data: FileData::Url { -url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into(), - original_url: None, - }, +original_url: None, +url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into() +}, media_type: "application/pdf".into(), filename: None, provider_options: None @@ -1838,7 +1823,7 @@ provider_options: None assert_eq!( body, json!({ - "model": "test-pdf-model", + "model": "gpt-4.1-nano", "input": [{ "type": "message", "role": "user", @@ -1856,7 +1841,7 @@ provider_options: None let server = MockServer::start().await; mock_json(&server, openai_pdf_json()).await; - let model = make_model(&server, "test-pdf-model"); + let model = make_model(&server, "gpt-4.1-nano"); let prompt = vec![LanguageModelMessage::User { content: vec![ UserPart::Text(TextPart { @@ -1865,9 +1850,9 @@ provider_options: None }), UserPart::File(FilePart { data: FileData::Url { -url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into(), - original_url: None, - }, +original_url: None, +url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into() +}, media_type: "application/pdf".into(), filename: None, provider_options: None @@ -1892,7 +1877,7 @@ provider_options: None let server = MockServer::start().await; mock_json(&server, openai_pdf_json()).await; - let model = make_model(&server, "test-pdf-model"); + let model = make_model(&server, "gpt-4.1-nano"); let prompt = vec![LanguageModelMessage::User { content: vec![ UserPart::Text(TextPart { @@ -1901,9 +1886,9 @@ provider_options: None }), UserPart::File(FilePart { data: FileData::Url { -url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into(), - original_url: None, - }, +original_url: None, +url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into() +}, media_type: "application/pdf".into(), filename: None, provider_options: None @@ -2176,7 +2161,7 @@ mod do_stream_tests { ); mock_sse(&server, body).await; - let model = make_model(&server, "test-pdf-model"); + let model = make_model(&server, "gpt-4.1-nano"); let prompt = vec![LanguageModelMessage::User { content: vec![ UserPart::Text(TextPart { @@ -2185,9 +2170,9 @@ provider_options: None }), UserPart::File(FilePart { data: FileData::Url { -url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into(), - original_url: None, - }, +original_url: None, +url: "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf".into() +}, media_type: "application/pdf".into(), filename: None, provider_options: None diff --git a/aimux-providers/tests/openai_compatible_factory_test.rs b/aimux-providers/tests/openai_compatible_factory_test.rs index 8a712b81b..dd5a88478 100644 --- a/aimux-providers/tests/openai_compatible_factory_test.rs +++ b/aimux-providers/tests/openai_compatible_factory_test.rs @@ -326,7 +326,7 @@ fn message(role: Role, text: &str) -> LanguageModelMessage { })], provider_options: None, }, - Role::Tool => unreachable!("fixture has no tool messages"), + Role::Tool => panic!("tool messages require tool results"), } } diff --git a/aimux-providers/tests/openai_compatible_upstream_e2e_test.rs b/aimux-providers/tests/openai_compatible_upstream_e2e_test.rs new file mode 100644 index 000000000..8167c6ccd --- /dev/null +++ b/aimux-providers/tests/openai_compatible_upstream_e2e_test.rs @@ -0,0 +1,394 @@ +//! End-to-end subset of the pinned openai-compatible upstream tests. + +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use std::sync::Arc; + +use aimux_core::AiMuxError; +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_core::result::GenerateContent; +use aimux_core::shared::{Size, provider_namespace}; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::FunctionTool; +use aimux_core::types::FinishReasonUnified; +use aimux_providers::openai_compatible::{ + OpenAICompatibleProvider, OpenAICompatibleProviderSettings, create_openai_compatible, +}; +use futures::StreamExt; +use mock_fetch::{Canned, MockFetch}; +use serde_json::{Value, json}; + +const BASE_URL: &str = "https://my.api.com/v1"; +const IMAGE_PROMPT: &str = "A photorealistic astronaut riding a horse"; + +fn provider(fetch: &Arc) -> OpenAICompatibleProvider { + create_openai_compatible(OpenAICompatibleProviderSettings { + name: "test-provider".into(), + base_url: BASE_URL.into(), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} + +fn options() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} + +fn chat_response(message: Value, usage: Value) -> Value { + json!({"id":"chatcmpl-test", "created":1711115037, "model":"test-model", + "choices":[{"index":0,"message":message,"finish_reason":"stop"}], + "usage":usage}) +} + +fn stream_response(values: &[Value]) -> Canned { + let mut body = values + .iter() + .map(|value| format!("data: {value}\n\n")) + .collect::(); + body.push_str("data: [DONE]\n\n"); + Canned { + status: 200, + headers: vec![("content-type".into(), "text/event-stream".into())], + body: body.into_bytes(), + } +} + +fn assert_request(fetch: &MockFetch, path: &str, body: Value) { + let requests = fetch.seen(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].method, "POST"); + assert_eq!(requests[0].url, format!("{BASE_URL}{path}")); + assert_eq!(requests[0].json_body(), body); +} + +fn stream_request() -> Value { + json!({"model":"test-model","messages":[{"role":"user","content":"Hello"}],"stream":true}) +} + +/// TS: "should use camelCase metadata key when camelCase provider options are used" (openai-compatible/src/chat/openai-compatible-chat-language-model.test.ts) +/// TS: "should extract detailed token usage when available" (openai-compatible/src/chat/openai-compatible-chat-language-model.test.ts) +#[tokio::test] +async fn chat_generate_metadata_and_usage() { + let usage = json!({"prompt_tokens":20,"completion_tokens":30,"total_tokens":50, + "prompt_tokens_details":{"cached_tokens":5}, + "completion_tokens_details":{"reasoning_tokens":10,"accepted_prediction_tokens":15,"rejected_prediction_tokens":5}}); + let fetch = MockFetch::new(vec![Canned::json(&chat_response( + json!({"content":"Hello!"}), + usage.clone(), + ))]); + let mut opts = options(); + opts.provider_options = Some( + provider_namespace("testProvider", json!({"reasoningEffort":"high"})) + .expect("provider metadata must be an object"), + ); + let result = provider(&fetch) + .chat("test-model") + .do_generate(&opts) + .await + .unwrap(); + let expected = json!({"model":"test-model","messages":[{"role":"user","content":"Hello"}],"reasoning_effort":"high"}); + assert_request(&fetch, "/chat/completions", expected.clone()); + assert_eq!(result.request.unwrap().body, Some(expected)); + assert_eq!(result.content.len(), 1); + assert!(matches!(&result.content[0], GenerateContent::Text { text, .. } if text == "Hello!")); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(result.finish_reason.raw.as_deref(), Some("stop")); + assert_eq!(result.usage.input_tokens.total, Some(20)); + assert_eq!(result.usage.input_tokens.no_cache, Some(15)); + assert_eq!(result.usage.input_tokens.cache_read, Some(5)); + assert_eq!(result.usage.input_tokens.cache_write, None); + assert_eq!(result.usage.output_tokens.total, Some(30)); + assert_eq!(result.usage.output_tokens.text, Some(20)); + assert_eq!(result.usage.output_tokens.reasoning, Some(10)); + assert_eq!(result.usage.raw, usage.as_object().cloned()); + let metadata = result.provider_metadata.unwrap(); + assert_eq!( + metadata, + provider_namespace( + "testProvider", + json!({"acceptedPredictionTokens":15,"rejectedPredictionTokens":5}) + ) + .expect("provider metadata must be an object") + ); + assert!(!metadata.contains_key("test-provider")); + let response = result.response.unwrap(); + assert_eq!(response.id.as_deref(), Some("chatcmpl-test")); + assert_eq!(response.model_id.as_deref(), Some("test-model")); + assert_eq!( + response.headers.unwrap()["content-type"], + "application/json" + ); + assert!(result.warnings.is_empty()); +} + +/// TS: "should stream reasoning content before text deltas" (openai-compatible/src/chat/openai-compatible-chat-language-model.test.ts) +#[tokio::test] +async fn chat_stream_reasoning_before_text() { + let delta = |value: Value| { + json!({"id":"chatcmpl-test","created":1711357598, + "model":"test-model","choices":[{"index":0,"delta":value,"finish_reason":null}]}) + }; + let fetch = MockFetch::new(vec![stream_response(&[ + delta(json!({"role":"assistant","content":"","reasoning_content":"Let me think"})), + delta(json!({"content":"","reasoning_content":" about this"})), + delta(json!({"content":"Here's"})), + delta(json!({"content":" my response"})), + json!({"choices":[{"index":0,"delta":{},"finish_reason":"stop"}], + "usage":{"prompt_tokens":18,"completion_tokens":439}}), + ])]); + let result = provider(&fetch) + .chat("test-model") + .do_stream(&options()) + .await + .unwrap(); + assert_request(&fetch, "/chat/completions", stream_request()); + assert_eq!(result.request.unwrap().body, Some(stream_request())); + let parts: Vec<_> = result.stream.map(|part| part.unwrap()).collect().await; + assert_eq!(parts.len(), 11); + assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); + assert!( + matches!(&parts[1], StreamPart::ResponseMetadata(metadata) if metadata.id.as_deref() == Some("chatcmpl-test") && metadata.model_id.as_deref() == Some("test-model")) + ); + assert!(matches!(&parts[2], StreamPart::ReasoningStart { id, .. } if id == "reasoning-0")); + assert!(matches!(&parts[5], StreamPart::ReasoningEnd { id, .. } if id == "reasoning-0")); + assert!(matches!(&parts[6], StreamPart::TextStart { id, .. } if id == "txt-0")); + assert!(matches!(&parts[9], StreamPart::TextEnd { id, .. } if id == "txt-0")); + let deltas: Vec<_> = parts + .iter() + .filter_map(|part| match part { + StreamPart::ReasoningDelta { id, delta, .. } => { + Some(("reasoning", id.as_str(), delta.as_str())) + } + StreamPart::TextDelta { id, delta, .. } => Some(("text", id.as_str(), delta.as_str())), + _ => None, + }) + .collect(); + assert_eq!( + deltas, + vec![ + ("reasoning", "reasoning-0", "Let me think"), + ("reasoning", "reasoning-0", " about this"), + ("text", "txt-0", "Here's"), + ("text", "txt-0", " my response") + ] + ); + let StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } = &parts[10] + else { + panic!("expected finish"); + }; + assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(usage.input_tokens.total, Some(18)); + assert_eq!(usage.input_tokens.no_cache, Some(18)); + assert_eq!(usage.input_tokens.cache_read, Some(0)); + assert_eq!(usage.output_tokens.total, Some(439)); + assert_eq!(usage.output_tokens.text, Some(439)); + assert_eq!(usage.output_tokens.reasoning, Some(0)); + assert_eq!( + provider_metadata.as_ref().unwrap(), + &provider_namespace("test-provider", json!({})) + .expect("provider metadata must be an object") + ); +} + +/// TS: "should parse thought signature from extra_content and include in providerMetadata" (openai-compatible/src/chat/openai-compatible-chat-language-model.test.ts) +#[tokio::test] +async fn chat_tool_call_preserves_arguments_and_signature() { + let fetch = MockFetch::new(vec![Canned::json(&chat_response( + json!({"tool_calls":[{ + "id":"function-call-1","type":"function", + "function":{"name":"check_flight","arguments":"{\"flight\":\"AA100\"}"}, + "extra_content":{"google":{"thought_signature":""}}}]}), + Value::Null, + ))]); + let schema = json!({"type":"object","properties":{"flight":{"type":"string"}}, + "required":["flight"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}); + let mut opts = options(); + opts.tools = Some(vec![ + FunctionTool::new("check_flight", schema.clone()).into(), + ]); + let result = provider(&fetch) + .chat("test-model") + .do_generate(&opts) + .await + .unwrap(); + assert_request( + &fetch, + "/chat/completions", + json!({"model":"test-model", + "messages":[{"role":"user","content":"Hello"}], + "tools":[{"type":"function","function":{"name":"check_flight","parameters":schema}}]}), + ); + assert_eq!(result.content.len(), 1); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(result.finish_reason.raw.as_deref(), Some("stop")); + let GenerateContent::ToolCall(call) = &result.content[0] else { + panic!("expected tool call"); + }; + assert_eq!(call.tool_call_id, "function-call-1"); + assert_eq!(call.tool_name, "check_flight"); + assert_eq!(call.input, "{\"flight\":\"AA100\"}"); + assert_eq!( + call.provider_metadata.as_ref().unwrap(), + &provider_namespace("test-provider", json!({"thoughtSignature":""})) + .expect("provider metadata must be an object") + ); +} + +/// TS: "should preserve structured error stream parts" (openai-compatible/src/chat/openai-compatible-chat-language-model.test.ts) +#[tokio::test] +async fn chat_stream_preserves_structured_error() { + let fetch = MockFetch::new(vec![stream_response(&[json!({"error":{ + "message":"Context length exceeded","code":"CONTEXT_LENGTH_EXCEEDED"}})])]); + let result = provider(&fetch) + .chat("test-model") + .do_stream(&options()) + .await + .unwrap(); + assert_request(&fetch, "/chat/completions", stream_request()); + let parts: Vec<_> = result.stream.map(|part| part.unwrap()).collect().await; + assert_eq!(parts.len(), 3); + assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); + let StreamPart::Error { + error: AiMuxError::ApiCall(error), + } = &parts[1] + else { + panic!("expected structured API error"); + }; + assert_eq!(error.message, "Context length exceeded"); + assert_eq!( + error.provider_code.as_deref(), + Some("CONTEXT_LENGTH_EXCEEDED") + ); + let StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } = &parts[2] + else { + panic!("expected finish"); + }; + assert_eq!(finish_reason.unified, FinishReasonUnified::Error); + assert_eq!(finish_reason.raw, None); + assert_eq!(usage.input_tokens.total, None); + assert_eq!(usage.output_tokens.total, None); + assert_eq!(usage.raw, None); + assert_eq!( + provider_metadata.as_ref().unwrap(), + &provider_namespace("test-provider", json!({})) + .expect("provider metadata must be an object") + ); +} + +/// TS: "should extract embedding" (openai-compatible/src/embedding/openai-compatible-embedding-model.test.ts) +/// TS: "should pass the model and the values" (openai-compatible/src/embedding/openai-compatible-embedding-model.test.ts) +#[tokio::test] +async fn embedding_request_and_result() { + let body = json!({"object":"list","data":[ + {"object":"embedding","index":0,"embedding":[0.1,0.2,0.3,0.4,0.5]}, + {"object":"embedding","index":1,"embedding":[0.6,0.7,0.8,0.9,1.0]}], + "model":"test-embedding","usage":{"prompt_tokens":8,"total_tokens":8}}); + let fetch = MockFetch::new(vec![Canned::json(&body)]); + let mut opts = EmbeddingCallOptions::new("sunny day at the beach"); + opts.values.push("rainy day in the city".into()); + let result = provider(&fetch) + .embedding("test-embedding") + .do_embed(&opts) + .await + .unwrap(); + assert_request( + &fetch, + "/embeddings", + json!({"model":"test-embedding", + "input":["sunny day at the beach","rainy day in the city"],"encoding_format":"float"}), + ); + assert_eq!( + result.embeddings, + vec![vec![0.1, 0.2, 0.3, 0.4, 0.5], vec![0.6, 0.7, 0.8, 0.9, 1.0]] + ); + assert_eq!(result.usage.unwrap().tokens, 8); + let response = result.response.unwrap(); + assert_eq!(response.body, Some(body)); + assert_eq!( + response.headers.unwrap()["content-type"], + "application/json" + ); + assert!(result.warnings.is_empty()); +} + +fn image_options() -> ImageCallOptions { + let mut opts = ImageCallOptions::new(IMAGE_PROMPT); + opts.size = Some(Size::new(1024, 1024)); + opts +} + +/// TS: "should return the raw b64_json content" (openai-compatible/src/image/openai-compatible-image-model.test.ts) +/// TS: "should map the usage object reported by the provider" (openai-compatible/src/image/openai-compatible-image-model.test.ts) +#[tokio::test] +async fn image_generate_raw_content_and_usage() { + let fetch = MockFetch::new(vec![Canned::json(&json!({"data":[ + {"b64_json":"test1234"},{"b64_json":"test5678"}], + "usage":{"input_tokens":12,"output_tokens":4,"total_tokens":16}}))]); + let mut opts = image_options(); + opts.n = 2; + let result = provider(&fetch) + .image("test-image") + .do_generate(&opts) + .await + .unwrap(); + assert_request( + &fetch, + "/images/generations", + json!({"model":"test-image", + "prompt":IMAGE_PROMPT,"n":2,"size":"1024x1024"}), + ); + let ImageOutputs::Base64(images) = result.images else { + panic!("expected base64 images"); + }; + assert_eq!(images, vec!["test1234", "test5678"]); + let usage = result.usage.unwrap(); + assert_eq!(usage.input_tokens, Some(12)); + assert_eq!(usage.output_tokens, Some(4)); + assert_eq!(usage.total_tokens, Some(16)); + assert_eq!(result.response.model_id.as_deref(), Some("test-image")); + assert_eq!( + result.response.headers.unwrap()["content-type"], + "application/json" + ); + assert!(result.warnings.is_empty()); +} + +/// TS: "should handle API errors with default error structure" (openai-compatible/src/image/openai-compatible-image-model.test.ts) +#[tokio::test] +async fn image_generate_api_error() { + let mut canned = Canned::json(&json!({"error":{"message":"Invalid prompt content", + "type":"invalid_request_error","param":null,"code":null}})); + canned.status = 400; + let fetch = MockFetch::new(vec![canned]); + let error = provider(&fetch) + .image("test-image") + .do_generate(&image_options()) + .await + .unwrap_err(); + assert_request( + &fetch, + "/images/generations", + json!({"model":"test-image", + "prompt":IMAGE_PROMPT,"n":1,"size":"1024x1024"}), + ); + let AiMuxError::ApiCall(error) = error else { + panic!("expected API error"); + }; + assert_eq!(error.message, "Invalid prompt content"); + assert_eq!(error.status_code, Some(400)); + assert_eq!(error.url, format!("{BASE_URL}/images/generations")); +} diff --git a/aimux-providers/tests/openai_factory_test.rs b/aimux-providers/tests/openai_factory_test.rs index f2f437578..872de6e86 100644 --- a/aimux-providers/tests/openai_factory_test.rs +++ b/aimux-providers/tests/openai_factory_test.rs @@ -304,18 +304,18 @@ impl Fixture { fn message(role: Role, text: &str) -> LanguageModelMessage { match role { Role::System => LanguageModelMessage::System { - content: text.into(), + content: text.to_string(), provider_options: None, }, Role::User => LanguageModelMessage::user_text(text), Role::Assistant => LanguageModelMessage::Assistant { content: vec![AssistantPart::Text(TextPart { - text: text.into(), + text: text.to_string(), provider_options: None, })], provider_options: None, }, - Role::Tool => panic!("unmapped tool role"), + Role::Tool => panic!("tool messages require tool results"), } } diff --git a/aimux-providers/tests/openai_upstream_e2e_test.rs b/aimux-providers/tests/openai_upstream_e2e_test.rs new file mode 100644 index 000000000..29c399312 --- /dev/null +++ b/aimux-providers/tests/openai_upstream_e2e_test.rs @@ -0,0 +1,574 @@ +//! Small end-to-end subset of the pinned upstream tests, using an in-process transport. +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::error::AiMuxError; +use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_core::result::GenerateContent; +use aimux_core::shared::{Size, provider_namespace}; +use aimux_core::speech_model::{AudioData, SpeechCallOptions, SpeechModel}; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::{FunctionTool, Tool}; +use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; +use aimux_core::types::FinishReasonUnified; +use aimux_providers::openai::{OpenAIProvider, OpenAIProviderSettings, create_openai}; +use futures::StreamExt; +use mock_fetch::{Canned, MockFetch}; +use serde_json::{Value, json}; + +fn provider(fetch: &std::sync::Arc) -> OpenAIProvider { + create_openai(OpenAIProviderSettings { + base_url: Some("https://example.test/v1".into()), + api_key: Some("test-key".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} +fn options() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("Hello")]) +} +fn chat_response() -> Value { + json!({"id":"chatcmpl-test","model":"test-model","created":1711115037,"choices":[{"index":0,"message":{"role":"assistant","content":"Hello, World!"},"finish_reason":"stop"}],"usage":{"prompt_tokens":4,"total_tokens":34,"completion_tokens":30}}) +} +fn responses_response() -> Value { + json!({"id":"resp_test","created_at":1741257730,"output":[{"id":"msg_test","type":"message","content":[{"type":"output_text","text":"answer text","annotations":[]}]}],"usage":{"input_tokens":345,"input_tokens_details":{"cached_tokens":234,"cache_write_tokens":45,"future_input_detail":{"tokens":7}},"output_tokens":538,"output_tokens_details":{"reasoning_tokens":123,"future_output_detail":["preserved"]},"total_tokens":572,"future_usage_field":{"value":true}}}) +} +fn sse(events: Vec) -> Canned { + Canned { + status: 200, + headers: vec![("content-type".into(), "text/event-stream".into())], + body: events + .into_iter() + .map(|event| format!("data: {event}\n\n")) + .collect::() + .into_bytes(), + } +} +fn assert_request(fetch: &MockFetch, endpoint: &str, expected: Value) { + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "POST"); + assert_eq!(seen[0].url, format!("https://example.test/v1{endpoint}")); + assert_eq!(seen[0].headers["authorization"], "Bearer test-key"); + assert_eq!(seen[0].json_body(), expected); +} + +/// TS: "should extract text response" (src/chat/openai-chat-language-model.test.ts) +/// TS: "should send request body" (src/chat/openai-chat-language-model.test.ts) +/// TS: "should extract usage" (src/chat/openai-chat-language-model.test.ts) +#[tokio::test] +async fn chat_generate() { + let fetch = MockFetch::new(vec![Canned::json(&chat_response())]); + let result = provider(&fetch) + .chat("test-model") + .do_generate(&options()) + .await + .unwrap(); + let body = json!({"model":"test-model","messages":[{"role":"user","content":"Hello"}]}); + assert_request(&fetch, "/chat/completions", body.clone()); + assert_eq!(result.request.unwrap().body, Some(body)); + assert!( + matches!(&result.content[..],[GenerateContent::Text{text,..}] if text=="Hello, World!") + ); + assert_eq!(result.usage.input_tokens.total, Some(4)); + assert_eq!(result.usage.input_tokens.no_cache, Some(4)); + assert_eq!(result.usage.output_tokens.text, Some(30)); + assert_eq!( + result.usage.raw, + chat_response()["usage"].as_object().cloned() + ); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!( + result.response.unwrap().id.as_deref(), + Some("chatcmpl-test") + ); +} + +/// TS: "should stream text deltas" (src/chat/openai-chat-language-model.test.ts) +#[tokio::test] +async fn chat_stream() { + let chunk = |delta: Value, finish: Value| json!({"id":"chatcmpl-test","model":"test-model","created":1702657020,"choices":[{"index":0,"delta":delta,"finish_reason":finish}]}); + let fetch = MockFetch::new(vec![sse(vec![ + chunk(json!({"role":"assistant","content":""}), Value::Null), + chunk(json!({"content":"Hello"}), Value::Null), + chunk(json!({"content":", "}), Value::Null), + chunk(json!({"content":"World!"}), Value::Null), + chunk(json!({}), json!("stop")), + json!({"id":"chatcmpl-test","model":"test-model","created":1702657020,"choices":[],"usage":{"prompt_tokens":17,"total_tokens":244,"completion_tokens":227}}), + ])]); + let result = provider(&fetch) + .chat("test-model") + .do_stream(&options()) + .await + .unwrap(); + let parts = result + .stream + .map(|part| part.unwrap()) + .collect::>() + .await; + assert_request( + &fetch, + "/chat/completions", + json!({"model":"test-model","messages":[{"role":"user","content":"Hello"}],"stream":true,"stream_options":{"include_usage":true}}), + ); + let text = parts + .iter() + .filter_map(|part| match part { + StreamPart::TextDelta { delta, .. } => Some(delta.as_str()), + _ => None, + }) + .collect::(); + assert_eq!(text, "Hello, World!"); + assert_eq!( + parts + .iter() + .filter(|part| matches!(part, StreamPart::TextStart { .. })) + .count(), + 1 + ); + assert_eq!( + parts + .iter() + .filter(|part| matches!(part, StreamPart::TextEnd { .. })) + .count(), + 1 + ); + assert!( + matches!(parts.last(),Some(StreamPart::Finish{finish_reason,usage,..}) if finish_reason.unified==FinishReasonUnified::Stop && usage.input_tokens.total==Some(17) && usage.output_tokens.total==Some(227)) + ); +} + +/// TS: "should parse tool results" (src/chat/openai-chat-language-model.test.ts) +#[tokio::test] +async fn chat_tool_call() { + let mut response = chat_response(); + response["choices"][0]["message"]["content"] = json!(""); + response["choices"][0]["message"]["tool_calls"] = json!([{"id":"call_test","type":"function","function":{"name":"test-tool","arguments":"{\"value\":\"Spark\"}"}}]); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let schema = json!({"type":"object","properties":{"value":{"type":"string"}},"required":["value"],"additionalProperties":false}); + let mut opts = options(); + opts.tools = Some(vec![Tool::Function(FunctionTool::new( + "test-tool", + schema.clone(), + ))]); + let result = provider(&fetch) + .chat("test-model") + .do_generate(&opts) + .await + .unwrap(); + let body = fetch.seen()[0].json_body(); + assert_eq!(body["model"], "test-model"); + assert_eq!(body["messages"], json!([{"role":"user","content":"Hello"}])); + assert_eq!(body["tools"][0]["function"]["name"], "test-tool"); + assert_eq!(body["tools"][0]["function"]["parameters"], schema); + assert!( + matches!(&result.content[..],[GenerateContent::ToolCall(call)] if call.tool_call_id=="call_test" && call.tool_name=="test-tool" && call.input==r#"{"value":"Spark"}"#) + ); +} + +/// TS: "should parse OpenRouter resource exhausted error" (src/openai-error.test.ts) +#[tokio::test] +async fn chat_error_response() { + let message = "{\n \"error\": {\n \"code\": 429,\n \"message\": \"Resource has been exhausted (e.g. check quota).\",\n \"status\": \"RESOURCE_EXHAUSTED\"\n }\n}\n"; + let response = json!({"error":{"message":message,"code":429}}); + let mut canned = Canned::json(&response); + canned.status = 429; + let fetch = MockFetch::new(vec![canned]); + let error = provider(&fetch) + .chat("test-model") + .do_generate(&options()) + .await + .unwrap_err(); + assert_request( + &fetch, + "/chat/completions", + json!({"model":"test-model","messages":[{"role":"user","content":"Hello"}]}), + ); + let AiMuxError::ApiCall(error) = error else { + panic!("expected API error") + }; + assert_eq!(error.message, message); + assert_eq!(error.status_code, Some(429)); + assert_eq!(error.data, Some(response)); +} + +/// TS: "should generate text" (src/responses/openai-responses-language-model.test.ts) +/// TS: "should extract usage" (src/responses/openai-responses-language-model.test.ts) +#[tokio::test] +async fn responses_generate() { + let response = responses_response(); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let result = provider(&fetch) + .responses("test-model") + .do_generate(&options()) + .await + .unwrap(); + assert_request( + &fetch, + "/responses", + json!({"model":"test-model","input":[{"role":"user","content":[{"type":"input_text","text":"Hello"}]}]}), + ); + assert_eq!( + result.content, + vec![GenerateContent::Text { + text: "answer text".into(), + provider_metadata: Some( + provider_namespace("openai", json!({"itemId":"msg_test"})) + .expect("provider metadata must be an object") + ) + }] + ); + assert_eq!(result.usage.input_tokens.total, Some(345)); + assert_eq!(result.usage.input_tokens.cache_read, Some(234)); + assert_eq!(result.usage.input_tokens.cache_write, Some(45)); + assert_eq!(result.usage.input_tokens.no_cache, Some(66)); + assert_eq!(result.usage.output_tokens.total, Some(538)); + assert_eq!(result.usage.output_tokens.reasoning, Some(123)); + assert_eq!(result.usage.output_tokens.text, Some(415)); + assert_eq!(result.usage.raw, response["usage"].as_object().cloned()); + assert_eq!(result.response.unwrap().id.as_deref(), Some("resp_test")); +} + +/// TS: "should stream text deltas" (src/responses/openai-responses-language-model.test.ts) +#[tokio::test] +async fn responses_stream() { + let events=json!([{"type": "response.created", "response": {"id": "resp_67c9a81b6a048190a9ee441c5755a4e8", "object": "response", "created_at": 1741269019, "status": "in_progress", "error": null, "incomplete_details": null, "input": [], "instructions": null, "max_output_tokens": null, "model": "test-model", "output": [], "parallel_tool_calls": true, "previous_response_id": null, "reasoning": {"effort": null, "summary": null}, "store": true, "temperature": 0.3, "text": {"format": {"type": "text"}}, "tool_choice": "auto", "tools": [], "top_p": 1, "truncation": "disabled", "usage": null, "user": null, "metadata": {}}}, {"type": "response.in_progress", "response": {"id": "resp_67c9a81b6a048190a9ee441c5755a4e8", "object": "response", "created_at": 1741269019, "status": "in_progress", "error": null, "incomplete_details": null, "input": [], "instructions": null, "max_output_tokens": null, "model": "test-model", "output": [], "parallel_tool_calls": true, "previous_response_id": null, "reasoning": {"effort": null, "summary": null}, "store": true, "temperature": 0.3, "text": {"format": {"type": "text"}}, "tool_choice": "auto", "tools": [], "top_p": 1, "truncation": "disabled", "usage": null, "user": null, "metadata": {}}}, {"type": "response.output_item.added", "output_index": 0, "item": {"id": "msg_67c9a81dea8c8190b79651a2b3adf91e", "type": "message", "status": "in_progress", "role": "assistant", "content": []}}, {"type": "response.content_part.added", "item_id": "msg_67c9a81dea8c8190b79651a2b3adf91e", "output_index": 0, "content_index": 0, "part": {"type": "output_text", "text": "", "annotations": [], "logprobs": []}}, {"type": "response.output_text.delta", "item_id": "msg_67c9a81dea8c8190b79651a2b3adf91e", "output_index": 0, "content_index": 0, "delta": "Hello,", "logprobs": []}, {"type": "response.output_text.delta", "item_id": "msg_67c9a81dea8c8190b79651a2b3adf91e", "output_index": 0, "content_index": 0, "delta": " World!", "logprobs": []}, {"type": "response.output_text.done", "item_id": "msg_67c9a8787f4c8190b49c858d4c1cf20c", "output_index": 0, "content_index": 0, "text": "Hello, World!"}, {"type": "response.content_part.done", "item_id": "msg_67c9a8787f4c8190b49c858d4c1cf20c", "output_index": 0, "content_index": 0, "part": {"type": "output_text", "text": "Hello, World!", "annotations": [], "logprobs": []}}, {"type": "response.output_item.done", "output_index": 0, "item": {"id": "msg_67c9a8787f4c8190b49c858d4c1cf20c", "type": "message", "status": "completed", "role": "assistant", "content": [{"type": "output_text", "text": "Hello, World!", "annotations": [], "logprobs": []}]}}, {"type": "response.completed", "response": {"id": "resp_67c9a878139c8190aa2e3105411b408b", "object": "response", "created_at": 1741269112, "status": "completed", "error": null, "incomplete_details": null, "input": [], "instructions": null, "max_output_tokens": null, "model": "test-model", "output": [{"id": "msg_67c9a8787f4c8190b49c858d4c1cf20c", "type": "message", "status": "completed", "role": "assistant", "content": [{"type": "output_text", "text": "Hello, World!", "annotations": []}]}], "parallel_tool_calls": true, "previous_response_id": null, "reasoning": {"effort": null, "summary": null}, "store": true, "temperature": 0.3, "text": {"format": {"type": "text"}}, "tool_choice": "auto", "tools": [], "top_p": 1, "truncation": "disabled", "usage": {"input_tokens": 543, "input_tokens_details": {"cached_tokens": 234}, "output_tokens": 478, "output_tokens_details": {"reasoning_tokens": 123}, "total_tokens": 512}, "user": null, "metadata": {}}}]).as_array().unwrap().clone(); + let fetch = MockFetch::new(vec![sse(events)]); + let result = provider(&fetch) + .responses("test-model") + .do_stream(&options()) + .await + .unwrap(); + let parts = result + .stream + .map(|part| part.unwrap()) + .collect::>() + .await; + assert_request( + &fetch, + "/responses", + json!({"model":"test-model","input":[{"role":"user","content":[{"type":"input_text","text":"Hello"}]}],"stream":true}), + ); + assert_eq!( + serde_json::to_value(parts).unwrap(), + json!([ + {"StreamStart": {"warnings": []}}, + {"ResponseMetadata": { + "id": "resp_67c9a81b6a048190a9ee441c5755a4e8", + "model_id": "test-model", + "timestamp": "2025-03-06T13:50:19+00:00" + }}, + {"TextStart": { + "id": "msg_67c9a81dea8c8190b79651a2b3adf91e", + "provider_metadata": {"openai": {"itemId": "msg_67c9a81dea8c8190b79651a2b3adf91e"}} + }}, + {"TextDelta": {"id": "msg_67c9a81dea8c8190b79651a2b3adf91e", "delta": "Hello,"}}, + {"TextDelta": {"id": "msg_67c9a81dea8c8190b79651a2b3adf91e", "delta": " World!"}}, + {"TextEnd": { + "id": "msg_67c9a81dea8c8190b79651a2b3adf91e", + "provider_metadata": {"openai": {"itemId": "msg_67c9a81dea8c8190b79651a2b3adf91e"}} + }}, + {"Finish": { + "finish_reason": {"unified": "stop", "raw": null}, + "provider_metadata": {"openai": {"responseId": "resp_67c9a81b6a048190a9ee441c5755a4e8"}}, + "usage": { + "input_tokens": {"total": 543, "no_cache": 309, "cache_read": 234}, + "output_tokens": {"total": 478, "reasoning": 123, "text": 355}, + "raw": { + "input_tokens": 543, + "input_tokens_details": {"cached_tokens": 234}, + "output_tokens": 478, + "output_tokens_details": {"reasoning_tokens": 123}, + "total_tokens": 512 + } + } + }} + ]) + ); +} + +/// TS: "should throw a descriptive error when the response has no output" (src/responses/openai-responses-language-model.test.ts) +#[tokio::test] +async fn responses_error_response() { + let fetch = MockFetch::new(vec![Canned::json( + &json!({"id":"resp_no_output","status":"incomplete","incomplete_details":{"reason":"content_filter"}}), + )]); + let error = provider(&fetch) + .responses("test-model") + .do_generate(&options()) + .await + .unwrap_err(); + assert_request( + &fetch, + "/responses", + json!({"model":"test-model","input":[{"role":"user","content":[{"type":"input_text","text":"Hello"}]}]}), + ); + assert!( + error + .to_string() + .contains("Responses API returned no output (content_filter)"), + "{error}" + ); +} + +fn embedding_response() -> Value { + serde_json::from_str( + r#"{ + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": [ + 0.0057293195, -0.012727811, 0.020042092, -0.013437585, 0.022833068 + ] + }, + { + "object": "embedding", + "index": 1, + "embedding": [ + -0.037104916, -0.05178114, -0.008340587, 0.001164541, -0.0035253682 + ] + } + ], + "model": "test-embedding", + "usage": { + "prompt_tokens": 12, + "total_tokens": 12 + } +} +"#, + ) + .unwrap() +} + +fn image_response() -> Value { + serde_json::from_str(r#"{ + "created": 1770935200, + "data": [ + { + "b64_json": "iVBORw0KGgoAAAANSUhEUgAABAAAAAQACAIAAADwf7zUAAA3CGNhQlgAADcIanVtYgAAAB5qdW1kYzJwYQARABCAAACqADibcQNj", + "revised_prompt": "A small and adorable baby sea otter. This little creature is covered in a thick and fluffy brown fur, its tiny paws are slightly visible. The otter has bright, curious eyes and it's floating on its back on a calm sea, surrounded by floating seaweed." + }, + { + "b64_json": "iVBORw0KGgoAAAANSUhEUgAABAAAAAQACAIAAADwf7zUAAEp2GNhQlgAASnYanVtYgAAAB5qdW1kYzJwYQARABCAAACqADibcQNj" + } + ] +} +"#).unwrap() +} + +fn transcription_response() -> Value { + serde_json::from_str(r#"{ + "task": "transcribe", + "language": "english", + "duration": 36.709999084472656, + "text": "Galileo was an American robotic space program that studied the planet Jupiter and its moons, as well as several other solar system bodies.", + "words": [ + { + "word": "Galileo", + "start": 0, + "end": 0.6600000262260437 + }, + { + "word": "was", + "start": 0.6600000262260437, + "end": 0.8999999761581421 + }, + { + "word": "an", + "start": 0.8999999761581421, + "end": 1.1399999856948853 + }, + { + "word": "American", + "start": 1.1399999856948853, + "end": 1.5 + }, + { + "word": "robotic", + "start": 1.5, + "end": 2.0199999809265137 + } + ], + "usage": { + "type": "duration", + "seconds": 37 + } +} +"#).unwrap() +} + +/// TS: "should extract embedding" (src/embedding/openai-embedding-model.test.ts) +/// TS: "should pass the model and the values" (src/embedding/openai-embedding-model.test.ts) +#[tokio::test] +async fn embedding_generate() { + let response = embedding_response(); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let mut opts = EmbeddingCallOptions::new("sunny day at the beach"); + opts.values.push("rainy day in the city".into()); + let result = provider(&fetch) + .embedding("test-embedding") + .do_embed(&opts) + .await + .unwrap(); + assert_request( + &fetch, + "/embeddings", + json!({"model":"test-embedding","input":opts.values,"encoding_format":"float"}), + ); + assert_eq!( + result.embeddings, + vec![ + vec![ + 0.0057293195, + -0.012727811, + 0.020042092, + -0.013437585, + 0.022833068 + ], + vec![ + -0.037104916, + -0.05178114, + -0.008340587, + 0.001164541, + -0.0035253682 + ] + ] + ); + assert_eq!(result.usage.unwrap().tokens, 12); + assert_eq!(result.response.unwrap().body, Some(response)); +} + +/// TS: "should pass the model and the settings" (src/image/openai-image-model.test.ts) +/// TS: "should extract the generated images" (src/image/openai-image-model.test.ts) +#[tokio::test] +async fn image_generate() { + let response = image_response(); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let mut opts = ImageCallOptions::new("A cute baby sea otter"); + opts.size = Some(Size::new(1024, 1024)); + opts.provider_options = provider_namespace("openai", json!({"style":"vivid"})) + .expect("provider metadata must be an object"); + let result = provider(&fetch) + .image("test-image") + .do_generate(&opts) + .await + .unwrap(); + assert_request( + &fetch, + "/images/generations", + json!({"model":"test-image","prompt":"A cute baby sea otter","n":1,"size":"1024x1024","style":"vivid","response_format":"b64_json"}), + ); + let expected = response["data"] + .as_array() + .unwrap() + .iter() + .map(|image| image["b64_json"].as_str().unwrap().to_owned()) + .collect::>(); + assert!(matches!(result.images,ImageOutputs::Base64(images) if images==expected)); +} + +/// TS: "should pass the model and text" (src/speech/openai-speech-model.test.ts) +/// TS: "should return audio data with correct content type" (src/speech/openai-speech-model.test.ts) +#[tokio::test] +async fn speech_generate() { + let audio = vec![1, 2, 3, 4]; + let fetch = MockFetch::new(vec![Canned { + status: 200, + headers: vec![("content-type".into(), "audio/mpeg".into())], + body: audio.clone(), + }]); + let result = provider(&fetch) + .speech("test-speech") + .do_generate(&SpeechCallOptions::new("Hello from the AI SDK!")) + .await + .unwrap(); + assert_request( + &fetch, + "/audio/speech", + json!({"model":"test-speech","input":"Hello from the AI SDK!","voice":"alloy","response_format":"mp3"}), + ); + assert!(matches!(result.audio,AudioData::Binary(bytes) if bytes==audio)); + assert_eq!( + result.response.headers.unwrap()["content-type"], + "audio/mpeg" + ); +} + +/// TS: "should extract the transcription text" (src/transcription/openai-transcription-model.test.ts) +#[tokio::test] +async fn transcription_generate() { + let response = transcription_response(); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let opts = TranscriptionCallOptions::new(AudioInput::Binary(vec![1, 2, 3]), "audio/wav"); + let result = provider(&fetch) + .transcription("test-transcription") + .do_generate(&opts) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].url, "https://example.test/v1/audio/transcriptions"); + assert_eq!(seen[0].method, "POST"); + assert!(seen[0].headers["content-type"].starts_with("multipart/form-data; boundary=")); + let body = String::from_utf8_lossy(&seen[0].body); + assert!(body.contains("name=\"model\"\r\n\r\ntest-transcription\r\n")); + assert!(body.contains("name=\"file\"; filename=\"audio.wav\"")); + assert!(seen[0].body.windows(3).any(|window| window == [1, 2, 3])); + assert_eq!(result.text, response["text"].as_str().unwrap()); + assert_eq!(result.language.as_deref(), Some("en")); + assert_eq!(result.duration_in_seconds, Some(36.709999084472656)); + assert_eq!(result.segments.len(), 5); +} + +/// TS: "should handle streaming web search without action" (src/responses/openai-responses-language-model.test.ts) +#[tokio::test] +async fn responses_web_search_stream() { + let fetch = MockFetch::new(vec![sse(vec![ + json!({"type":"response.created","response":{"id":"resp_missing_action","created_at":1741630255,"model":"test-model"}}), + json!({"type":"response.output_item.added","output_index":0,"item":{"type":"web_search_call","id":"ws_missing_action","status":"in_progress"}}), + json!({"type":"response.web_search_call.in_progress","output_index":0,"item_id":"ws_missing_action"}), + json!({"type":"response.web_search_call.completed","output_index":0,"item_id":"ws_missing_action"}), + json!({"type":"response.output_item.done","output_index":0,"item":{"type":"web_search_call","id":"ws_missing_action","status":"completed"}}), + json!({"type":"response.completed","response":{"id":"resp_missing_action","usage":{"input_tokens":10,"input_tokens_details":{"cached_tokens":0},"output_tokens":2,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":12}}}), + ])]); + let mut call = options(); + call.tools = Some(vec![Tool::Provider(aimux_core::tool::ProviderTool { + id: "openai.web_search".into(), + name: "webSearch".into(), + args: serde_json::Map::new(), + })]); + let result = provider(&fetch) + .responses("test-model") + .do_stream(&call) + .await + .unwrap(); + let parts = result + .stream + .map(|part| part.unwrap()) + .collect::>() + .await; + assert_request( + &fetch, + "/responses", + json!({"model":"test-model","input":[{"role":"user","content":[{"type":"input_text","text":"Hello"}]}],"stream":true,"tools":[{"type":"web_search"}],"include":["web_search_call.action.sources"]}), + ); + assert_eq!( + serde_json::to_value(parts).unwrap(), + json!([ + {"StreamStart":{"warnings":[]}}, + {"ResponseMetadata":{"id":"resp_missing_action","model_id":"test-model","timestamp":"2025-03-10T18:10:55+00:00"}}, + {"ToolInputStart":{"id":"ws_missing_action","tool_name":"webSearch","provider_executed":true}}, + {"ToolInputEnd":{"id":"ws_missing_action"}}, + {"ToolCall":{"tool_call_id":"ws_missing_action","tool_name":"webSearch","input":"{}","provider_executed":true}}, + {"ToolResult":{"tool_call_id":"ws_missing_action","tool_name":"webSearch","result":{}}}, + {"Finish":{"finish_reason":{"unified":"stop","raw":null},"usage":{"input_tokens":{"total":10,"no_cache":10,"cache_read":0},"output_tokens":{"total":2,"reasoning":0,"text":2},"raw":{"input_tokens":10,"input_tokens_details":{"cached_tokens":0},"output_tokens":2,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":12}},"provider_metadata":{"openai":{"responseId":"resp_missing_action"}}}} + ]) + ); +} diff --git a/aimux-providers/tests/openrouter_test.rs b/aimux-providers/tests/openrouter_test.rs index 73bfc3a7a..ab9e1eb66 100644 --- a/aimux-providers/tests/openrouter_test.rs +++ b/aimux-providers/tests/openrouter_test.rs @@ -131,10 +131,6 @@ fn make_provider(server: &MockServer) -> OpenAICompatibleProvider { openrouter_at(server.uri(), "test-api-key") } -// ════════════════════════════════════════════════════════════════════════════ -// Provider configuration -// ════════════════════════════════════════════════════════════════════════════ - /// Custom API key is sent in the `Authorization: Bearer` header. #[tokio::test] async fn custom_api_key_used_in_auth_header() { diff --git a/aimux-providers/tests/parallel_ai_test.rs b/aimux-providers/tests/parallel_ai_test.rs index 36a44dcd2..bff08d9b3 100644 --- a/aimux-providers/tests/parallel_ai_test.rs +++ b/aimux-providers/tests/parallel_ai_test.rs @@ -194,8 +194,6 @@ async fn status_403_maps_to_provider_error() { ); } -// -- Provider trait ---------------------------------------------------------- - #[test] fn model_id_is_parallel_search() { let config = ParallelAiProviderSettings { diff --git a/aimux-providers/tests/provider_error_test.rs b/aimux-providers/tests/provider_error_test.rs index 0c917e66c..dcc472aea 100644 --- a/aimux-providers/tests/provider_error_test.rs +++ b/aimux-providers/tests/provider_error_test.rs @@ -536,103 +536,4 @@ mod anthropic_stream_errors { "expected retryable 529 ApiCall error, got {error:?}" ); } - - /// TS: "should forward error chunks" + "should forward overloaded error - /// during streaming". - /// - /// A normal text delta is emitted first, then an error event mid-stream. - /// The error must be forwarded as a `StreamPart::Error`. - #[tokio::test] - async fn mid_stream_error_is_forwarded() { - let server = MockServer::start().await; - let body = concat!( - "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_01\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-haiku-20240307\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":17,\"output_tokens\":1}}}\n\n", - "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", - "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\n", - "event: error\n", - "data: {\"type\":\"error\",\"error\":{\"details\":null,\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n", - ); - Mock::given(method("POST")) - .and(path("/v1/messages")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(body), - ) - .mount(&server) - .await; - - let stream = model(&server).do_stream(&options()).await.unwrap(); - let parts = collect_stream(stream).await; - - let saw_text = parts.iter().any(|p| { - matches!(p, - StreamPart::TextDelta { delta, .. } if delta == "Hello") - }); - assert!( - saw_text, - "expected a 'Hello' text delta before the error, got {parts:?}" - ); - - assert!( - parts.iter().any(|p| matches!(p, - StreamPart::Error { error: AiMuxError::ApiCall(m) } if m.message == "Overloaded")), - "expected a StreamPart::Error carrying 'Overloaded', got {parts:?}" - ); - - // Finish=final chunk: a mid-stream error must still be followed by a - // Finish part (error finish), never leave the stream unterminated. - assert!( - matches!(parts.last(), - Some(StreamPart::Finish { finish_reason, .. }) - if matches!(finish_reason.unified, - aimux_core::prelude::FinishReasonUnified::Error)), - "expected the last part to be Finish with unified=Error, got {parts:?}" - ); - // Error-finish zeroes usage: the message_start input_tokens (17) - // must not leak into the final part (mirrors openai behaviour). - if let Some(StreamPart::Finish { usage, .. }) = parts.last() { - assert_eq!( - usage.input_tokens.total, None, - "usage must be zeroed on error finish" - ); - } - } - - /// TS: "should forward error chunks" — a generic (non-overloaded) error - /// event mid-stream must also be forwarded. - #[tokio::test] - async fn mid_stream_generic_error_is_forwarded() { - let server = MockServer::start().await; - let body = concat!( - "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_01\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-haiku-20240307\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":17,\"output_tokens\":1}}}\n\n", - "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", - "data: {\"type\":\"ping\"}\n\n", - "data: {\"type\":\"error\",\"error\":{\"type\":\"error\",\"message\":\"test error\"}}\n\n", - ); - Mock::given(method("POST")) - .and(path("/v1/messages")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(body), - ) - .mount(&server) - .await; - - let stream = model(&server).do_stream(&options()).await.unwrap(); - let parts = collect_stream(stream).await; - assert!( - parts.iter().any(|p| matches!(p, - StreamPart::Error { error: AiMuxError::ApiCall(m) } if m.message == "test error")), - "expected a StreamPart::Error carrying 'test error', got {parts:?}" - ); - assert!( - matches!(parts.last(), - Some(StreamPart::Finish { finish_reason, .. }) - if matches!(finish_reason.unified, - aimux_core::prelude::FinishReasonUnified::Error)), - "expected the last part to be Finish with unified=Error, got {parts:?}" - ); - } } diff --git a/aimux-providers/tests/recraft_test.rs b/aimux-providers/tests/recraft_test.rs index dfb81dfbf..de7a3443e 100644 --- a/aimux-providers/tests/recraft_test.rs +++ b/aimux-providers/tests/recraft_test.rs @@ -209,10 +209,6 @@ async fn status_500_maps_to_provider_error() { ); } -// ════════════════════════════════════════════════════════════════════════════ -// Provider trait -// ════════════════════════════════════════════════════════════════════════════ - // ════════════════════════════════════════════════════════════════════════════ // Auth header matcher (verifies the header is sent on the wire) // ════════════════════════════════════════════════════════════════════════════ diff --git a/aimux-providers/tests/replicate_image_test.rs b/aimux-providers/tests/replicate_image_test.rs index 385cf110a..8e8665309 100644 --- a/aimux-providers/tests/replicate_image_test.rs +++ b/aimux-providers/tests/replicate_image_test.rs @@ -143,6 +143,32 @@ async fn should_pass_auth_headers() { ); } +#[tokio::test] +async fn should_use_custom_wait_time() { + let server = MockServer::start().await; + mock_replicate(&server, json!(format!("{}/image.png", server.uri()))).await; + let config = ReplicateProviderSettings { + api_key: Some(Resolvable::Value("test-token".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_replicate(config) + .unwrap() + .image("black-forest-labs/flux-schnell"); + let mut opts = options(PROMPT); + opts.n = 1; + opts.provider_options.extend( + aimux_core::shared::provider_namespace("replicate", json!({ "maxWaitTimeInSeconds": 30 })) + .unwrap(), + ); + model.do_generate(&opts).await.unwrap(); + let reqs = server.received_requests().await.unwrap(); + assert_eq!( + reqs[0].headers.get("prefer").unwrap().to_str().unwrap(), + "wait=30" + ); +} + #[tokio::test] async fn should_handle_array_output() { let server = MockServer::start().await; diff --git a/aimux-providers/tests/searxng_test.rs b/aimux-providers/tests/searxng_test.rs index 9f8150b9c..8cf8b2303 100644 --- a/aimux-providers/tests/searxng_test.rs +++ b/aimux-providers/tests/searxng_test.rs @@ -174,8 +174,6 @@ async fn status_401_maps_to_auth_error() { ); } -// -- Provider trait ---------------------------------------------------------- - #[test] fn model_id_is_searxng_search() { let config = SearxngProviderSettings { @@ -187,5 +185,3 @@ fn model_id_is_searxng_search() { assert_eq!(model.model_id(), "searxng-search"); assert_eq!(model.provider(), "searxng.search"); } - -// -- instance URL ----------------------------------------------------------- diff --git a/aimux-providers/tests/sse_terminal_error_test.rs b/aimux-providers/tests/sse_terminal_error_test.rs index 8fa09556d..cb2781a0e 100644 --- a/aimux-providers/tests/sse_terminal_error_test.rs +++ b/aimux-providers/tests/sse_terminal_error_test.rs @@ -48,13 +48,17 @@ async fn corrupted_stream(prefix: Value) -> Vec> fn assert_terminal_decode_error(parts: &[Result]) { let errors: Vec<_> = parts .iter() - .filter_map(|part| part.as_ref().err()) + .filter_map(|part| match part { + Ok(StreamPart::Error { error }) => Some(error), + _ => None, + }) .collect(); assert_eq!(errors.len(), 1); assert!(matches!(errors[0], AiMuxError::ApiCall(_))); assert!(!errors[0].is_recoverable_stream_error()); assert!(!errors[0].is_retryable()); - assert!(parts.last().is_some_and(Result::is_err)); + assert!(matches!(parts.last(), Some(Ok(StreamPart::Error { .. })))); + assert!(parts.iter().all(Result::is_ok)); assert!( !parts .iter() diff --git a/aimux-providers/tests/tinyfish_test.rs b/aimux-providers/tests/tinyfish_test.rs index e011f0b52..fee7e166d 100644 --- a/aimux-providers/tests/tinyfish_test.rs +++ b/aimux-providers/tests/tinyfish_test.rs @@ -198,5 +198,3 @@ async fn status_401_maps_to_auth_error() { "expected Auth error, got {result:?}" ); } - -// -- Provider trait ---------------------------------------------------------- diff --git a/aimux-providers/tests/vertex_upstream_e2e_test.rs b/aimux-providers/tests/vertex_upstream_e2e_test.rs new file mode 100644 index 000000000..1bd3feff0 --- /dev/null +++ b/aimux-providers/tests/vertex_upstream_e2e_test.rs @@ -0,0 +1,426 @@ +//! In-process end-to-end subset of the pinned Vertex upstream ports. + +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::{CallOptions, Tool}; +use aimux_core::result::GenerateContent; +use aimux_core::shared::{AspectRatio, provider_namespace}; +use aimux_core::tool::FunctionTool; +use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions}; +use aimux_core::types::FinishReasonUnified; +use aimux_core::video_model::{VideoCallOptions, VideoData, VideoModel, VideoOperationStatus}; +use aimux_provider_utils::Resolvable; +use aimux_providers::vertex::{VertexProvider, VertexProviderSettings, create_google_vertex}; +use mock_fetch::{Canned, MockFetch}; +use serde_json::json; +use std::sync::Arc; + +fn provider(fetch: &Arc) -> VertexProvider { + create_google_vertex(VertexProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".into())), + base_url: Some("https://api.example.com".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} + +/// TS: "should preserve local JSON Schema references in tool requests" (@ai-sdk/google-vertex/src/google-vertex-language-model.test.ts) +#[tokio::test] +async fn language_generate_preserves_tool_schema_references() { + let response = json!({ + "candidates": [{"content": {"parts": [{"text": "Done"}], "role": "model"}, "finishReason": "STOP", "index": 0}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2} + }); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let schema = json!({ + "type": "object", + "properties": {"locale": {"$ref": "#/$defs/Locale", "description": "Locale for formatting"}}, + "required": ["locale"], "additionalProperties": false, + "$defs": {"Locale": {"type": "string", "enum": ["de", "en"]}} + }); + 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: schema.clone(), + strict: None, + provider_options: None, + input_examples: None, + })]); + let result = provider(&fetch) + .chat("gemini-2.5-flash") + .do_generate(&options) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "POST"); + assert_eq!( + seen[0].url, + "https://api.example.com/models/gemini-2.5-flash:generateContent" + ); + assert_eq!(seen[0].headers["x-goog-api-key"], "test-api-key"); + let body = seen[0].json_body(); + assert_eq!( + body["contents"], + json!([{"role": "user", "parts": [{"text": "Hello"}]}]) + ); + assert_eq!( + body["tools"][0]["functionDeclarations"][0], + json!({ + "name": "format-date", "description": "Format a date", "parametersJsonSchema": schema + }) + ); + assert_eq!( + result.content, + vec![GenerateContent::Text { + text: "Done".into(), + provider_metadata: None + }] + ); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(result.usage.input_tokens.total, Some(1)); + assert_eq!(result.usage.output_tokens.total, Some(1)); + assert_eq!( + result.usage.raw, + response["usageMetadata"].as_object().cloned() + ); + let metadata = serde_json::to_value(result.provider_metadata.unwrap()).unwrap(); + assert_eq!( + metadata["googleVertex"]["usageMetadata"], + response["usageMetadata"] + ); + assert_eq!(metadata["vertex"], metadata["googleVertex"]); + assert_eq!(result.request.unwrap().body.unwrap(), body); + 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) +#[tokio::test] +async fn embedding_request_and_result() { + let response = json!({"predictions": [ + {"embeddings": {"statistics": {"token_count": 5, "truncated": false}, "values": [-0.017999587580561638, -0.006893285550177097, -0.036766719073057175, -0.017558680847287178, -0.019938766956329346]}}, + {"embeddings": {"statistics": {"token_count": 6, "truncated": false}, "values": [-0.06007182598114014, 0.004907649010419846, -0.00690646655857563, -0.007314121350646019, -0.048464205116033554]}} + ], "metadata": {"billableCharacterCount": 35}}); + let fetch = MockFetch::new(vec![Canned::json(&response)]); + let mut options = EmbeddingCallOptions::new("test text one"); + options.values.push("test text two".into()); + options.provider_options = Some(provider_namespace( + "google", + json!({ + "outputDimensionality": 768, "taskType": "SEMANTIC_SIMILARITY", "title": "test title", "autoTruncate": false + }), + ).expect("provider metadata must be an object")); + let result = provider(&fetch) + .embedding("textembedding-gecko@001") + .do_embed(&options) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!( + seen[0].url, + "https://api.example.com/models/textembedding-gecko@001:predict" + ); + assert_eq!( + seen[0].json_body(), + 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": {"outputDimensionality": 768, "autoTruncate": false} + }) + ); + let expected: Vec> = response["predictions"] + .as_array() + .unwrap() + .iter() + .map(|prediction| { + prediction["embeddings"]["values"] + .as_array() + .unwrap() + .iter() + .map(|value| value.as_f64().unwrap() as f32) + .collect() + }) + .collect(); + assert_eq!(result.embeddings, expected); + assert_eq!(result.usage.unwrap().tokens, 11); + assert!(result.warnings.is_empty()); + assert!(result.provider_metadata.is_none()); + 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) +#[tokio::test] +async fn image_generate_request_and_result() { + let response = json!({ + "candidates": [{"content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": "base64-generated-image"}}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": 200, "totalTokenCount": 220} + }); + let mut canned = Canned::json(&response); + canned + .headers + .push(("request-id".into(), "test-request-id".into())); + let fetch = MockFetch::new(vec![canned]); + let provider = create_google_vertex(VertexProviderSettings { + api_key: Some(Resolvable::Value("test-key".into())), + base_url: Some("https://api.example.com".into()), + headers: Some(Resolvable::Value( + [( + "Custom-Provider-Header".into(), + Some("provider-header-value".into()), + )] + .into(), + )), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + 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".into(), + "request-header-value".into(), + )] + .into(), + ); + let result = provider + .image("gemini-2.5-flash-image") + .do_generate(&options) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!( + seen[0].url, + "https://api.example.com/models/gemini-2.5-flash-image:generateContent" + ); + assert_eq!( + seen[0].headers["custom-provider-header"], + "provider-header-value" + ); + assert_eq!( + seen[0].headers["custom-request-header"], + "request-header-value" + ); + assert_eq!( + seen[0].json_body(), + json!({ + "contents": [{"role": "user", "parts": [{"text": "A beautiful sunset"}]}], + "generationConfig": {"responseModalities": ["IMAGE"], "imageConfig": {"aspectRatio": "21:9"}, "seed": 12345} + }) + ); + assert!( + matches!(result.images, ImageOutputs::Base64(images) if images == ["base64-generated-image"]) + ); + assert_eq!( + serde_json::to_value(result.provider_metadata.unwrap()).unwrap(), + json!({"googleVertex": {"images": [{}]}, "vertex": {"images": [{}]}}) + ); + let usage = result.usage.unwrap(); + assert_eq!(usage.input_tokens, Some(20)); + assert_eq!(usage.output_tokens, Some(200)); + assert_eq!(usage.total_tokens, Some(220)); + assert_eq!( + result.response.model_id.as_deref(), + Some("gemini-2.5-flash-image") + ); + assert_eq!( + result.response.headers.unwrap()["request-id"], + "test-request-id" + ); +} + +/// 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) +#[tokio::test] +async fn transcription_generate_request_and_result() { + let fetch = MockFetch::new(vec![Canned::json(&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"} + }))]); + let provider = create_google_vertex(VertexProviderSettings { + access_token: Some(Resolvable::Value("test-token".into())), + project: Some("test-project".into()), + location: Some("us-central1".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + let options = TranscriptionCallOptions::new( + AudioInput::Binary(vec![1, 2, 3, 4, 5, 6, 7, 8]), + "audio/wav", + ); + let result = provider + .transcription("chirp_2") + .do_generate(&options) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!( + seen[0].url, + "https://us-central1-speech.googleapis.com/v2/projects/test-project/locations/us-central1/recognizers/_:recognize" + ); + assert_eq!(seen[0].headers["authorization"], "Bearer test-token"); + assert_eq!( + seen[0].json_body(), + json!({ + "config": {"model": "chirp_2", "languageCodes": ["auto"], "autoDecodingConfig": {}, "features": {"enableWordTimeOffsets": true, "enableAutomaticPunctuation": true}}, + "content": "AQIDBAUGBwg=" + }) + ); + assert_eq!(result.text, "hello world"); + assert_eq!( + serde_json::to_value(result.segments).unwrap(), + json!([ + {"text": "hello", "start_second": 0.0, "end_second": 0.5}, + {"text": "world", "start_second": 0.5, "end_second": 1.0} + ]) + ); + assert_eq!(result.language.as_deref(), Some("en")); + 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) +#[tokio::test] +async fn video_start_and_completed_result() { + let fetch = MockFetch::new(vec![ + Canned::json(&json!({"name": "operations/my-op-123", "done": false})), + Canned::json( + &json!({"done": true, "response": {"videos": [{"bytesBase64Encoded": "base64-video-data", "mimeType": "video/mp4"}]}}), + ), + ]); + let model = provider(&fetch).video("veo-2.0-generate-001"); + let options = VideoCallOptions::new("A futuristic city with flying cars"); + let start = model.do_start(&options).await.unwrap(); + assert_eq!( + start.operation, + json!({"operationName": "operations/my-op-123"}) + ); + let status = model.do_status(&start.operation, &options).await.unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 2); + assert_eq!( + seen[0].url, + "https://api.example.com/models/veo-2.0-generate-001:predictLongRunning" + ); + assert_eq!( + seen[0].json_body(), + json!({"instances": [{"prompt": "A futuristic city with flying cars"}], "parameters": {"sampleCount": 1}}) + ); + assert_eq!(seen[1].method, "POST"); + assert_eq!( + seen[1].url, + "https://api.example.com/models/veo-2.0-generate-001:fetchPredictOperation" + ); + assert_eq!( + seen[1].json_body(), + json!({"operationName": "operations/my-op-123"}) + ); + let VideoOperationStatus::Completed(result) = status else { + panic!("expected completed video") + }; + assert!( + matches!(&result.videos[..], [VideoData::Base64 { data, media_type }] if data == "base64-video-data" && media_type == "video/mp4") + ); + let metadata = json!({"videos": [{"mimeType": "video/mp4"}]}); + assert_eq!( + serde_json::to_value(result.provider_metadata.unwrap()).unwrap(), + json!({"googleVertex": metadata, "vertex": metadata, "google-vertex": metadata}) + ); + assert_eq!( + result.response.model_id.as_deref(), + Some("veo-2.0-generate-001") + ); +} + +/// TS: "should throw when no operation name returned" (@ai-sdk/google-vertex/src/google-vertex-video-model.test.ts) +#[tokio::test] +async fn video_missing_operation_name_is_an_error() { + let fetch = MockFetch::new(vec![Canned::json(&json!({"done": false}))]); + let error = provider(&fetch) + .video("veo-2.0-generate-001") + .do_start(&VideoCallOptions::new("A futuristic city with flying cars")) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("No operation name returned from API") + ); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!( + seen[0].json_body(), + json!({"instances": [{"prompt": "A futuristic city with flying cars"}], "parameters": {"sampleCount": 1}}) + ); +} + +/// TS: "transcribes audio via Vertex generateContent with audioTranscriptionConfig" +/// (google-vertex/src/gemini-transcription/google-vertex-gemini-transcription-model.test.ts) +#[tokio::test] +async fn gemini_transcription_generate_request_and_result() { + let usage = json!({"promptTokenCount": 10, "candidatesTokenCount": 4}); + let fetch = MockFetch::new(vec![Canned::json(&json!({ + "candidates": [{"content": {"parts": [{"text": "Hello "}, {"text": "world."}]}}], + "usageMetadata": usage, + }))]); + let provider = create_google_vertex(VertexProviderSettings { + access_token: Some(Resolvable::Value("test-oauth-token".into())), + project: Some("test-project".into()), + location: Some("us-central1".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap(); + let mut options = + TranscriptionCallOptions::new(AudioInput::Binary(vec![1, 2, 3, 4]), "audio/wav"); + options.provider_options = Some(provider_namespace( + "googleVertex", + json!({ + "customVocabulary": ["Gemini", "Kubernetes"], "languageCodes": ["es-ES"], "mode": "SMART", + }), + ).expect("provider metadata must be an object")); + let result = provider + .transcription("gemini-3.5-transcribe") + .do_generate(&options) + .await + .unwrap(); + let seen = fetch.seen(); + assert_eq!(seen.len(), 1); + assert_eq!( + seen[0].url, + "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-transcribe:generateContent" + ); + assert_eq!(seen[0].headers["authorization"], "Bearer test-oauth-token"); + 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!(result.text, "Hello world."); + assert_eq!( + serde_json::to_value(result.provider_metadata.unwrap()).unwrap(), + json!({"google": {"usageMetadata": usage}}) + ); +} diff --git a/aimux-providers/tests/voyage_embedding_test.rs b/aimux-providers/tests/voyage_embedding_test.rs index 332268881..329b7ad59 100644 --- a/aimux-providers/tests/voyage_embedding_test.rs +++ b/aimux-providers/tests/voyage_embedding_test.rs @@ -178,6 +178,68 @@ async fn should_pass_model_and_values() { assert!(body.get("input_type").is_none()); } +/// TS: "should pass the input_type setting" +#[tokio::test] +async fn should_pass_input_type_setting() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/embeddings")) + .respond_with(ResponseTemplate::new(200).set_body_json(embedding_response_body())) + .mount(&server) + .await; + + let provider = test_provider("test-api-key", server.uri()); + let model = provider.embedding("voyage-3.5"); + + let provider_options = + aimux_core::shared::provider_namespace("voyage", json!({"inputType": "document"})).unwrap(); + let options = EmbeddingCallOptions { + values: test_values(), + abort_signal: None, + provider_options: Some(provider_options), + headers: None, + max_retries: None, + timeout: None, + }; + + let _ = model.do_embed(&options).await.expect("should succeed"); + + let requests = server.received_requests().await.expect("requests recorded"); + let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!(body["input_type"], "document"); +} + +/// TS: "should pass the output_dimension setting" +#[tokio::test] +async fn should_pass_output_dimension_setting() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/embeddings")) + .respond_with(ResponseTemplate::new(200).set_body_json(embedding_response_body())) + .mount(&server) + .await; + + let provider = test_provider("test-api-key", server.uri()); + let model = provider.embedding("voyage-3.5"); + + let provider_options = + aimux_core::shared::provider_namespace("voyage", json!({"outputDimension": 256})).unwrap(); + let options = EmbeddingCallOptions { + values: test_values(), + abort_signal: None, + provider_options: Some(provider_options), + headers: None, + max_retries: None, + timeout: None, + }; + + let _ = model.do_embed(&options).await.expect("should succeed"); + + let requests = server.received_requests().await.expect("requests recorded"); + let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!(body["output_dimension"], 256); +} + /// TS: "should pass headers" #[tokio::test] async fn should_pass_headers() { diff --git a/aimux-providers/tests/xai_upstream_e2e_test.rs b/aimux-providers/tests/xai_upstream_e2e_test.rs new file mode 100644 index 000000000..89adc16f6 --- /dev/null +++ b/aimux-providers/tests/xai_upstream_e2e_test.rs @@ -0,0 +1,266 @@ +//! Small Responses API subset from the ported xAI upstream cases. + +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use aimux_core::AiMuxError; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_core::result::GenerateContent; +use aimux_core::shared::provider_namespace; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::{ProviderTool, Tool}; +use aimux_core::types::FinishReasonUnified; +use aimux_providers::xai::{XAIProviderSettings, XaiResponsesModel, create_xai}; +use futures::TryStreamExt; +use mock_fetch::{Canned, MockFetch}; +use serde_json::{Value, json}; + +fn model(fetch: &std::sync::Arc) -> XaiResponsesModel { + create_xai(XAIProviderSettings { + base_url: Some("https://example.test/v1/".into()), + api_key: Some("test-key".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() + .responses("grok-4-fast-non-reasoning") +} + +fn options() -> CallOptions { + CallOptions { + prompt: vec![LanguageModelMessage::user_text("hello")], + ..Default::default() + } +} + +fn assert_request(fetch: &MockFetch, stream: bool, tools: bool) { + let requests = fetch.seen(); + assert_eq!(requests.len(), 1); + let request = &requests[0]; + assert_eq!(request.method, "POST"); + assert_eq!(request.url, "https://example.test/v1/responses"); + assert_eq!(request.headers["authorization"], "Bearer test-key"); + let mut expected = json!({ + "model": "grok-4-fast-non-reasoning", + "input": [{"role": "user", "content": [{"type": "input_text", "text": "hello"}]}] + }); + if stream { + expected["stream"] = json!(true); + } + if tools { + expected["tools"] = json!([{"type": "web_search"}]); + } + assert_eq!(request.json_body(), expected); +} + +fn sse(events: &[Value]) -> Canned { + Canned { + status: 200, + headers: vec![("content-type".into(), "text/event-stream".into())], + body: (events + .iter() + .map(|event| format!("data: {event}\n\n")) + .collect::() + + "data: [DONE]\n\n") + .into_bytes(), + } +} + +/// TS: "should extract usage correctly" (xai/src/responses/xai-responses-language-model.test.ts) +#[tokio::test] +async fn generate_converts_usage() { + // The port used an unrelated search fixture for this case; restore the TS response. + let usage = json!({ + "input_tokens": 345, "output_tokens": 538, "total_tokens": 883, + "output_tokens_details": {"reasoning_tokens": 123} + }); + let body = json!({ + "id": "resp_123", "object": "response", "status": "completed", + "model": "grok-4-fast-non-reasoning", "output": [], "usage": usage + }); + let fetch = MockFetch::new(vec![Canned::json(&body)]); + let result = model(&fetch).do_generate(&options()).await.unwrap(); + assert_request(&fetch, false, false); + assert!(result.content.is_empty()); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(result.finish_reason.raw.as_deref(), Some("completed")); + assert_eq!( + serde_json::to_value(result.usage).unwrap(), + json!({ + "input_tokens": {"total": 345, "no_cache": 345, "cache_read": 0}, + "output_tokens": {"total": 538, "text": 415, "reasoning": 123}, + "raw": usage + }) + ); + assert!(result.provider_metadata.is_none()); + assert_eq!(result.response.unwrap().id.as_deref(), Some("resp_123")); +} + +/// TS: "should expose cost_in_usd_ticks in finish providerMetadata" (xai/src/responses/xai-responses-language-model.test.ts) +#[tokio::test] +async fn stream_reports_cost_and_text() { + let events = vec![ + json!({"type": "response.created", "response": {"id": "resp_123", "object": "response", "model": "grok-4-fast-non-reasoning", "output": []}}), + json!({"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "Hello"}), + json!({"type": "response.completed", "response": {"id": "resp_123", "object": "response", "model": "grok-4-fast-non-reasoning", "status": "completed", "output": [{"type": "message", "id": "msg_001", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": "Hello"}]}], "usage": {"input_tokens": 10, "output_tokens": 5, "cost_in_usd_ticks": 113500}}}), + ]; + let fetch = MockFetch::new(vec![sse(&events)]); + let parts: Vec<_> = model(&fetch) + .do_stream(&options()) + .await + .unwrap() + .stream + .try_collect() + .await + .unwrap(); + assert_request(&fetch, true, false); + let text: String = parts + .iter() + .filter_map(|part| match part { + StreamPart::TextDelta { delta, .. } => Some(delta.as_str()), + _ => None, + }) + .collect(); + assert_eq!(text, "Hello"); + let finishes: Vec<_> = parts + .iter() + .filter_map(|part| match part { + StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } => Some((finish_reason, usage, provider_metadata)), + _ => None, + }) + .collect(); + assert_eq!(finishes.len(), 1); + let (reason, usage, metadata) = finishes[0]; + assert_eq!(reason.unified, FinishReasonUnified::Stop); + assert_eq!(reason.raw.as_deref(), Some("completed")); + assert_eq!(usage.input_tokens.total, Some(10)); + assert_eq!(usage.output_tokens.total, Some(5)); + assert_eq!( + metadata.as_ref().unwrap(), + &provider_namespace("xai", json!({"costInUsdTicks": 113500})).unwrap() + ); + assert!( + !parts + .iter() + .any(|part| matches!(part, StreamPart::Error { .. })) + ); +} + +/// TS: "should forward web_search_call action and sources to tool-result" (xai/src/responses/xai-responses-language-model.test.ts) +#[tokio::test] +async fn tool_results_preserve_actions_and_sources() { + let body = json!({ + "id": "resp_123", + "object": "response", + "status": "completed", + "model": "grok-4-fast-non-reasoning", + "output": [ + {"type": "web_search_call", "id": "ws_action_1", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "search", "query": "latest AI news", "sources": [{"type": "url", "url": "https://example.com/a"}, {"type": "url", "url": "https://example.com/b"}]}}, + {"type": "web_search_call", "id": "ws_action_2", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "open_page", "url": "https://example.com/a"}}, + {"type": "web_search_call", "id": "ws_action_3", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "find_in_page", "url": "https://example.com/a", "pattern": "climate"}}, + {"type": "web_search_call", "id": "ws_action_4", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "search", "query": null, "sources": null}}, + {"type": "web_search_call", "id": "ws_action_5", "name": "web_search", "call_id": "", "status": "completed", "action": {"type": "open_page", "url": null}}, + ], + "usage": {"input_tokens": 10, "output_tokens": 5}, + }); + let fetch = MockFetch::new(vec![Canned::json(&body)]); + let mut options = options(); + options.tools = Some(vec![Tool::Provider(ProviderTool { + id: "xai.web_search".into(), + name: "web_search".into(), + args: serde_json::Map::new(), + })]); + let result = model(&fetch).do_generate(&options).await.unwrap(); + assert_request(&fetch, false, true); + let expected = vec![ + json!({"action": {"type": "search", "query": "latest AI news"}, "sources": [{"type": "url", "url": "https://example.com/a"}, {"type": "url", "url": "https://example.com/b"}]}), + json!({"action": {"type": "openPage", "url": "https://example.com/a"}}), + json!({"action": {"type": "findInPage", "url": "https://example.com/a", "pattern": "climate"}}), + json!({"action": {"type": "search"}}), + json!({"action": {"type": "openPage", "url": null}}), + ]; + assert_eq!(result.content.len(), expected.len() * 2); + for (index, (pair, expected_result)) in result + .content + .as_chunks::<2>() + .0 + .iter() + .zip(expected) + .enumerate() + { + let id = format!("ws_action_{}", index + 1); + let GenerateContent::ToolCall(call) = &pair[0] else { + panic!("missing tool call") + }; + assert_eq!(call.tool_call_id, id); + assert_eq!(call.tool_name, "web_search"); + assert_eq!(call.input, ""); + assert_eq!(call.provider_executed, Some(true)); + let GenerateContent::ToolResult(output) = &pair[1] else { + panic!("missing tool result") + }; + assert_eq!(output.tool_call_id, id); + assert_eq!(output.tool_name, "web_search"); + assert_eq!(output.result, expected_result); + assert_eq!(output.is_error, None); + } + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); +} + +/// TS: "should set finish reason to error" (xai/src/responses/xai-responses-language-model.test.ts) +#[tokio::test] +async fn failed_response_emits_error_and_finish() { + let events = vec![ + json!({"type": "response.created", "response": {"id": "resp_123", "object": "response", "model": "grok-4-fast-non-reasoning", "output": []}}), + json!({"type": "response.failed", "response": {"error": {"code": "server_error", "message": "Internal server error"}, "usage": {"input_tokens": 50, "output_tokens": 0}}}), + ]; + let fetch = MockFetch::new(vec![sse(&events)]); + let parts: Vec<_> = model(&fetch) + .do_stream(&options()) + .await + .unwrap() + .stream + .try_collect() + .await + .unwrap(); + assert_request(&fetch, true, false); + let errors: Vec<_> = parts + .iter() + .filter_map(|part| match part { + StreamPart::Error { error } => Some(error), + _ => None, + }) + .collect(); + assert_eq!(errors.len(), 1); + let AiMuxError::ApiCall(error) = errors[0] else { + panic!("expected API error") + }; + assert_eq!(error.message, "Internal server error"); + assert_eq!(error.provider_code.as_deref(), Some("server_error")); + assert_eq!(error.status_code, Some(500)); + assert!(error.is_retryable); + assert_eq!(error.data.as_ref(), Some(&events[1])); + let finishes: Vec<_> = parts + .iter() + .filter_map(|part| match part { + StreamPart::Finish { + finish_reason, + usage, + .. + } => Some((finish_reason, usage)), + _ => None, + }) + .collect(); + assert_eq!(finishes.len(), 1); + let (reason, usage) = finishes[0]; + assert_eq!(reason.unified, FinishReasonUnified::Error); + assert_eq!(reason.raw.as_deref(), Some("error")); + assert_eq!(usage.input_tokens.total, Some(50)); + assert_eq!(usage.output_tokens.total, Some(0)); +} diff --git a/aimux-providers/tests/you_com_test.rs b/aimux-providers/tests/you_com_test.rs index 9d0accd97..dfb3b93ad 100644 --- a/aimux-providers/tests/you_com_test.rs +++ b/aimux-providers/tests/you_com_test.rs @@ -179,5 +179,3 @@ async fn status_401_maps_to_auth_error() { "expected Auth error, got {result:?}" ); } - -// -- Provider trait ---------------------------------------------------------- diff --git a/docs/api/providers.md b/docs/api/providers.md index 2c8d116b5..60e6029fb 100644 --- a/docs/api/providers.md +++ b/docs/api/providers.md @@ -330,7 +330,7 @@ These providers are **not** name-addressable: `provider("anthropic", ...)` fails | `jev` | `JevConfig` / `JevProvider` | | `mistral` | `MistralProvider` / `MistralProviderSettings` / `create_mistral` | | `openai` | `OpenAIProvider` / `OpenAIProviderSettings` / `create_openai` | -| `vertex` | `VertexProvider` / `VertexProviderSettings` / `create_google_vertex` | +| `vertex` | `VertexAnthropicProvider` / `VertexAnthropicProviderSettings` / `VertexProvider` / `VertexProviderSettings` / `create_google_vertex` / `create_google_vertex_anthropic` | | `voyage` | `VoyageProvider` / `VoyageProviderSettings` / `create_voyage` | | `codex` | `CodexProvider` / `CodexProviderSettings` / `create_codex` | | `xai` | `XAIProvider` / `XAIProviderSettings` / `create_xai` |